@Override
  public List<Set<TopicPartition>> assignPartitions(Cluster cluster, int numMetricFetchers) {
    if (numMetricFetchers != SUPPORTED_NUM_METRIC_FETCHER) {
      throw new IllegalArgumentException("DefaultMetricSamplerPartitionAssignor supports only a single metric fetcher.");
    }
    // Create an array to host the assignment of all the metric fetchers.
    List<Set<TopicPartition>> assignments = new ArrayList<>();
    assignments.add(assignPartitions(cluster));
    return assignments;
  }
@Test
  public void testAssignment() {
    int maxNumPartitionsForTopic = -1;
    int totalNumPartitions = 0;
    // Prepare the metadata
    Set<PartitionInfo> partitions = new HashSet<>();
    for (int i = 0; i < NUM_TOPICS; i++) {
      // Random number of partitions ranging from 4 to 400
      int randomNumPartitions = 4 * (RANDOM.nextInt(100) + 1);
      maxNumPartitionsForTopic = Math.max(randomNumPartitions, maxNumPartitionsForTopic);
      totalNumPartitions += randomNumPartitions;
      for (int j = 0; j < randomNumPartitions; j++) {
        partitions.add(new PartitionInfo(TOPIC_PREFIX + i, j, NODE_0, nodes(), nodes()));
      }
    }
    Cluster cluster = new Cluster("cluster", Arrays.asList(nodes()), partitions, Collections.emptySet(), Collections.emptySet());
    Metadata metadata = new Metadata(METADATA_REFRESH_BACKOFF,
                                     METADATA_EXPIRY_MS,
                                     new LogContext(),
                                     new ClusterResourceListeners());

    Map<String, Set<PartitionInfo>> topicToTopicPartitions = new HashMap<>();
    for (PartitionInfo tp : partitions) {
      topicToTopicPartitions.putIfAbsent(tp.topic(), new HashSet<>());
      topicToTopicPartitions.get(tp.topic()).add(tp);
    }

    List<MetadataResponse.TopicMetadata> topicMetadata = new ArrayList<>(partitions.size());
    for (Map.Entry<String, Set<PartitionInfo>> entry : topicToTopicPartitions.entrySet()) {
      List<MetadataResponse.PartitionMetadata> partitionMetadata = new ArrayList<>(entry.getValue().size());
      for (PartitionInfo tp : entry.getValue()) {
        partitionMetadata.add(new MetadataResponse.PartitionMetadata(Errors.NONE,
                                                                     new TopicPartition(
                                                                             tp.topic(), tp.partition()), Optional.of(NODE_0.id()),
                                                                     Optional.of(RecordBatch.NO_PARTITION_LEADER_EPOCH),
                                                                     nodeIds(), nodeIds(),
                                                                     Collections.emptyList()));
      }
      topicMetadata.add(new MetadataResponse.TopicMetadata(Errors.NONE, entry.getKey(), false, partitionMetadata));
    }

    MetadataResponse metadataResponse = KafkaCruiseControlUtils.prepareMetadataResponse(cluster.nodes(),
                                                                                        cluster.clusterResource().clusterId(),
                                                                                        MetadataResponse.NO_CONTROLLER_ID,
                                                                                        topicMetadata);
    metadata.update(KafkaCruiseControlUtils.REQUEST_VERSION_UPDATE, metadataResponse, false, 0);

    MetricSamplerPartitionAssignor assignor = new DefaultMetricSamplerPartitionAssignor();
    Set<TopicPartition> assignment = assignor.assignPartitions(metadata.fetch());

    int maxAssignedNumPartitionsForFetcher = -1;
    int minAssignedNumPartitionsForFetcher = Integer.MAX_VALUE;
    int totalAssignedNumPartitions = 0;
    maxAssignedNumPartitionsForFetcher = Math.max(maxAssignedNumPartitionsForFetcher, assignment.size());
    minAssignedNumPartitionsForFetcher = Math.min(minAssignedNumPartitionsForFetcher, assignment.size());
    Set<TopicPartition> uniqueAssignedPartitions = new HashSet<>(assignment);
    totalAssignedNumPartitions += assignment.size();
    // Make sure all the partitions are assigned and there is no double assignment.
    assertEquals("Total assigned number of partitions should be " + totalNumPartitions,
        totalNumPartitions, totalAssignedNumPartitions);
    assertEquals("Total number of unique assigned partitions should be " + totalNumPartitions,
        totalNumPartitions, uniqueAssignedPartitions.size());

    int avgAssignedPartitionsPerFetcher = totalNumPartitions;
    assertTrue("In the worst case the max number of partitions assigned to a metric fetchers should not differ by "
               + "more than the partition number of the biggest topic, which is " + maxNumPartitionsForTopic,
        maxAssignedNumPartitionsForFetcher - avgAssignedPartitionsPerFetcher <= maxNumPartitionsForTopic);
    assertTrue("In the worst case the min number of partitions assigned to a metric fetchers should not differ by "
               + "more than the partition number of the biggest topic, which is " + maxNumPartitionsForTopic,
        avgAssignedPartitionsPerFetcher - minAssignedNumPartitionsForFetcher <= maxNumPartitionsForTopic);
  }