Source:
public void testUpdateAssignmentDeleteUnusedTasks() { String testCluster = "testUpdateAssignmentDeleteUnusedTasks"; String connectorType = "connectorType"; int totalTasks = 10; ZkClientInterceptingAdapter adapter = createInterceptingZkAdapter(testCluster); adapter.connect(); List<DatastreamTask> tasks1 = new ArrayList<>(); Set<DatastreamTask> unusedTaskSet = new HashSet<>(); for (int i = 0; i < totalTasks; i++) { DatastreamTaskImpl dsTask = new DatastreamTaskImpl(); dsTask.setId("task1" + i); String taskPrefix = "taskPrefix1" + i; dsTask.setTaskPrefix(taskPrefix); dsTask.setConnectorType(connectorType); dsTask.setZkAdapter(adapter); tasks1.add(dsTask); if (i % 2 == 0) { unusedTaskSet.add(dsTask); } } updateInstanceAssignment(adapter, adapter.getInstanceName(), tasks1); Assert.assertEquals(unusedTaskSet.size(), totalTasks / 2); adapter.cleanupDeadInstanceAssignmentsAndUnusedTasks(new ArrayList<>(), unusedTaskSet); Assert.assertEquals(unusedTaskSet.size(), 0); List<String> leftOverTasks = adapter.getZkClient().getChildren(KeyBuilder.connector(testCluster, connectorType)); Assert.assertEquals(leftOverTasks.size(), totalTasks / 2); Assert.assertFalse(adapter.getZkClient().exists(KeyBuilder.instance(testCluster, adapter.getInstanceName()))); adapter.disconnect(); }

Target:
public void testUpdateAssignmentDeleteUnusedTasks() { String testCluster = "testUpdateAssignmentDeleteUnusedTasks"; String connectorType = "connectorType"; int totalTasks = 10; ZkClientInterceptingAdapter adapter = createInterceptingZkAdapter(testCluster); adapter.connect(); List<DatastreamTask> tasks1 = new ArrayList<>(); List<DatastreamTask> newTasks1 = new ArrayList<>(); List<DatastreamTask> tasks2 = new ArrayList<>(); for (int i = 0; i < totalTasks; i++) { DatastreamTaskImpl dsTask = new DatastreamTaskImpl(); dsTask.setId("task1" + i); String taskPrefix = "taskPrefix1" + i; dsTask.setTaskPrefix(taskPrefix); dsTask.setConnectorType(connectorType); dsTask.setZkAdapter(adapter); if (i % 2 == 0) { tasks1.add(dsTask); newTasks1.add(dsTask); } else { tasks2.add(dsTask); } } Map<String, List<DatastreamTask>> oldAssignments = new HashMap<>(); oldAssignments.put(adapter.getInstanceName(), tasks1); oldAssignments.put("deadInstance-999", tasks2); adapter.updateAllAssignments(oldAssignments); Map<String, Set<DatastreamTask>> currentAssignments = adapter.getAllAssignedDatastreamTasks(); Assert.assertEquals(oldAssignments.keySet(), currentAssignments.keySet()); currentAssignments.forEach((k, v) -> Assert.assertEquals(v, new HashSet<>(oldAssignments.get(k)))); newTasks1.remove(0); Map<String, List<DatastreamTask>> newAssignments = new HashMap<>(); newAssignments.put(adapter.getInstanceName(), newTasks1); adapter.cleanUpDeadInstanceDataAndOtherUnusedTasks(currentAssignments, newAssignments, new ArrayList<>(newAssignments.keySet())); List<String> leftOverTasks = adapter.getZkClient().getChildren(KeyBuilder.connector(testCluster, connectorType)); Assert.assertEquals(leftOverTasks.size(), newTasks1.size()); Assert.assertTrue(adapter.getZkClient().exists(KeyBuilder.instance(testCluster, adapter.getInstanceName()))); Assert.assertFalse(adapter.getZkClient().exists(KeyBuilder.instance(testCluster, "deadInstance-999"))); adapter.disconnect(); }