Source:
public void testAddPartitions() { StickyPartitionAssignmentStrategy strategy = new StickyPartitionAssignmentStrategy(Optional.empty(), Optional.empty(), Optional.empty()); Set<DatastreamTask> taskSet = new HashSet<>(); List<DatastreamGroup> datastreams = generateDatastreams("ds", 1); Map<String, Set<DatastreamTask>> assignment = generateEmptyAssignment(datastreams, 1, 3, true); assignment.put("instance1", taskSet); List<String> partitions = ImmutableList.of("t-0", "t-1", "t1-0"); DatastreamGroupPartitionsMetadata partitionsMetadata = new DatastreamGroupPartitionsMetadata(datastreams.get(0), partitions); assignment = strategy.assignPartitions(assignment, partitionsMetadata); List<String> newPartitions = ImmutableList.of("t-0", "t-1", "t1-0", "t2-0", "t2-1", "t2-2"); DatastreamGroupPartitionsMetadata newPartitionsMetadata = new DatastreamGroupPartitionsMetadata(datastreams.get(0), newPartitions); Map<String, Set<DatastreamTask>> newAssignment = strategy.assignPartitions(assignment, newPartitionsMetadata); for (DatastreamTask task : newAssignment.get("instance0")) { Assert.assertEquals(task.getPartitionsV2().size(), 2); } Map<String, List<DatastreamTask>> taskToCleanup = strategy.getTasksToCleanUp(datastreams, newAssignment); Assert.assertEquals(taskToCleanup.size(), 0); Map<String, Set<DatastreamTask>> finalAssignment = assignment; newAssignment.forEach((instance, taskSet1) -> taskSet1.addAll(finalAssignment.get(instance))); taskToCleanup = strategy.getTasksToCleanUp(datastreams, newAssignment); Assert.assertEquals(taskToCleanup.size(), 1); taskToCleanup.forEach((instance, taskList1) -> Assert.assertEquals(taskList1.size(), 3)); }

Target:
public void testAddPartitions() { StickyPartitionAssignmentStrategy strategy = new StickyPartitionAssignmentStrategy(Optional.empty(), Optional.empty(), Optional.empty()); Set<DatastreamTask> taskSet = new HashSet<>(); List<DatastreamGroup> datastreams = generateDatastreams("ds", 1); Map<String, Set<DatastreamTask>> assignment = generateEmptyAssignment(datastreams, 1, 3, true); assignment.put("instance1", taskSet); List<String> partitions = ImmutableList.of("t-0", "t-1", "t1-0"); DatastreamGroupPartitionsMetadata partitionsMetadata = new DatastreamGroupPartitionsMetadata(datastreams.get(0), partitions); assignment = strategy.assignPartitions(assignment, partitionsMetadata); List<String> newPartitions = ImmutableList.of("t-0", "t-1", "t1-0", "t2-0", "t2-1", "t2-2"); DatastreamGroupPartitionsMetadata newPartitionsMetadata = new DatastreamGroupPartitionsMetadata(datastreams.get(0), newPartitions); Map<String, Set<DatastreamTask>> newAssignment = strategy.assignPartitions(assignment, newPartitionsMetadata); for (DatastreamTask task : newAssignment.get("instance0")) { Assert.assertEquals(task.getPartitionsV2().size(), 2); } Map<String, List<DatastreamTask>> taskToCleanup = strategy.getTasksToCleanUp(datastreams, newAssignment); Assert.assertEquals(taskToCleanup.size(), 0); Map<String, Set<DatastreamTask>> finalAssignment = assignment; newAssignment.forEach((instance, taskSet1) -> taskSet1.addAll(finalAssignment.get(instance))); taskToCleanup = strategy.getTasksToCleanUp(datastreams, newAssignment); Assert.assertEquals(taskToCleanup.size(), 1); taskToCleanup.forEach((instance, taskList1) -> Assert.assertEquals(taskList1.size(), 3)); Assert.assertEquals(new HashSet<>(taskToCleanup.get("instance0")), new HashSet<>(assignment.get("instance0"))); }