@Test
  public void testBalancerAfterContainerIncrease() {
    Set<TaskModel> taskModels = generateTaskModels(9);
    Set<ContainerModel> prevContainers = new GroupByContainerCount(2).group(taskModels);
    Map<TaskName, String> prevTaskToContainerMapping = generateTaskContainerMapping(prevContainers);
    GrouperMetadataImpl grouperMetadata = new GrouperMetadataImpl(new HashMap<>(), new HashMap<>(), new HashMap<>(), prevTaskToContainerMapping);

    Set<ContainerModel> containers = new GroupByContainerCount(4).group(taskModels, grouperMetadata);

    Map<String, ContainerModel> containersMap = new HashMap<>();
    for (ContainerModel container : containers) {
      containersMap.put(container.getId(), container);
    }

    assertEquals(4, containers.size());
    ContainerModel container0 = containersMap.get("0");
    ContainerModel container1 = containersMap.get("1");
    ContainerModel container2 = containersMap.get("2");
    ContainerModel container3 = containersMap.get("3");
    assertNotNull(container0);
    assertNotNull(container1);
    assertNotNull(container2);
    assertNotNull(container3);
    assertEquals("0", container0.getId());
    assertEquals("1", container1.getId());
    assertEquals(3, container0.getTasks().size());
    assertEquals(2, container1.getTasks().size());
    assertEquals(2, container2.getTasks().size());
    assertEquals(2, container3.getTasks().size());

    // Tasks 0-4 should stay on the same original containers
    assertTrue(container0.getTasks().containsKey(getTaskName(0)));
    assertTrue(container0.getTasks().containsKey(getTaskName(2)));
    assertTrue(container0.getTasks().containsKey(getTaskName(4)));
    assertTrue(container1.getTasks().containsKey(getTaskName(1)));
    assertTrue(container1.getTasks().containsKey(getTaskName(3)));
    // Tasks 5-8 should be reassigned to the new containers.
    // Consistency is the goal with these reassignments
    assertTrue(container2.getTasks().containsKey(getTaskName(8)));
    assertTrue(container2.getTasks().containsKey(getTaskName(6)));
    assertTrue(container3.getTasks().containsKey(getTaskName(5)));
    assertTrue(container3.getTasks().containsKey(getTaskName(7)));
  }
Before:
 C0  C1
--------
 T0  T1
 T2  T3
 T4  T5
 T6  T7
 T8

After:
 C0  C1  C2  C3
----------------
 T0  T1  T6  T5
 T2  T3  T8  T7
 T4

 NOTE for host affinity, it would help to have some additional logic to reassign tasks
 from C0 and C1 to containers that were on the same respective hosts, it wasn't implemented
 because the scenario is infrequent, the benefits are not guaranteed, and the code complexity
 wasn't worth it. It certainly could be implemented in the future.
samza-core/src/test/java/org/apache/samza/container/grouper/task/TestGroupByContainerCount.java