@Override
  public List<String> assignSegment(String segmentName, Map<String, Map<String, String>> currentAssignment,
      Map<InstancePartitionsType, InstancePartitions> instancePartitionsMap) {
    Preconditions.checkState(instancePartitionsMap.size() == 1, "One instance partition type should be provided");
    InstancePartitions instancePartitions = instancePartitionsMap.get(InstancePartitionsType.CONSUMING);
    Preconditions.checkState(instancePartitions != null, "Failed to find CONSUMING instance partitions for table: %s",
        _tableNameWithType);
    _logger.info("Assigning segment: {} with instance partitions: {} for table: {}", segmentName, instancePartitions,
        _tableNameWithType);

    int partitionId = getPartitionId(segmentName);
    List<String> instancesAssigned = assignConsumingSegment(partitionId, instancePartitions);
    Set<String> existingAssignment = getExistingAssignment(partitionId, currentAssignment);
    // Check if the candidate assignment is consistent with existing assignment. Use existing assignment if not.
    if (existingAssignment == null) {
      _logger.info("No existing assignment from idealState, using the one decided by instancePartitions");
    } else if (!isSameAssignment(existingAssignment, instancesAssigned)) {
      _logger.warn("Assignment: {} is inconsistent with idealState: {}, using the one from idealState",
          instancesAssigned, existingAssignment);
      instancesAssigned = new ArrayList<>(existingAssignment);
      if (_controllerMetrics != null) {
        _controllerMetrics.addMeteredTableValue(_tableNameWithType,
            ControllerMeter.CONTROLLER_REALTIME_TABLE_SEGMENT_ASSIGNMENT_MISMATCH, 1L);
      }
    }
    _logger.info("Assigned segment: {} to instances: {} for table: {}", segmentName, instancesAssigned,
        _tableNameWithType);
    return instancesAssigned;
  }
@Test(expectedExceptions = IllegalStateException.class)
  public void testAssignSegmentToCompletedServers() {
    _segmentAssignment.assignSegment("seg01", new TreeMap<>(), new TreeMap<>());
  }