@Test
  public void testStatefulSamzaApplicationShouldRedistributeInputPartitionsToCorrectTasksWhenMultipleInputStreamsAreExpanded() throws Exception {
    // Setup the two input kafka topics.
    String statefulInputKafkaTopic1 = String.format("test-input-topic-%s", UUID.randomUUID().toString());
    String statefulInputKafkaTopic2 = String.format("test-input-topic-%s", UUID.randomUUID().toString());
    createTopic(statefulInputKafkaTopic1, 32, 1);
    createTopic(statefulInputKafkaTopic2, 32, 1);

    // Generate configuration for the test.
    Map<String, String> configMap =
        buildStreamApplicationConfigMap(testStreamAppName, testStreamAppId, false, Optional.of("test-store-redistribute-partitions"));
    configMap.put(JobConfig.PROCESSOR_ID, PROCESSOR_IDS[0]);
    Config applicationConfig1 = new ApplicationConfig(new MapConfig(configMap));

    // Publish events into the input kafka topics.
    publishKafkaEvents(statefulInputKafkaTopic1, 0, NUM_KAFKA_EVENTS, PROCESSOR_IDS[0]);
    publishKafkaEvents(statefulInputKafkaTopic2, 0, NUM_KAFKA_EVENTS, PROCESSOR_IDS[0]);

    // Create and launch the StreamApplication from configuration.
    CountDownLatch kafkaEventsConsumedLatch1 = new CountDownLatch(NUM_KAFKA_EVENTS);
    CountDownLatch processedMessagesLatch1 = new CountDownLatch(1);

    ApplicationRunner appRunner1 = ApplicationRunners.getApplicationRunner(TestStreamApplication.getInstance(
        TEST_SYSTEM, ImmutableList.of(statefulInputKafkaTopic1, statefulInputKafkaTopic2), outputKafkaTopic, processedMessagesLatch1, null, kafkaEventsConsumedLatch1,
        applicationConfig1), applicationConfig1);

    executeRun(appRunner1, applicationConfig1);
    processedMessagesLatch1.await();
    kafkaEventsConsumedLatch1.await();

    // Generate the correct task assignments before the input stream expansion.
    Map<TaskName, Set<SystemStreamPartition>> expectedTaskAssignments = new HashMap<>();
    for (int partition = 0; partition < 32; ++partition) {
      TaskName taskName = new TaskName(String.format("Partition %d", partition));
      SystemStreamPartition systemStreamPartition1 = new SystemStreamPartition(TEST_SYSTEM, statefulInputKafkaTopic1, new Partition(partition));
      SystemStreamPartition systemStreamPartition2 = new SystemStreamPartition(TEST_SYSTEM, statefulInputKafkaTopic2, new Partition(partition));
      expectedTaskAssignments.put(taskName, ImmutableSet.of(systemStreamPartition1, systemStreamPartition2));
    }

    // Read the latest JobModel for validation.
    String jobModelVersion = zkUtils.getJobModelVersion();
    JobModel jobModel = JobModelUtil.readJobModel(jobModelVersion, zkMetadataStore);
    List<SystemStreamPartition> ssps = getSystemStreamPartitions(jobModel);

    // Validate that the input 64 partitions are present in JobModel.
    Assert.assertEquals(64, ssps.size());

    // Validate that the new JobModel has the expected task assignments before the input stream expansion.
    Map<TaskName, Set<SystemStreamPartition>> actualTaskAssignments = getTaskAssignments(jobModel);
    Assert.assertEquals(expectedTaskAssignments, actualTaskAssignments);

    // Increase the partition count of the input kafka topic1 to 64.
    increasePartitionsTo(statefulInputKafkaTopic1, 64);

    // Increase the partition count of the input kafka topic2 to 64.
    increasePartitionsTo(statefulInputKafkaTopic2, 64);

    // Wait for the JobModel version to change due to the increase in the input partition count.
    long jobModelWaitTimeInMillis = 10;
    while (true) {
      LOGGER.info("Waiting for new jobModel to be published");
      jobModelVersion = zkUtils.getJobModelVersion();
      jobModel = JobModelUtil.readJobModel(jobModelVersion, zkMetadataStore);
      ssps = getSystemStreamPartitions(jobModel);

      if (ssps.size() == 128) {
        break;
      }

      Thread.sleep(jobModelWaitTimeInMillis);
    }

    // Validate that the input partition count is 128 in the new JobModel.
    Assert.assertEquals(128, ssps.size());

    // Generate the correct task assignments after the input stream expansion.
    expectedTaskAssignments = new HashMap<>();
    for (int partition = 0; partition < 32; ++partition) {
      TaskName taskName = new TaskName(String.format("Partition %d", partition));
      SystemStreamPartition ssp1 = new SystemStreamPartition(TEST_SYSTEM, statefulInputKafkaTopic1, new Partition(partition));
      SystemStreamPartition expandedSSP1 = new SystemStreamPartition(TEST_SYSTEM, statefulInputKafkaTopic1, new Partition(partition + 32));
      SystemStreamPartition ssp2 = new SystemStreamPartition(TEST_SYSTEM, statefulInputKafkaTopic2, new Partition(partition));
      SystemStreamPartition expandedSSP2 = new SystemStreamPartition(TEST_SYSTEM, statefulInputKafkaTopic2, new Partition(partition + 32));
      expectedTaskAssignments.put(taskName, ImmutableSet.of(ssp1, expandedSSP1, ssp2, expandedSSP2));
    }

    // Validate that the new JobModel has the expected task assignments.
    actualTaskAssignments = getTaskAssignments(jobModel);
    Assert.assertEquals(expectedTaskAssignments, actualTaskAssignments);
    Assert.assertEquals(32, jobModel.getMaxChangeLogStreamPartitions());
  }
A. Create a input kafka topic: T1 with partition count set to 32.
B. Create a input kafka topic: T2 with partition count set to 32.
C. Create and launch a stateful samza application which consumes events from the kafka topics: T1 and T2.
D. Validate that the {@link JobModel} contains 32 {@link SystemStreamPartition}'s of T1 and 32 {@link SystemStreamPartition}'s of T2.
E. Increase the partition count of the input kafka topic: T1 to 64.
F. Increase the partition count of the input kafka topic: T2 to 64.
G. Validate that the new {@link JobModel} contains 64 {@link SystemStreamPartition}'s of T1, T2 and the input
SystemStreamPartitions are mapped to the correct task.
samza-test/src/test/java/org/apache/samza/test/processor/TestZkLocalApplicationRunner.java