@Override
  public List<String> assignSegment(String segmentName, Map<String, Map<String, String>> currentAssignment,
      InstancePartitions instancePartitions, InstancePartitionsType instancePartitionsType) {
    String serverTag = _tenantConfig.getServer();
    Set<String> instances = HelixHelper.getServerInstancesForTenant(_helixManager, serverTag);
    int numInstances = instances.size();
    Preconditions.checkState(numInstances > 0, "No instance found with tag: %s or %s",
        TagNameUtils.getOfflineTagForTenant(serverTag), TagNameUtils.getRealtimeTagForTenant(serverTag));
    return new ArrayList<>(instances);
  }
@Test
  public void testSegmentAssignmentAndRebalance() {
    List<HelixProperty> instanceConfigList = new ArrayList<>();
    for (String instance : INSTANCES) {
      ZNRecord znRecord = new ZNRecord(instance);
      znRecord.setListField(TAG_LIST.name(), ImmutableList.of(OFFLINE_SERVER_TAG, REALTIME_SERVER_TAG));
      instanceConfigList.add(new InstanceConfig(znRecord));
    }
    HelixDataAccessor dataAccessor = mock(HelixDataAccessor.class);
    PropertyKey.Builder builder = new PropertyKey.Builder("cluster");
    when(dataAccessor.keyBuilder()).thenReturn(builder);
    when(dataAccessor.getChildValues(builder.instanceConfigs(), true)).thenReturn(instanceConfigList);
    when(_helixManager.getHelixDataAccessor()).thenReturn(dataAccessor);

    List<String> instances = _segmentAssignment.assignSegment(SEGMENT_NAME, new TreeMap(), _instancePartitionsMap);
    assertEquals(instances.size(), NUM_INSTANCES);
    assertEqualsNoOrder(instances.toArray(), INSTANCES.toArray());

    // Remove one instance and rebalance table
    Map<String, Map<String, String>> currentAssignment = new TreeMap<>();
    Map<String, String> segment1Assginment = new TreeMap<>();
    instances.stream().forEach(instance -> segment1Assginment.put(instance, "ONLINE"));
    currentAssignment.put(SEGMENT_NAME, segment1Assginment);
    ZNRecord znRecord = new ZNRecord(instanceConfigList.get(0).getId());
    znRecord.setListField(TAG_LIST.name(), ImmutableList.of(BROKER_TAG));
    InstanceConfig newInstanceConfig = new InstanceConfig(znRecord);
    instanceConfigList.set(0, newInstanceConfig);
    when(dataAccessor.getChildValues(builder.instanceConfigs(), true)).thenReturn(instanceConfigList);

    Map<String, Map<String, String>> newAssignment =
        _segmentAssignment.rebalanceTable(currentAssignment, _instancePartitionsMap, null, null, new RebalanceConfig());
    assertEquals(newAssignment.get(SEGMENT_NAME).size(), NUM_INSTANCES - 1);
  }