@Test
public void testRemoveTreatment() throws ValidationException  {

    doReturn(0L).when(identity).computeHash(anyInt(), Mockito.<Set<String>>any(), any(AttributesMap.class));
    final Experiment experiment =
        new Experiment(null, "foo")
            .addTreatment("control")
            .addTreatment("cake")
            .addTreatment("pie")
            .allocate("control", 10)
            .allocate("cake", 10)
            .allocate("pie", 10)
            .allocate("control", 10)
            .allocate("cake", 10)
            .allocate("pie", 10)
            .addOverride("control_override", "control", "true");
    List<Treatment> treatments = Lists.newArrayList(experiment.getTreatments());
    assertEquals("should contain expected number of treatments", 3, treatments.size());
    List<Allocation> allocations = Lists.newArrayList(experiment.getAllocations());
    assertEquals("should contain expected number of allocations", 6, allocations.size());
    List<TreatmentOverride> overrides = Lists.newArrayList(experiment.getOverrides());
    assertEquals("should contain override", 1, overrides.size());
    experiment.removeTreatment("pie");
    treatments = Lists.newArrayList(experiment.getTreatments());
    assertEquals("should contain fewer treatments", 2, treatments.size());
    allocations = Lists.newArrayList(experiment.getAllocations());
    assertEquals("should contain fewer allocations", 4, allocations.size());
    overrides = Lists.newArrayList(experiment.getOverrides());
    assertEquals("should contain same override", 1, overrides.size());
    experiment.removeTreatment("control");
    treatments = Lists.newArrayList(experiment.getTreatments());
    assertEquals("should contain fewer treatments", 1, treatments.size());
    allocations = Lists.newArrayList(experiment.getAllocations());
    // (it's 1 allocation rather than 2, because the remaining 'cake' allocations merge into one)
    assertEquals("should contain fewer allocations", 1, allocations.size());
    overrides = Lists.newArrayList(experiment.getOverrides());
    assertEquals("should contain no overrides", 0, overrides.size());
}