@SuppressWarnings("unchecked")
@Test
public void testPartitionAssignmentChange() throws Exception {

    StreamsConfig config = new StreamsConfig(configProps());
    TopologyBuilder builder = new TopologyBuilder();
    builder.addSource("source1", "topic1");
    builder.addSource("source2", "topic2");
    builder.addSource("source3", "topic3");
    builder.addProcessor("processor", new MockProcessorSupplier(), "source2", "source3");
    StreamThread thread = new StreamThread(builder, config, new MockClientSupplier(), applicationId, clientId, processId, new Metrics(), new SystemTime()) {
        @Override
        protected StreamTask createStreamTask(TaskId id, Collection<TopicPartition> partitionsForTask) {
            ProcessorTopology topology = builder.build("X", id.topicGroupId);
            return new TestStreamTask(id, applicationId, partitionsForTask, topology, consumer, producer, restoreConsumer, config);
        }
    };
    initPartitionGrouper(config, thread);
    ConsumerRebalanceListener rebalanceListener = thread.rebalanceListener;
    assertTrue(thread.tasks().isEmpty());
    List<TopicPartition> revokedPartitions;
    List<TopicPartition> assignedPartitions;
    Set<TopicPartition> expectedGroup1;
    Set<TopicPartition> expectedGroup2;
    revokedPartitions = Collections.emptyList();
    assignedPartitions = Collections.singletonList(t1p1);
    expectedGroup1 = new HashSet<>(Arrays.asList(t1p1));
    rebalanceListener.onPartitionsRevoked(revokedPartitions);
    rebalanceListener.onPartitionsAssigned(assignedPartitions);
    assertTrue(thread.tasks().containsKey(task1));
    assertEquals(expectedGroup1, thread.tasks().get(task1).partitions());
    assertEquals(1, thread.tasks().size());
    revokedPartitions = assignedPartitions;
    assignedPartitions = Collections.singletonList(t1p2);
    expectedGroup2 = new HashSet<>(Arrays.asList(t1p2));
    rebalanceListener.onPartitionsRevoked(revokedPartitions);
    rebalanceListener.onPartitionsAssigned(assignedPartitions);
    assertTrue(thread.tasks().containsKey(task2));
    assertEquals(expectedGroup2, thread.tasks().get(task2).partitions());
    assertEquals(1, thread.tasks().size());
    revokedPartitions = assignedPartitions;
    assignedPartitions = Arrays.asList(t1p1, t1p2);
    expectedGroup1 = new HashSet<>(Collections.singleton(t1p1));
    expectedGroup2 = new HashSet<>(Collections.singleton(t1p2));
    rebalanceListener.onPartitionsRevoked(revokedPartitions);
    rebalanceListener.onPartitionsAssigned(assignedPartitions);
    assertTrue(thread.tasks().containsKey(task1));
    assertTrue(thread.tasks().containsKey(task2));
    assertEquals(expectedGroup1, thread.tasks().get(task1).partitions());
    assertEquals(expectedGroup2, thread.tasks().get(task2).partitions());
    assertEquals(2, thread.tasks().size());
    revokedPartitions = assignedPartitions;
    assignedPartitions = Arrays.asList(t2p1, t2p2, t3p1, t3p2);
    expectedGroup1 = new HashSet<>(Arrays.asList(t2p1, t3p1));
    expectedGroup2 = new HashSet<>(Arrays.asList(t2p2, t3p2));
    rebalanceListener.onPartitionsRevoked(revokedPartitions);
    rebalanceListener.onPartitionsAssigned(assignedPartitions);
    assertTrue(thread.tasks().containsKey(task4));
    assertTrue(thread.tasks().containsKey(task5));
    assertEquals(expectedGroup1, thread.tasks().get(task4).partitions());
    assertEquals(expectedGroup2, thread.tasks().get(task5).partitions());
    assertEquals(2, thread.tasks().size());
    revokedPartitions = assignedPartitions;
    assignedPartitions = Arrays.asList(t1p1, t2p1, t3p1);
    expectedGroup1 = new HashSet<>(Arrays.asList(t1p1));
    expectedGroup2 = new HashSet<>(Arrays.asList(t2p1, t3p1));
    rebalanceListener.onPartitionsRevoked(revokedPartitions);
    rebalanceListener.onPartitionsAssigned(assignedPartitions);
    assertTrue(thread.tasks().containsKey(task1));
    assertTrue(thread.tasks().containsKey(task4));
    assertEquals(expectedGroup1, thread.tasks().get(task1).partitions());
    assertEquals(expectedGroup2, thread.tasks().get(task4).partitions());
    assertEquals(2, thread.tasks().size());
    revokedPartitions = assignedPartitions;
    assignedPartitions = Collections.emptyList();
    rebalanceListener.onPartitionsRevoked(revokedPartitions);
    rebalanceListener.onPartitionsAssigned(assignedPartitions);
    assertTrue(thread.tasks().isEmpty());
}