@Test
public void testGrouping() {

    PartitionGrouper grouper = new DefaultPartitionGrouper();
    int topicGroupId;
    Map<TaskId, Set<TopicPartition>> expectedPartitionsForTask;
    Map<Integer, Set<String>> topicGroups;
    topicGroups = new HashMap<>();
    topicGroups.put(0, mkSet("topic1"));
    topicGroups.put(1, mkSet("topic2"));
    expectedPartitionsForTask = new HashMap<>();
    topicGroupId = 0;
    expectedPartitionsForTask.put(new TaskId(topicGroupId, 0), mkSet(new TopicPartition("topic1", 0)));
    expectedPartitionsForTask.put(new TaskId(topicGroupId, 1), mkSet(new TopicPartition("topic1", 1)));
    expectedPartitionsForTask.put(new TaskId(topicGroupId, 2), mkSet(new TopicPartition("topic1", 2)));
    topicGroupId++;
    expectedPartitionsForTask.put(new TaskId(topicGroupId, 0), mkSet(new TopicPartition("topic2", 0)));
    expectedPartitionsForTask.put(new TaskId(topicGroupId, 1), mkSet(new TopicPartition("topic2", 1)));
    assertEquals(expectedPartitionsForTask, grouper.partitionGroups(topicGroups, metadata));
    topicGroups = new HashMap<>();
    topicGroups.put(0, mkSet("topic1", "topic2"));
    expectedPartitionsForTask = new HashMap<>();
    topicGroupId = 0;
    expectedPartitionsForTask.put(new TaskId(topicGroupId, 0), mkSet(new TopicPartition("topic1", 0), new TopicPartition("topic2", 0)));
    expectedPartitionsForTask.put(new TaskId(topicGroupId, 1), mkSet(new TopicPartition("topic1", 1), new TopicPartition("topic2", 1)));
    expectedPartitionsForTask.put(new TaskId(topicGroupId, 2), mkSet(new TopicPartition("topic1", 2)));
    assertEquals(expectedPartitionsForTask, grouper.partitionGroups(topicGroups, metadata));
}