Source:
public void testFetchPartitionChange() throws Exception { String clusterName = "testGroupIdAssignment"; Properties config = getDefaultConfig(Optional.empty()); config.put(AbstractKafkaConnector.IS_GROUP_ID_HASHING_ENABLED, Boolean.toString(true)); config.put(KafkaMirrorMakerConnector.PARTITION_FETCH_INTERVAL, "1000"); KafkaMirrorMakerConnector connector = new KafkaMirrorMakerConnector("MirrorMakerConnector", config, clusterName); String yummyTopic = "YummyPizza"; createTopic(_zkUtils, yummyTopic, 1); Datastream datastream = KafkaMirrorMakerConnectorTestUtils.createDatastream("pizzaStream", _broker, "\\w+Pizza"); DatastreamGroup group = new DatastreamGroup(ImmutableList.of(datastream)); List<DatastreamGroup> datastreamGroups1 = new ArrayList<>(); datastreamGroups1.add(group); AtomicInteger partitionChangeCalls = new AtomicInteger(0); connector.onPartitionChange(groupName -> { if (groupName.equals(group.getTaskPrefix())) { partitionChangeCalls.incrementAndGet(); } }); connector.start(null); connector.handleDatastream(datastreamGroups1); PollUtils.poll(() -> partitionChangeCalls.get() == 1, POLL_PERIOD_MS, POLL_TIMEOUT_MS); Map<String, Optional<DatastreamPartitionsMetadata>> partitionInfo = connector.getDatastreamPartitions(); Assert.assertEquals(partitionInfo.get(group.getTaskPrefix()).get().getDatastreamGroupName(), group.getTaskPrefix()); Assert.assertEquals(new HashSet<String>(partitionInfo.get(group.getTaskPrefix()).get().getPartitions()), ImmutableSet.of(yummyTopic + "-0")); String saltyTopic = "SaltyPizza"; createTopic(_zkUtils, saltyTopic, 2); PollUtils.poll(() -> partitionChangeCalls.get() == 2, POLL_PERIOD_MS, POLL_TIMEOUT_MS); partitionInfo = connector.getDatastreamPartitions(); Assert.assertEquals(new HashSet<String>(partitionInfo.get(group.getTaskPrefix()).get().getPartitions()), ImmutableSet.of(yummyTopic + "-0", saltyTopic + "-0", saltyTopic + "-1")); connector.stop(); }

Target:
public void testFetchPartitionChange() throws Exception { String clusterName = "testGroupIdAssignment"; Properties config = getDefaultConfig(Optional.empty()); config.put(AbstractKafkaConnector.IS_GROUP_ID_HASHING_ENABLED, Boolean.toString(true)); config.put(KafkaMirrorMakerConnector.PARTITION_FETCH_INTERVAL, "1000"); config.put(KafkaBasedConnectorConfig.ENABLE_PARTITION_ASSIGNMENT, Boolean.toString(true)); KafkaMirrorMakerConnector connector = new KafkaMirrorMakerConnector("MirrorMakerConnector", config, clusterName); String yummyTopic = "YummyPizza"; createTopic(_zkUtils, yummyTopic, 1); Datastream datastream = KafkaMirrorMakerConnectorTestUtils.createDatastream("pizzaStream", _broker, "\\w+Pizza"); DatastreamGroup group = new DatastreamGroup(ImmutableList.of(datastream)); List<DatastreamGroup> datastreamGroups1 = new ArrayList<>(); datastreamGroups1.add(group); AtomicInteger partitionChangeCalls = new AtomicInteger(0); connector.onPartitionChange(callbackGroup -> { if (callbackGroup.equals(group)) { partitionChangeCalls.incrementAndGet(); } }); connector.start(null); connector.handleDatastream(datastreamGroups1); Assert.assertTrue(PollUtils.poll(() -> partitionChangeCalls.get() == 1, POLL_PERIOD_MS, POLL_TIMEOUT_MS)); Map<String, Optional<DatastreamGroupPartitionsMetadata>> partitionInfo = connector.getDatastreamPartitions(); Assert.assertEquals(partitionInfo.get(group.getTaskPrefix()).get().getDatastreamGroup().getName(), group.getTaskPrefix()); Assert.assertEquals(new HashSet<String>(partitionInfo.get(group.getTaskPrefix()).get().getPartitions()), ImmutableSet.of(yummyTopic + "-0")); String saltyTopic = "SaltyPizza"; createTopic(_zkUtils, saltyTopic, 2); Assert.assertTrue(PollUtils.poll(() -> partitionChangeCalls.get() == 2, POLL_PERIOD_MS, POLL_TIMEOUT_MS)); partitionInfo = connector.getDatastreamPartitions(); Assert.assertEquals(new HashSet<String>(partitionInfo.get(group.getTaskPrefix()).get().getPartitions()), ImmutableSet.of(yummyTopic + "-0", saltyTopic + "-0", saltyTopic + "-1")); connector.stop(); }