public static Map<String, Map<String, String>> rebalanceTableWithHelixAutoRebalanceStrategy(
      Map<String, Map<String, String>> currentAssignment, List<String> instances, int replication) {
    // Use Helix AutoRebalanceStrategy to rebalance the table
    LinkedHashMap<String, Integer> states = new LinkedHashMap<>();
    states.put(SegmentStateModel.ONLINE, replication);
    AutoRebalanceStrategy autoRebalanceStrategy =
        new AutoRebalanceStrategy(null, new ArrayList<>(currentAssignment.keySet()), states);
    // Make a copy of the current assignment because this step might change the passed in assignment
    Map<String, Map<String, String>> currentAssignmentCopy = new TreeMap<>();
    for (Map.Entry<String, Map<String, String>> entry : currentAssignment.entrySet()) {
      String segmentName = entry.getKey();
      Map<String, String> instanceStateMap = entry.getValue();
      currentAssignmentCopy.put(segmentName, new TreeMap<>(instanceStateMap));
    }
    return autoRebalanceStrategy.computePartitionAssignment(instances, instances, currentAssignmentCopy, null)
        .getMapFields();
  }
@Test
  public void testRebalanceTableWithHelixAutoRebalanceStrategy() {
    int numSegments = 100;
    List<String> segments = SegmentAssignmentTestUtils.getNameList(SEGMENT_NAME_PREFIX, numSegments);
    int numInstances = 10;
    List<String> instances = SegmentAssignmentTestUtils.getNameList(INSTANCE_NAME_PREFIX, numInstances);

    // Uniformly spray segments to the instances:
    // [instance_0,   instance_1,   instance_2,   instance_3,   instance_4,   instance_5,   instance_6,   instance_7,
    // instance_8,   instance_9]
    //  segment_0(r0) segment_0(r1) segment_0(r2) segment_1(r0) segment_1(r1) segment_1(r2) segment_2(r0) segment_2
    //  (r1) segment_2(r2) segment_3(r0)
    //  segment_3(r1) segment_3(r2) segment_4(r0) segment_4(r1) segment_4(r2) segment_5(r0) segment_5(r1) segment_5
    //  (r2) segment_6(r0) segment_6(r1)
    //  ...
    Map<String, Map<String, String>> currentAssignment = new TreeMap<>();
    int assignedInstanceId = 0;
    for (String segmentName : segments) {
      List<String> instancesAssigned = new ArrayList<>(NUM_REPLICAS);
      for (int replicaId = 0; replicaId < NUM_REPLICAS; replicaId++) {
        instancesAssigned.add(instances.get(assignedInstanceId));
        assignedInstanceId = (assignedInstanceId + 1) % numInstances;
      }
      currentAssignment
          .put(segmentName, SegmentAssignmentUtils.getInstanceStateMap(instancesAssigned, SegmentStateModel.ONLINE));
    }

    // There should be 100 segments assigned
    assertEquals(currentAssignment.size(), numSegments);
    // Each segment should have 3 replicas
    for (Map<String, String> instanceStateMap : currentAssignment.values()) {
      assertEquals(instanceStateMap.size(), NUM_REPLICAS);
    }
    // Each instance should have 30 segments assigned
    int[] numSegmentsAssignedPerInstance =
        SegmentAssignmentUtils.getNumSegmentsAssignedPerInstance(currentAssignment, instances);
    int[] expectedNumSegmentsAssignedPerInstance = new int[numInstances];
    int numSegmentsPerInstance = numSegments * NUM_REPLICAS / numInstances;
    Arrays.fill(expectedNumSegmentsAssignedPerInstance, numSegmentsPerInstance);
    assertEquals(numSegmentsAssignedPerInstance, expectedNumSegmentsAssignedPerInstance);
    // Current assignment should already be balanced
    assertEquals(
        SegmentAssignmentUtils.rebalanceTableWithHelixAutoRebalanceStrategy(currentAssignment, instances, NUM_REPLICAS),
        currentAssignment);

    // Replace instance_0 with instance_10
    // {
    //   0_0=[instance_10, instance_1, instance_2, instance_3, instance_4, instance_5, instance_6, instance_7,
    //   instance_8, instance_9]
    // }
    List<String> newInstances = new ArrayList<>(instances);
    String newInstanceName = INSTANCE_NAME_PREFIX + 10;
    newInstances.set(0, newInstanceName);
    Map<String, Map<String, String>> newAssignment = SegmentAssignmentUtils
        .rebalanceTableWithHelixAutoRebalanceStrategy(currentAssignment, newInstances, NUM_REPLICAS);
    // There should be 100 segments assigned
    assertEquals(currentAssignment.size(), numSegments);
    // Each segment should have 3 replicas
    for (Map<String, String> instanceStateMap : newAssignment.values()) {
      assertEquals(instanceStateMap.size(), NUM_REPLICAS);
    }
    // Each instance should have 30 segments assigned
    numSegmentsAssignedPerInstance =
        SegmentAssignmentUtils.getNumSegmentsAssignedPerInstance(newAssignment, newInstances);
    assertEquals(numSegmentsAssignedPerInstance, expectedNumSegmentsAssignedPerInstance);
    // All segments on instance_0 should be moved to instance_10
    Map<String, Integer> numSegmentsToBeMovedPerInstance =
        SegmentAssignmentUtils.getNumSegmentsToBeMovedPerInstance(currentAssignment, newAssignment);
    assertEquals(numSegmentsToBeMovedPerInstance.size(), 1);
    assertEquals((int) numSegmentsToBeMovedPerInstance.get(newInstanceName), numSegmentsPerInstance);
    String oldInstanceName = INSTANCE_NAME_PREFIX + 0;
    for (String segmentName : segments) {
      if (currentAssignment.get(segmentName).containsKey(oldInstanceName)) {
        assertTrue(newAssignment.get(segmentName).containsKey(newInstanceName));
      }
    }

    // Remove 5 instances
    // {
    //   0_0=[instance_0, instance_1, instance_2, instance_3, instance_4]
    // }
    int newNumInstances = numInstances - 5;
    newInstances = SegmentAssignmentTestUtils.getNameList(INSTANCE_NAME_PREFIX, newNumInstances);
    newAssignment = SegmentAssignmentUtils
        .rebalanceTableWithHelixAutoRebalanceStrategy(currentAssignment, newInstances, NUM_REPLICAS);
    // There should be 100 segments assigned
    assertEquals(newAssignment.size(), numSegments);
    // Each segment should have 3 replicas
    for (Map<String, String> instanceStateMap : newAssignment.values()) {
      assertEquals(instanceStateMap.size(), NUM_REPLICAS);
    }
    // The segments are not perfectly balanced, but should be deterministic
    numSegmentsAssignedPerInstance =
        SegmentAssignmentUtils.getNumSegmentsAssignedPerInstance(newAssignment, newInstances);
    assertEquals(numSegmentsAssignedPerInstance[0], 56);
    assertEquals(numSegmentsAssignedPerInstance[1], 60);
    assertEquals(numSegmentsAssignedPerInstance[2], 60);
    assertEquals(numSegmentsAssignedPerInstance[3], 60);
    assertEquals(numSegmentsAssignedPerInstance[4], 64);
    numSegmentsToBeMovedPerInstance =
        SegmentAssignmentUtils.getNumSegmentsToBeMovedPerInstance(currentAssignment, newAssignment);
    assertEquals(numSegmentsToBeMovedPerInstance.size(), newNumInstances);
    assertEquals((int) numSegmentsToBeMovedPerInstance.get(newInstances.get(0)), 26);
    assertEquals((int) numSegmentsToBeMovedPerInstance.get(newInstances.get(1)), 30);
    assertEquals((int) numSegmentsToBeMovedPerInstance.get(newInstances.get(2)), 30);
    assertEquals((int) numSegmentsToBeMovedPerInstance.get(newInstances.get(3)), 30);
    assertEquals((int) numSegmentsToBeMovedPerInstance.get(newInstances.get(4)), 34);

    // Add 5 instances
    // {
    //   0_0=[
    //       instance_0, instance_1, instance_2, instance_3, instance_4, instance_5, instance_6, instance_7,
    //       instance_8, instance_9,
    //       instance_10, instance_11, instance_12, instance_13, instance_14]
    // }
    newNumInstances = numInstances + 5;
    newInstances = SegmentAssignmentTestUtils.getNameList(INSTANCE_NAME_PREFIX, newNumInstances);
    newAssignment = SegmentAssignmentUtils
        .rebalanceTableWithHelixAutoRebalanceStrategy(currentAssignment, newInstances, NUM_REPLICAS);
    // There should be 100 segments assigned
    assertEquals(newAssignment.size(), numSegments);
    // Each segment should have 3 replicas
    for (Map<String, String> instanceStateMap : newAssignment.values()) {
      assertEquals(instanceStateMap.size(), NUM_REPLICAS);
    }
    // Each instance should have 20 segments assigned
    numSegmentsAssignedPerInstance =
        SegmentAssignmentUtils.getNumSegmentsAssignedPerInstance(newAssignment, newInstances);
    expectedNumSegmentsAssignedPerInstance = new int[newNumInstances];
    int newNumSegmentsPerInstance = numSegments * NUM_REPLICAS / newNumInstances;
    Arrays.fill(expectedNumSegmentsAssignedPerInstance, newNumSegmentsPerInstance);
    assertEquals(numSegmentsAssignedPerInstance, expectedNumSegmentsAssignedPerInstance);
    // Each new added instance should have 20 segments to be moved to it
    numSegmentsToBeMovedPerInstance =
        SegmentAssignmentUtils.getNumSegmentsToBeMovedPerInstance(currentAssignment, newAssignment);
    assertEquals(numSegmentsToBeMovedPerInstance.size(), 5);
    for (int instanceId = numInstances; instanceId < newNumInstances; instanceId++) {
      assertEquals((int) numSegmentsToBeMovedPerInstance.get(newInstances.get(instanceId)), newNumSegmentsPerInstance);
    }

    // Change all instances
    // {
    //   0_0=[i_0, i_1, i_2, i_3, i_4, i_5, i_6, i_7, i_8, i_9]
    // }
    String newInstanceNamePrefix = "i_";
    newInstances = SegmentAssignmentTestUtils.getNameList(newInstanceNamePrefix, numInstances);
    newAssignment = SegmentAssignmentUtils
        .rebalanceTableWithHelixAutoRebalanceStrategy(currentAssignment, newInstances, NUM_REPLICAS);
    // There should be 100 segments assigned
    assertEquals(newAssignment.size(), numSegments);
    // Each segment should have 3 replicas
    for (Map<String, String> instanceStateMap : newAssignment.values()) {
      assertEquals(instanceStateMap.size(), NUM_REPLICAS);
    }
    // Each instance should have 30 segments assigned
    numSegmentsAssignedPerInstance =
        SegmentAssignmentUtils.getNumSegmentsAssignedPerInstance(newAssignment, newInstances);
    expectedNumSegmentsAssignedPerInstance = new int[numInstances];
    Arrays.fill(expectedNumSegmentsAssignedPerInstance, numSegmentsPerInstance);
    assertEquals(numSegmentsAssignedPerInstance, expectedNumSegmentsAssignedPerInstance);
    // Each instance should have 30 segments to be moved to it
    numSegmentsToBeMovedPerInstance =
        SegmentAssignmentUtils.getNumSegmentsToBeMovedPerInstance(currentAssignment, newAssignment);
    assertEquals(numSegmentsToBeMovedPerInstance.size(), numInstances);
    for (String instanceName : newInstances) {
      assertEquals((int) numSegmentsToBeMovedPerInstance.get(instanceName), numSegmentsPerInstance);
    }
  }