@Test
public void testUpdateTreatment() {

    final UpdateTreatmentRequest request = new UpdateTreatmentRequest(Optional.of("new description"));
    post(TREATMENT_ENDPOINT, EXPERIMENT_BAD, EXP_1_TREATMENT_1)
        .entity(request)
        .assertStatus(Status.NOT_FOUND);
    post(TREATMENT_ENDPOINT, EXPERIMENT_1, TREATMENT_BAD)
        .entity(request)
        .assertStatus(Status.NOT_FOUND);
    post(TREATMENT_ENDPOINT, EXPERIMENT_1, EXP_1_TREATMENT_1)
        .entity(request)
        .assertStatus(Status.NO_CONTENT);
    final TreatmentDto treatment =
        get(TREATMENT_ENDPOINT, EXPERIMENT_1, EXP_1_TREATMENT_1)
            .assertStatus(Status.OK)
            .result(TreatmentDto.class);
    assertEquals(request.getDescription().orElse(null), treatment.getDescription());
    final Iterable<AllocationDto> allocations =
        get(ALLOCATIONS_ENDPOINT, EXPERIMENT_1)
            .assertStatus(Status.OK)
            .result(iterable(AllocationDto.class));
    // Make sure we didn't lose our allocations for this treatment
    assertNotNull(
        Iterables.find(allocations, new Predicate<AllocationDto>() {
            @Override
            public boolean apply(@Nullable AllocationDto input) {
                return input != null && input.getTreatment().equals(EXP_1_TREATMENT_1);
            }
        })
    );
}