@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 testSegmentAssignmentToRealtimeHosts() {
    List<HelixProperty> instanceConfigList = new ArrayList<>();
    for (String instance : INSTANCES) {
      ZNRecord znRecord = new ZNRecord(instance);
      znRecord.setListField(TAG_LIST.name(), ImmutableList.of(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());
  }