org.apache.flink.streaming.connectors.kafka.internals.KafkaTopicPartitionStateSentinel Java Examples

/**
 * Tests reassignment works correctly in the case when:
 *  - the consumer has no initial assignments
 *  - new unassigned partitions have undefined offsets
 *  - the consumer was woken up prior to the reassignment
 *
 * <p>In this case, reassignment should not have occurred at all, and the consumer retains the original assignment.
 *
 * <p>Setting a timeout because the test will not finish if there is logic error with
 * the reassignment flow.
 */
@SuppressWarnings("unchecked")
@Test(timeout = 10000)
public void testReassignPartitionsDefinedOffsetsWithoutInitialAssignmentsWhenEarlyWakeup() throws Exception {
	final String testTopic = "test-topic";

	// -------- new partitions with defined offsets --------

	KafkaTopicPartitionState<TopicPartition> newPartition1 = new KafkaTopicPartitionState<>(
		new KafkaTopicPartition(testTopic, 0), new TopicPartition(testTopic, 0));
	newPartition1.setOffset(KafkaTopicPartitionStateSentinel.EARLIEST_OFFSET);

	KafkaTopicPartitionState<TopicPartition> newPartition2 = new KafkaTopicPartitionState<>(
		new KafkaTopicPartition(testTopic, 1), new TopicPartition(testTopic, 1));
	newPartition2.setOffset(KafkaTopicPartitionStateSentinel.EARLIEST_OFFSET);

	List<KafkaTopicPartitionState<TopicPartition>> newPartitions = new ArrayList<>(2);
	newPartitions.add(newPartition1);
	newPartitions.add(newPartition2);

	// -------- setup mock KafkaConsumer --------

	// no initial assignments
	final Map<TopicPartition, Long> mockConsumerAssignmentsAndPositions = new LinkedHashMap<>();

	// mock retrieved values that should replace the EARLIEST_OFFSET sentinels
	final Map<TopicPartition, Long> mockRetrievedPositions = new HashMap<>();
	mockRetrievedPositions.put(newPartition1.getKafkaPartitionHandle(), 23L);
	mockRetrievedPositions.put(newPartition2.getKafkaPartitionHandle(), 32L);

	final KafkaConsumer<byte[], byte[]> mockConsumer = createMockConsumer(
			mockConsumerAssignmentsAndPositions,
			mockRetrievedPositions,
			true,
			null,
			null);

	// -------- setup new partitions to be polled from the unassigned partitions queue --------

	final ClosableBlockingQueue<KafkaTopicPartitionState<TopicPartition>> unassignedPartitionsQueue =
		new ClosableBlockingQueue<>();

	for (KafkaTopicPartitionState<TopicPartition> newPartition : newPartitions) {
		unassignedPartitionsQueue.add(newPartition);
	}

	// -------- start test --------

	final TestKafkaConsumerThread testThread =
		new TestKafkaConsumerThread(mockConsumer, unassignedPartitionsQueue, new Handover());
	testThread.start();

	// pause just before the reassignment so we can inject the wakeup
	testThread.waitPartitionReassignmentInvoked();

	testThread.setOffsetsToCommit(new HashMap<TopicPartition, OffsetAndMetadata>(), mock(KafkaCommitCallback.class));

	// make sure the consumer was actually woken up
	verify(mockConsumer, times(1)).wakeup();

	testThread.startPartitionReassignment();
	testThread.waitPartitionReassignmentComplete();

	// the consumer's assignment should have remained untouched (in this case, empty)
	assertEquals(0, mockConsumerAssignmentsAndPositions.size());

	// the new partitions should have been re-added to the unassigned partitions queue
	assertEquals(2, unassignedPartitionsQueue.size());
}
 

/**
 * Test restoring from an empty state taken using a previous Flink version, when some partitions could be
 * found for topics.
 */
@Test
public void testRestoreFromEmptyStateWithPartitions() throws Exception {
	final List<KafkaTopicPartition> partitions = new ArrayList<>(PARTITION_STATE.keySet());

	final DummyFlinkKafkaConsumer<String> consumerFunction =
		new DummyFlinkKafkaConsumer<>(TOPICS, partitions, FlinkKafkaConsumerBase.PARTITION_DISCOVERY_DISABLED);

	StreamSource<String, DummyFlinkKafkaConsumer<String>> consumerOperator =
			new StreamSource<>(consumerFunction);

	final AbstractStreamOperatorTestHarness<String> testHarness =
			new AbstractStreamOperatorTestHarness<>(consumerOperator, 1, 1, 0);

	testHarness.setTimeCharacteristic(TimeCharacteristic.ProcessingTime);

	testHarness.setup();

	// restore state from binary snapshot file
	testHarness.initializeState(
		OperatorSnapshotUtil.getResourceFilename(
			"kafka-consumer-migration-test-flink" + testMigrateVersion + "-empty-state-snapshot"));

	testHarness.open();

	// the expected state in "kafka-consumer-migration-test-flink1.x-snapshot-empty-state";
	// all new partitions after the snapshot are considered as partitions that were created while the
	// consumer wasn't running, and should start from the earliest offset.
	final HashMap<KafkaTopicPartition, Long> expectedSubscribedPartitionsWithStartOffsets = new HashMap<>();
	for (KafkaTopicPartition partition : PARTITION_STATE.keySet()) {
		expectedSubscribedPartitionsWithStartOffsets.put(partition, KafkaTopicPartitionStateSentinel.EARLIEST_OFFSET);
	}

	// assert that there are partitions and is identical to expected list
	assertTrue(consumerFunction.getSubscribedPartitionsToStartOffsets() != null);
	assertTrue(!consumerFunction.getSubscribedPartitionsToStartOffsets().isEmpty());
	assertEquals(expectedSubscribedPartitionsWithStartOffsets, consumerFunction.getSubscribedPartitionsToStartOffsets());

	// the new partitions should have been considered as restored state
	assertTrue(consumerFunction.getRestoredState() != null);
	assertTrue(!consumerFunction.getSubscribedPartitionsToStartOffsets().isEmpty());
	for (Map.Entry<KafkaTopicPartition, Long> expectedEntry : expectedSubscribedPartitionsWithStartOffsets.entrySet()) {
		assertEquals(expectedEntry.getValue(), consumerFunction.getRestoredState().get(expectedEntry.getKey()));
	}

	consumerOperator.close();
	consumerOperator.cancel();
}
 

@Test
public void testCancellationWhenEmitBlocks() throws Exception {

	// ----- some test data -----

	final String topic = "test-topic";
	final int partition = 3;
	final byte[] payload = new byte[] {1, 2, 3, 4};

	final List<ConsumerRecord<byte[], byte[]>> records = Arrays.asList(
			new ConsumerRecord<>(topic, partition, 15, payload, payload),
			new ConsumerRecord<>(topic, partition, 16, payload, payload),
			new ConsumerRecord<>(topic, partition, 17, payload, payload));

	final Map<TopicPartition, List<ConsumerRecord<byte[], byte[]>>> data = new HashMap<>();
	data.put(new TopicPartition(topic, partition), records);

	final ConsumerRecords<byte[], byte[]> consumerRecords = new ConsumerRecords<>(data);

	// ----- the test consumer -----

	final KafkaConsumer<?, ?> mockConsumer = mock(KafkaConsumer.class);
	when(mockConsumer.poll(anyLong())).thenAnswer(new Answer<ConsumerRecords<?, ?>>() {
		@Override
		public ConsumerRecords<?, ?> answer(InvocationOnMock invocation) {
			return consumerRecords;
		}
	});

	whenNew(KafkaConsumer.class).withAnyArguments().thenReturn(mockConsumer);

	// ----- build a fetcher -----

	BlockingSourceContext<String> sourceContext = new BlockingSourceContext<>();
	Map<KafkaTopicPartition, Long> partitionsWithInitialOffsets =
		Collections.singletonMap(new KafkaTopicPartition(topic, partition), KafkaTopicPartitionStateSentinel.GROUP_OFFSET);
	KafkaDeserializationSchema<String> schema = new KafkaDeserializationSchemaWrapper<>(new SimpleStringSchema());

	final Kafka010Fetcher<String> fetcher = new Kafka010Fetcher<>(
			sourceContext,
			partitionsWithInitialOffsets,
			null, /* watermark strategy */
			new TestProcessingTimeService(),
			10, /* watermark interval */
			this.getClass().getClassLoader(),
			"task_name",
			schema,
			new Properties(),
			0L,
			new UnregisteredMetricsGroup(),
			new UnregisteredMetricsGroup(),
			false, null);

	// ----- run the fetcher -----

	final AtomicReference<Throwable> error = new AtomicReference<>();
	final Thread fetcherRunner = new Thread("fetcher runner") {

		@Override
		public void run() {
			try {
				fetcher.runFetchLoop();
			} catch (Throwable t) {
				error.set(t);
			}
		}
	};
	fetcherRunner.start();

	// wait until the thread started to emit records to the source context
	sourceContext.waitTillHasBlocker();

	// now we try to cancel the fetcher, including the interruption usually done on the task thread
	// once it has finished, there must be no more thread blocked on the source context
	fetcher.cancel();
	fetcherRunner.interrupt();
	fetcherRunner.join();

	assertFalse("fetcher threads did not properly finish", sourceContext.isStillBlocking());
}
 

@Test(timeout = 10000)
public void testRatelimiting() throws Exception {
	final String testTopic = "test-topic-ratelimit";

	// -------- setup mock KafkaConsumer with test data --------
	final int partition = 0;
	final byte[] payload = new byte[] {1};

	final List<ConsumerRecord<byte[], byte[]>> records = Arrays.asList(
			new ConsumerRecord<>(testTopic, partition, 15, payload, payload),
			new ConsumerRecord<>(testTopic, partition, 16, payload, payload));

	final Map<TopicPartition, List<ConsumerRecord<byte[], byte[]>>> data = new HashMap<>();
	data.put(new TopicPartition(testTopic, partition), records);

	final ConsumerRecords<byte[], byte[]> consumerRecords = new ConsumerRecords<>(data);

	// Sleep for one second in each consumer.poll() call to return 24 bytes / second
	final KafkaConsumer<byte[], byte[]> mockConsumer = mock(KafkaConsumer.class);
	PowerMockito.when(mockConsumer.poll(anyLong())).thenAnswer(
			invocationOnMock -> consumerRecords
	);

	whenNew(KafkaConsumer.class).withAnyArguments().thenReturn(mockConsumer);

	// -------- new partitions with defined offsets --------

	KafkaTopicPartitionState<Object, TopicPartition> newPartition1 = new KafkaTopicPartitionState<>(
			new KafkaTopicPartition(testTopic, 0), new TopicPartition(testTopic, 0));
	newPartition1.setOffset(KafkaTopicPartitionStateSentinel.EARLIEST_OFFSET);

	List<KafkaTopicPartitionState<Object, TopicPartition>> newPartitions = new ArrayList<>(1);
	newPartitions.add(newPartition1);

	final ClosableBlockingQueue<KafkaTopicPartitionState<Object, TopicPartition>> unassignedPartitionsQueue =
			new ClosableBlockingQueue<>();

	for (KafkaTopicPartitionState<Object, TopicPartition> newPartition : newPartitions) {
		unassignedPartitionsQueue.add(newPartition);
	}

	// --- ratelimiting properties ---
	StreamingRuntimeContext mockRuntimeContext = mock(StreamingRuntimeContext.class);
	when(mockRuntimeContext.getNumberOfParallelSubtasks()).thenReturn(1);
	Properties properties = new Properties();

	// -- mock Handover and logger ---
	Handover mockHandover = PowerMockito.mock(Handover.class);
	doNothing().when(mockHandover).produce(any());
	Logger mockLogger = mock(Logger.class);

	MetricGroup metricGroup = new UnregisteredMetricsGroup();
	FlinkConnectorRateLimiter rateLimiter = new GuavaFlinkConnectorRateLimiter();
	rateLimiter.setRate(1L);
	rateLimiter.open(mockRuntimeContext);

	// -- Test Kafka Consumer thread ---

	KafkaConsumerThread testThread = new TestKafkaConsumerThreadRateLimit(
			mockLogger,
			mockHandover,
			properties,
			unassignedPartitionsQueue,
			"test",
			30L,
			false,
			metricGroup,
			metricGroup,
			mockConsumer,
			rateLimiter
	);

	testThread.start();
	// Wait for 4 seconds to ensure atleast 2 calls to consumer.poll()
	testThread.join(5000);
	assertNotNull(testThread.getRateLimiter());
	assertEquals(testThread.getRateLimiter().getRate(), 1, 0);

	// In a period of 5 seconds, no more than 3 calls to poll should be made.
	// The expected rate is 1 byte / second and we read 4 bytes in every consumer.poll()
	// call. The rate limiter should thus slow down the call by 4 seconds when the rate takes
	// effect.
	verify(mockConsumer, times(3)).poll(anyLong());
	testThread.shutdown();

}
 

/**
 * Tests reassignment works correctly in the case when:
 *  - the consumer has no initial assignments
 *  - new unassigned partitions have undefined offsets
 *  - the consumer was woken up prior to the reassignment
 *
 * <p>In this case, reassignment should not have occurred at all, and the consumer retains the original assignment.
 *
 * <p>Setting a timeout because the test will not finish if there is logic error with
 * the reassignment flow.
 */
@SuppressWarnings("unchecked")
@Test(timeout = 10000)
public void testReassignPartitionsDefinedOffsetsWithoutInitialAssignmentsWhenEarlyWakeup() throws Exception {
	final String testTopic = "test-topic";

	// -------- new partitions with defined offsets --------

	KafkaTopicPartitionState<Object, TopicPartition> newPartition1 = new KafkaTopicPartitionState<>(
		new KafkaTopicPartition(testTopic, 0), new TopicPartition(testTopic, 0));
	newPartition1.setOffset(KafkaTopicPartitionStateSentinel.EARLIEST_OFFSET);

	KafkaTopicPartitionState<Object, TopicPartition> newPartition2 = new KafkaTopicPartitionState<>(
		new KafkaTopicPartition(testTopic, 1), new TopicPartition(testTopic, 1));
	newPartition2.setOffset(KafkaTopicPartitionStateSentinel.EARLIEST_OFFSET);

	List<KafkaTopicPartitionState<Object, TopicPartition>> newPartitions = new ArrayList<>(2);
	newPartitions.add(newPartition1);
	newPartitions.add(newPartition2);

	// -------- setup mock KafkaConsumer --------

	// no initial assignments
	final Map<TopicPartition, Long> mockConsumerAssignmentsAndPositions = new LinkedHashMap<>();

	// mock retrieved values that should replace the EARLIEST_OFFSET sentinels
	final Map<TopicPartition, Long> mockRetrievedPositions = new HashMap<>();
	mockRetrievedPositions.put(newPartition1.getKafkaPartitionHandle(), 23L);
	mockRetrievedPositions.put(newPartition2.getKafkaPartitionHandle(), 32L);

	final TestConsumer mockConsumer = createMockConsumer(
			mockConsumerAssignmentsAndPositions,
			mockRetrievedPositions,
			true,
			null,
			null);

	// -------- setup new partitions to be polled from the unassigned partitions queue --------

	final ClosableBlockingQueue<KafkaTopicPartitionState<Object, TopicPartition>> unassignedPartitionsQueue =
		new ClosableBlockingQueue<>();

	for (KafkaTopicPartitionState<Object, TopicPartition> newPartition : newPartitions) {
		unassignedPartitionsQueue.add(newPartition);
	}

	// -------- start test --------

	final TestKafkaConsumerThread testThread =
		new TestKafkaConsumerThread(mockConsumer, unassignedPartitionsQueue, new Handover());
	testThread.start();

	// pause just before the reassignment so we can inject the wakeup
	testThread.waitPartitionReassignmentInvoked();

	testThread.setOffsetsToCommit(new HashMap<TopicPartition, OffsetAndMetadata>(), mock(KafkaCommitCallback.class));

	// make sure the consumer was actually woken up
	assertEquals(1, mockConsumer.getNumWakeupCalls());

	testThread.startPartitionReassignment();
	testThread.waitPartitionReassignmentComplete();

	// the consumer's assignment should have remained untouched (in this case, empty)
	assertEquals(0, mockConsumerAssignmentsAndPositions.size());

	// the new partitions should have been re-added to the unassigned partitions queue
	assertEquals(2, unassignedPartitionsQueue.size());
}
 

/**
 * Tests reassignment works correctly in the case when:
 *  - the consumer already have some assignments
 *  - new unassigned partitions have undefined offsets (e.g. EARLIEST_OFFSET sentinel value)
 *
 * <p>Setting a timeout because the test will not finish if there is logic error with
 * the reassignment flow.
 */
@SuppressWarnings("unchecked")
@Test(timeout = 10000)
public void testReassigningPartitionsWithoutDefinedOffsets() throws Exception {
	final String testTopic = "test-topic";

	// -------- old partitions --------

	KafkaTopicPartitionState<Object, TopicPartition> oldPartition1 = new KafkaTopicPartitionState<>(
		new KafkaTopicPartition(testTopic, 0), new TopicPartition(testTopic, 0));
	oldPartition1.setOffset(23L);

	KafkaTopicPartitionState<Object, TopicPartition> oldPartition2 = new KafkaTopicPartitionState<>(
		new KafkaTopicPartition(testTopic, 1), new TopicPartition(testTopic, 1));
	oldPartition2.setOffset(32L);

	List<KafkaTopicPartitionState<Object, TopicPartition>> oldPartitions = new ArrayList<>(2);
	oldPartitions.add(oldPartition1);
	oldPartitions.add(oldPartition2);

	// -------- new partitions with undefined offsets --------

	KafkaTopicPartitionState<Object, TopicPartition> newPartition = new KafkaTopicPartitionState<>(
		new KafkaTopicPartition(testTopic, 2), new TopicPartition(testTopic, 2));
	newPartition.setOffset(KafkaTopicPartitionStateSentinel.EARLIEST_OFFSET);

	List<KafkaTopicPartitionState<Object, TopicPartition>> totalPartitions = new ArrayList<>(3);
	totalPartitions.add(oldPartition1);
	totalPartitions.add(oldPartition2);
	totalPartitions.add(newPartition);

	// -------- setup mock KafkaConsumer --------

	// has initial assignments
	final Map<TopicPartition, Long> mockConsumerAssignmentsAndPositions = new HashMap<>();
	for (KafkaTopicPartitionState<Object, TopicPartition> oldPartition : oldPartitions) {
		mockConsumerAssignmentsAndPositions.put(oldPartition.getKafkaPartitionHandle(), oldPartition.getOffset() + 1);
	}

	// mock retrieved values that should replace the EARLIEST_OFFSET sentinels
	final Map<TopicPartition, Long> mockRetrievedPositions = new HashMap<>();
	mockRetrievedPositions.put(newPartition.getKafkaPartitionHandle(), 30L);

	final Consumer<byte[], byte[]> mockConsumer = createMockConsumer(
			mockConsumerAssignmentsAndPositions,
			mockRetrievedPositions,
			false,
			null,
			null);

	// -------- setup new partitions to be polled from the unassigned partitions queue --------

	final ClosableBlockingQueue<KafkaTopicPartitionState<Object, TopicPartition>> unassignedPartitionsQueue =
		new ClosableBlockingQueue<>();

	unassignedPartitionsQueue.add(newPartition);

	// -------- start test --------

	final TestKafkaConsumerThread testThread =
		new TestKafkaConsumerThread(mockConsumer, unassignedPartitionsQueue, new Handover());
	testThread.start();

	testThread.startPartitionReassignment();
	testThread.waitPartitionReassignmentComplete();

	// the sentinel offset states should have been replaced with defined values according to the retrieved positions
	assertEquals(mockRetrievedPositions.get(newPartition.getKafkaPartitionHandle()) - 1, newPartition.getOffset());

	// verify that the consumer called assign() with all new partitions, and that positions are correctly advanced

	assertEquals(totalPartitions.size(), mockConsumerAssignmentsAndPositions.size());

	// old partitions should be re-seeked to their previous positions
	for (KafkaTopicPartitionState<Object, TopicPartition> partition : totalPartitions) {
		assertTrue(mockConsumerAssignmentsAndPositions.containsKey(partition.getKafkaPartitionHandle()));

		// should be seeked to (offset in state + 1) because offsets in state represent the last processed record
		assertEquals(
			partition.getOffset() + 1,
			mockConsumerAssignmentsAndPositions.get(partition.getKafkaPartitionHandle()).longValue());
	}

	assertEquals(0, unassignedPartitionsQueue.size());
}
 

/**
 * Test restoring from an empty state taken using a previous Flink version, when some partitions could be
 * found for topics.
 */
@Test
public void testRestoreFromEmptyStateWithPartitions() throws Exception {
	final List<KafkaTopicPartition> partitions = new ArrayList<>(PARTITION_STATE.keySet());

	final DummyFlinkKafkaConsumer<String> consumerFunction =
		new DummyFlinkKafkaConsumer<>(TOPICS, partitions, FlinkKafkaConsumerBase.PARTITION_DISCOVERY_DISABLED);

	StreamSource<String, DummyFlinkKafkaConsumer<String>> consumerOperator =
			new StreamSource<>(consumerFunction);

	final AbstractStreamOperatorTestHarness<String> testHarness =
			new AbstractStreamOperatorTestHarness<>(consumerOperator, 1, 1, 0);

	testHarness.setTimeCharacteristic(TimeCharacteristic.ProcessingTime);

	testHarness.setup();

	// restore state from binary snapshot file
	testHarness.initializeState(
		OperatorSnapshotUtil.getResourceFilename(
			"kafka-consumer-migration-test-flink" + testMigrateVersion + "-empty-state-snapshot"));

	testHarness.open();

	// the expected state in "kafka-consumer-migration-test-flink1.2-snapshot-empty-state";
	// all new partitions after the snapshot are considered as partitions that were created while the
	// consumer wasn't running, and should start from the earliest offset.
	final HashMap<KafkaTopicPartition, Long> expectedSubscribedPartitionsWithStartOffsets = new HashMap<>();
	for (KafkaTopicPartition partition : PARTITION_STATE.keySet()) {
		expectedSubscribedPartitionsWithStartOffsets.put(partition, KafkaTopicPartitionStateSentinel.EARLIEST_OFFSET);
	}

	// assert that there are partitions and is identical to expected list
	assertTrue(consumerFunction.getSubscribedPartitionsToStartOffsets() != null);
	assertTrue(!consumerFunction.getSubscribedPartitionsToStartOffsets().isEmpty());
	assertEquals(expectedSubscribedPartitionsWithStartOffsets, consumerFunction.getSubscribedPartitionsToStartOffsets());

	// the new partitions should have been considered as restored state
	assertTrue(consumerFunction.getRestoredState() != null);
	assertTrue(!consumerFunction.getSubscribedPartitionsToStartOffsets().isEmpty());
	for (Map.Entry<KafkaTopicPartition, Long> expectedEntry : expectedSubscribedPartitionsWithStartOffsets.entrySet()) {
		assertEquals(expectedEntry.getValue(), consumerFunction.getRestoredState().get(expectedEntry.getKey()));
	}

	consumerOperator.close();
	consumerOperator.cancel();
}
 

@Test
public void testCancellationWhenEmitBlocks() throws Exception {

	// ----- some test data -----

	final String topic = "test-topic";
	final int partition = 3;
	final byte[] payload = new byte[] {1, 2, 3, 4};

	final List<ConsumerRecord<byte[], byte[]>> records = Arrays.asList(
			new ConsumerRecord<>(topic, partition, 15, payload, payload),
			new ConsumerRecord<>(topic, partition, 16, payload, payload),
			new ConsumerRecord<>(topic, partition, 17, payload, payload));

	final Map<TopicPartition, List<ConsumerRecord<byte[], byte[]>>> data = new HashMap<>();
	data.put(new TopicPartition(topic, partition), records);

	final ConsumerRecords<byte[], byte[]> consumerRecords = new ConsumerRecords<>(data);

	// ----- the test consumer -----

	final KafkaConsumer<?, ?> mockConsumer = mock(KafkaConsumer.class);
	when(mockConsumer.poll(anyLong())).thenAnswer(new Answer<ConsumerRecords<?, ?>>() {
		@Override
		public ConsumerRecords<?, ?> answer(InvocationOnMock invocation) {
			return consumerRecords;
		}
	});

	whenNew(KafkaConsumer.class).withAnyArguments().thenReturn(mockConsumer);

	// ----- build a fetcher -----

	BlockingSourceContext<String> sourceContext = new BlockingSourceContext<>();
	Map<KafkaTopicPartition, Long> partitionsWithInitialOffsets =
		Collections.singletonMap(new KafkaTopicPartition(topic, partition), KafkaTopicPartitionStateSentinel.GROUP_OFFSET);
	KafkaDeserializationSchema<String> schema = new KafkaDeserializationSchemaWrapper<>(new SimpleStringSchema());

	final Kafka010Fetcher<String> fetcher = new Kafka010Fetcher<>(
			sourceContext,
			partitionsWithInitialOffsets,
			null, /* periodic watermark extractor */
			null, /* punctuated watermark extractor */
			new TestProcessingTimeService(),
			10, /* watermark interval */
			this.getClass().getClassLoader(),
			"task_name",
			schema,
			new Properties(),
			0L,
			new UnregisteredMetricsGroup(),
			new UnregisteredMetricsGroup(),
			false, null);

	// ----- run the fetcher -----

	final AtomicReference<Throwable> error = new AtomicReference<>();
	final Thread fetcherRunner = new Thread("fetcher runner") {

		@Override
		public void run() {
			try {
				fetcher.runFetchLoop();
			} catch (Throwable t) {
				error.set(t);
			}
		}
	};
	fetcherRunner.start();

	// wait until the thread started to emit records to the source context
	sourceContext.waitTillHasBlocker();

	// now we try to cancel the fetcher, including the interruption usually done on the task thread
	// once it has finished, there must be no more thread blocked on the source context
	fetcher.cancel();
	fetcherRunner.interrupt();
	fetcherRunner.join();

	assertFalse("fetcher threads did not properly finish", sourceContext.isStillBlocking());
}
 

@Test(timeout = 10000)
public void testRatelimiting() throws Exception {
	final String testTopic = "test-topic-ratelimit";

	// -------- setup mock KafkaConsumer with test data --------
	final int partition = 0;
	final byte[] payload = new byte[] {1};

	final List<ConsumerRecord<byte[], byte[]>> records = Arrays.asList(
			new ConsumerRecord<>(testTopic, partition, 15, payload, payload),
			new ConsumerRecord<>(testTopic, partition, 16, payload, payload));

	final Map<TopicPartition, List<ConsumerRecord<byte[], byte[]>>> data = new HashMap<>();
	data.put(new TopicPartition(testTopic, partition), records);

	final ConsumerRecords<byte[], byte[]> consumerRecords = new ConsumerRecords<>(data);

	// Sleep for one second in each consumer.poll() call to return 24 bytes / second
	final KafkaConsumer<byte[], byte[]> mockConsumer = mock(KafkaConsumer.class);
	PowerMockito.when(mockConsumer.poll(anyLong())).thenAnswer(
			invocationOnMock -> consumerRecords
	);

	whenNew(KafkaConsumer.class).withAnyArguments().thenReturn(mockConsumer);

	// -------- new partitions with defined offsets --------

	KafkaTopicPartitionState<TopicPartition> newPartition1 = new KafkaTopicPartitionState<>(
			new KafkaTopicPartition(testTopic, 0), new TopicPartition(testTopic, 0));
	newPartition1.setOffset(KafkaTopicPartitionStateSentinel.EARLIEST_OFFSET);

	List<KafkaTopicPartitionState<TopicPartition>> newPartitions = new ArrayList<>(1);
	newPartitions.add(newPartition1);

	final ClosableBlockingQueue<KafkaTopicPartitionState<TopicPartition>> unassignedPartitionsQueue =
			new ClosableBlockingQueue<>();

	for (KafkaTopicPartitionState<TopicPartition> newPartition : newPartitions) {
		unassignedPartitionsQueue.add(newPartition);
	}

	// --- ratelimiting properties ---
	StreamingRuntimeContext mockRuntimeContext = mock(StreamingRuntimeContext.class);
	when(mockRuntimeContext.getNumberOfParallelSubtasks()).thenReturn(1);
	Properties properties = new Properties();
	KafkaConsumerCallBridge09 mockBridge = mock(KafkaConsumerCallBridge09.class);

	// -- mock Handover and logger ---
	Handover mockHandover = PowerMockito.mock(Handover.class);
	doNothing().when(mockHandover).produce(any());
	Logger mockLogger = mock(Logger.class);

	MetricGroup metricGroup = new UnregisteredMetricsGroup();
	FlinkConnectorRateLimiter rateLimiter = new GuavaFlinkConnectorRateLimiter();
	rateLimiter.setRate(1L);
	rateLimiter.open(mockRuntimeContext);

	// -- Test Kafka Consumer thread ---

	KafkaConsumerThread testThread = new TestKafkaConsumerThreadRateLimit(
			mockLogger,
			mockHandover,
			properties,
			unassignedPartitionsQueue,
			mockBridge,
			"test",
			30L,
			false,
			metricGroup,
			metricGroup,
			mockConsumer,
			rateLimiter
	);

	testThread.start();
	// Wait for 4 seconds to ensure atleast 2 calls to consumer.poll()
	testThread.join(5000);
	assertNotNull(testThread.getRateLimiter());
	assertEquals(testThread.getRateLimiter().getRate(), 1, 0);

	// In a period of 5 seconds, no more than 3 calls to poll should be made.
	// The expected rate is 1 byte / second and we read 4 bytes in every consumer.poll()
	// call. The rate limiter should thus slow down the call by 4 seconds when the rate takes
	// effect.
	verify(mockConsumer, times(3)).poll(anyLong());
	testThread.shutdown();

}
 

@Test
public void testCancellationWhenEmitBlocks() throws Exception {

	// ----- some test data -----

	final String topic = "test-topic";
	final int partition = 3;
	final byte[] payload = new byte[] {1, 2, 3, 4};

	final List<ConsumerRecord<byte[], byte[]>> records = Arrays.asList(
			new ConsumerRecord<>(topic, partition, 15, payload, payload),
			new ConsumerRecord<>(topic, partition, 16, payload, payload),
			new ConsumerRecord<>(topic, partition, 17, payload, payload));

	final Map<TopicPartition, List<ConsumerRecord<byte[], byte[]>>> data = new HashMap<>();
	data.put(new TopicPartition(topic, partition), records);

	final ConsumerRecords<byte[], byte[]> consumerRecords = new ConsumerRecords<>(data);

	// ----- the test consumer -----

	final KafkaConsumer<?, ?> mockConsumer = mock(KafkaConsumer.class);
	when(mockConsumer.poll(anyLong())).thenAnswer(new Answer<ConsumerRecords<?, ?>>() {
		@Override
		public ConsumerRecords<?, ?> answer(InvocationOnMock invocation) {
			return consumerRecords;
		}
	});

	whenNew(KafkaConsumer.class).withAnyArguments().thenReturn(mockConsumer);

	// ----- build a fetcher -----

	BlockingSourceContext<String> sourceContext = new BlockingSourceContext<>();
	Map<KafkaTopicPartition, Long> partitionsWithInitialOffsets =
		Collections.singletonMap(new KafkaTopicPartition(topic, partition), KafkaTopicPartitionStateSentinel.GROUP_OFFSET);
	KafkaDeserializationSchema<String> schema = new KafkaDeserializationSchemaWrapper<>(new SimpleStringSchema());

	final Kafka09Fetcher<String> fetcher = new Kafka09Fetcher<>(
			sourceContext,
			partitionsWithInitialOffsets,
			null, /* periodic watermark extractor */
			null, /* punctuated watermark extractor */
			new TestProcessingTimeService(),
			10, /* watermark interval */
			this.getClass().getClassLoader(),
			"task_name",
			schema,
			new Properties(),
			0L,
			new UnregisteredMetricsGroup(),
			new UnregisteredMetricsGroup(),
			false, null);

	// ----- run the fetcher -----

	final AtomicReference<Throwable> error = new AtomicReference<>();
	final Thread fetcherRunner = new Thread("fetcher runner") {

		@Override
		public void run() {
			try {
				fetcher.runFetchLoop();
			} catch (Throwable t) {
				error.set(t);
			}
		}
	};
	fetcherRunner.start();

	// wait until the thread started to emit records to the source context
	sourceContext.waitTillHasBlocker();

	// now we try to cancel the fetcher, including the interruption usually done on the task thread
	// once it has finished, there must be no more thread blocked on the source context
	fetcher.cancel();
	fetcherRunner.interrupt();
	fetcherRunner.join();

	assertFalse("fetcher threads did not properly finish", sourceContext.isStillBlocking());
}
 

/**
 * Tests reassignment works correctly in the case when:
 *  - the consumer already have some assignments
 *  - new unassigned partitions have undefined offsets (e.g. EARLIEST_OFFSET sentinel value)
 *
 * <p>Setting a timeout because the test will not finish if there is logic error with
 * the reassignment flow.
 */
@SuppressWarnings("unchecked")
@Test(timeout = 10000)
public void testReassigningPartitionsWithoutDefinedOffsets() throws Exception {
	final String testTopic = "test-topic";

	// -------- old partitions --------

	KafkaTopicPartitionState<TopicPartition> oldPartition1 = new KafkaTopicPartitionState<>(
		new KafkaTopicPartition(testTopic, 0), new TopicPartition(testTopic, 0));
	oldPartition1.setOffset(23L);

	KafkaTopicPartitionState<TopicPartition> oldPartition2 = new KafkaTopicPartitionState<>(
		new KafkaTopicPartition(testTopic, 1), new TopicPartition(testTopic, 1));
	oldPartition2.setOffset(32L);

	List<KafkaTopicPartitionState<TopicPartition>> oldPartitions = new ArrayList<>(2);
	oldPartitions.add(oldPartition1);
	oldPartitions.add(oldPartition2);

	// -------- new partitions with undefined offsets --------

	KafkaTopicPartitionState<TopicPartition> newPartition = new KafkaTopicPartitionState<>(
		new KafkaTopicPartition(testTopic, 2), new TopicPartition(testTopic, 2));
	newPartition.setOffset(KafkaTopicPartitionStateSentinel.EARLIEST_OFFSET);

	List<KafkaTopicPartitionState<TopicPartition>> totalPartitions = new ArrayList<>(3);
	totalPartitions.add(oldPartition1);
	totalPartitions.add(oldPartition2);
	totalPartitions.add(newPartition);

	// -------- setup mock KafkaConsumer --------

	// has initial assignments
	final Map<TopicPartition, Long> mockConsumerAssignmentsAndPositions = new HashMap<>();
	for (KafkaTopicPartitionState<TopicPartition> oldPartition : oldPartitions) {
		mockConsumerAssignmentsAndPositions.put(oldPartition.getKafkaPartitionHandle(), oldPartition.getOffset() + 1);
	}

	// mock retrieved values that should replace the EARLIEST_OFFSET sentinels
	final Map<TopicPartition, Long> mockRetrievedPositions = new HashMap<>();
	mockRetrievedPositions.put(newPartition.getKafkaPartitionHandle(), 30L);

	final KafkaConsumer<byte[], byte[]> mockConsumer = createMockConsumer(
			mockConsumerAssignmentsAndPositions,
			mockRetrievedPositions,
			false,
			null,
			null);

	// -------- setup new partitions to be polled from the unassigned partitions queue --------

	final ClosableBlockingQueue<KafkaTopicPartitionState<TopicPartition>> unassignedPartitionsQueue =
		new ClosableBlockingQueue<>();

	unassignedPartitionsQueue.add(newPartition);

	// -------- start test --------

	final TestKafkaConsumerThread testThread =
		new TestKafkaConsumerThread(mockConsumer, unassignedPartitionsQueue, new Handover());
	testThread.start();

	testThread.startPartitionReassignment();
	testThread.waitPartitionReassignmentComplete();

	// the sentinel offset states should have been replaced with defined values according to the retrieved positions
	assertEquals(mockRetrievedPositions.get(newPartition.getKafkaPartitionHandle()) - 1, newPartition.getOffset());

	// verify that the consumer called assign() with all new partitions, and that positions are correctly advanced

	assertEquals(totalPartitions.size(), mockConsumerAssignmentsAndPositions.size());

	// old partitions should be re-seeked to their previous positions
	for (KafkaTopicPartitionState<TopicPartition> partition : totalPartitions) {
		assertTrue(mockConsumerAssignmentsAndPositions.containsKey(partition.getKafkaPartitionHandle()));

		// should be seeked to (offset in state + 1) because offsets in state represent the last processed record
		assertEquals(
			partition.getOffset() + 1,
			mockConsumerAssignmentsAndPositions.get(partition.getKafkaPartitionHandle()).longValue());
	}

	assertEquals(0, unassignedPartitionsQueue.size());
}
 

@Test
public void testCancellationWhenEmitBlocks() throws Exception {

	// ----- some test data -----

	final String topic = "test-topic";
	final int partition = 3;
	final byte[] payload = new byte[] {1, 2, 3, 4};

	final List<ConsumerRecord<byte[], byte[]>> records = Arrays.asList(
			new ConsumerRecord<>(topic, partition, 15, payload, payload),
			new ConsumerRecord<>(topic, partition, 16, payload, payload),
			new ConsumerRecord<>(topic, partition, 17, payload, payload));

	final Map<TopicPartition, List<ConsumerRecord<byte[], byte[]>>> data = new HashMap<>();
	data.put(new TopicPartition(topic, partition), records);

	final ConsumerRecords<byte[], byte[]> consumerRecords = new ConsumerRecords<>(data);

	// ----- the test consumer -----

	final KafkaConsumer<?, ?> mockConsumer = mock(KafkaConsumer.class);
	when(mockConsumer.poll(anyLong())).thenAnswer(new Answer<ConsumerRecords<?, ?>>() {
		@Override
		public ConsumerRecords<?, ?> answer(InvocationOnMock invocation) {
			return consumerRecords;
		}
	});

	whenNew(KafkaConsumer.class).withAnyArguments().thenReturn(mockConsumer);

	// ----- build a fetcher -----

	BlockingSourceContext<String> sourceContext = new BlockingSourceContext<>();
	Map<KafkaTopicPartition, Long> partitionsWithInitialOffsets =
		Collections.singletonMap(new KafkaTopicPartition(topic, partition), KafkaTopicPartitionStateSentinel.GROUP_OFFSET);
	KafkaDeserializationSchema<String> schema = new KafkaDeserializationSchemaWrapper<>(new SimpleStringSchema());

	final Kafka09Fetcher<String> fetcher = new Kafka09Fetcher<>(
			sourceContext,
			partitionsWithInitialOffsets,
			null, /* periodic watermark extractor */
			null, /* punctuated watermark extractor */
			new TestProcessingTimeService(),
			10, /* watermark interval */
			this.getClass().getClassLoader(),
			"task_name",
			schema,
			new Properties(),
			0L,
			new UnregisteredMetricsGroup(),
			new UnregisteredMetricsGroup(),
			false, null);

	// ----- run the fetcher -----

	final AtomicReference<Throwable> error = new AtomicReference<>();
	final Thread fetcherRunner = new Thread("fetcher runner") {

		@Override
		public void run() {
			try {
				fetcher.runFetchLoop();
			} catch (Throwable t) {
				error.set(t);
			}
		}
	};
	fetcherRunner.start();

	// wait until the thread started to emit records to the source context
	sourceContext.waitTillHasBlocker();

	// now we try to cancel the fetcher, including the interruption usually done on the task thread
	// once it has finished, there must be no more thread blocked on the source context
	fetcher.cancel();
	fetcherRunner.interrupt();
	fetcherRunner.join();

	assertFalse("fetcher threads did not properly finish", sourceContext.isStillBlocking());
}
 

/**
 * Test restoring from an empty state taken using a previous Flink version, when some partitions could be
 * found for topics.
 */
@Test
public void testRestoreFromEmptyStateWithPartitions() throws Exception {
	final List<KafkaTopicPartition> partitions = new ArrayList<>(PARTITION_STATE.keySet());

	final DummyFlinkKafkaConsumer<String> consumerFunction =
		new DummyFlinkKafkaConsumer<>(TOPICS, partitions, FlinkKafkaConsumerBase.PARTITION_DISCOVERY_DISABLED);

	StreamSource<String, DummyFlinkKafkaConsumer<String>> consumerOperator =
			new StreamSource<>(consumerFunction);

	final AbstractStreamOperatorTestHarness<String> testHarness =
			new AbstractStreamOperatorTestHarness<>(consumerOperator, 1, 1, 0);

	testHarness.setTimeCharacteristic(TimeCharacteristic.ProcessingTime);

	testHarness.setup();

	// restore state from binary snapshot file
	testHarness.initializeState(
		OperatorSnapshotUtil.getResourceFilename(
			"kafka-consumer-migration-test-flink" + testMigrateVersion + "-empty-state-snapshot"));

	testHarness.open();

	// the expected state in "kafka-consumer-migration-test-flink1.2-snapshot-empty-state";
	// all new partitions after the snapshot are considered as partitions that were created while the
	// consumer wasn't running, and should start from the earliest offset.
	final HashMap<KafkaTopicPartition, Long> expectedSubscribedPartitionsWithStartOffsets = new HashMap<>();
	for (KafkaTopicPartition partition : PARTITION_STATE.keySet()) {
		expectedSubscribedPartitionsWithStartOffsets.put(partition, KafkaTopicPartitionStateSentinel.EARLIEST_OFFSET);
	}

	// assert that there are partitions and is identical to expected list
	assertTrue(consumerFunction.getSubscribedPartitionsToStartOffsets() != null);
	assertTrue(!consumerFunction.getSubscribedPartitionsToStartOffsets().isEmpty());
	assertEquals(expectedSubscribedPartitionsWithStartOffsets, consumerFunction.getSubscribedPartitionsToStartOffsets());

	// the new partitions should have been considered as restored state
	assertTrue(consumerFunction.getRestoredState() != null);
	assertTrue(!consumerFunction.getSubscribedPartitionsToStartOffsets().isEmpty());
	for (Map.Entry<KafkaTopicPartition, Long> expectedEntry : expectedSubscribedPartitionsWithStartOffsets.entrySet()) {
		assertEquals(expectedEntry.getValue(), consumerFunction.getRestoredState().get(expectedEntry.getKey()));
	}

	consumerOperator.close();
	consumerOperator.cancel();
}
 

@Test
public void testCancellationWhenEmitBlocks() throws Exception {

	// ----- some test data -----

	final String topic = "test-topic";
	final int partition = 3;
	final byte[] payload = new byte[] {1, 2, 3, 4};

	final List<ConsumerRecord<byte[], byte[]>> records = Arrays.asList(
			new ConsumerRecord<>(topic, partition, 15, payload, payload),
			new ConsumerRecord<>(topic, partition, 16, payload, payload),
			new ConsumerRecord<>(topic, partition, 17, payload, payload));

	final Map<TopicPartition, List<ConsumerRecord<byte[], byte[]>>> data = new HashMap<>();
	data.put(new TopicPartition(topic, partition), records);

	final ConsumerRecords<byte[], byte[]> consumerRecords = new ConsumerRecords<>(data);

	// ----- the test consumer -----

	final KafkaConsumer<?, ?> mockConsumer = mock(KafkaConsumer.class);
	when(mockConsumer.poll(anyLong())).thenAnswer(new Answer<ConsumerRecords<?, ?>>() {
		@Override
		public ConsumerRecords<?, ?> answer(InvocationOnMock invocation) {
			return consumerRecords;
		}
	});

	whenNew(KafkaConsumer.class).withAnyArguments().thenReturn(mockConsumer);

	// ----- build a fetcher -----

	BlockingSourceContext<String> sourceContext = new BlockingSourceContext<>();
	Map<KafkaTopicPartition, Long> partitionsWithInitialOffsets =
		Collections.singletonMap(new KafkaTopicPartition(topic, partition), KafkaTopicPartitionStateSentinel.GROUP_OFFSET);
	KafkaDeserializationSchema<String> schema = new KafkaDeserializationSchemaWrapper<>(new SimpleStringSchema());

	final Kafka010Fetcher<String> fetcher = new Kafka010Fetcher<>(
			sourceContext,
			partitionsWithInitialOffsets,
			null, /* periodic watermark extractor */
			null, /* punctuated watermark extractor */
			new TestProcessingTimeService(),
			10, /* watermark interval */
			this.getClass().getClassLoader(),
			"task_name",
			schema,
			new Properties(),
			0L,
			new UnregisteredMetricsGroup(),
			new UnregisteredMetricsGroup(),
			false, null);

	// ----- run the fetcher -----

	final AtomicReference<Throwable> error = new AtomicReference<>();
	final Thread fetcherRunner = new Thread("fetcher runner") {

		@Override
		public void run() {
			try {
				fetcher.runFetchLoop();
			} catch (Throwable t) {
				error.set(t);
			}
		}
	};
	fetcherRunner.start();

	// wait until the thread started to emit records to the source context
	sourceContext.waitTillHasBlocker();

	// now we try to cancel the fetcher, including the interruption usually done on the task thread
	// once it has finished, there must be no more thread blocked on the source context
	fetcher.cancel();
	fetcherRunner.interrupt();
	fetcherRunner.join();

	assertFalse("fetcher threads did not properly finish", sourceContext.isStillBlocking());
}
 

@Test(timeout = 10000)
public void testRatelimiting() throws Exception {
	final String testTopic = "test-topic-ratelimit";

	// -------- setup mock KafkaConsumer with test data --------
	final int partition = 0;
	final byte[] payload = new byte[] {1};

	final List<ConsumerRecord<byte[], byte[]>> records = Arrays.asList(
			new ConsumerRecord<>(testTopic, partition, 15, payload, payload),
			new ConsumerRecord<>(testTopic, partition, 16, payload, payload));

	final Map<TopicPartition, List<ConsumerRecord<byte[], byte[]>>> data = new HashMap<>();
	data.put(new TopicPartition(testTopic, partition), records);

	final ConsumerRecords<byte[], byte[]> consumerRecords = new ConsumerRecords<>(data);

	// Sleep for one second in each consumer.poll() call to return 24 bytes / second
	final KafkaConsumer<byte[], byte[]> mockConsumer = mock(KafkaConsumer.class);
	PowerMockito.when(mockConsumer.poll(anyLong())).thenAnswer(
			invocationOnMock -> consumerRecords
	);

	whenNew(KafkaConsumer.class).withAnyArguments().thenReturn(mockConsumer);

	// -------- new partitions with defined offsets --------

	KafkaTopicPartitionState<TopicPartition> newPartition1 = new KafkaTopicPartitionState<>(
			new KafkaTopicPartition(testTopic, 0), new TopicPartition(testTopic, 0));
	newPartition1.setOffset(KafkaTopicPartitionStateSentinel.EARLIEST_OFFSET);

	List<KafkaTopicPartitionState<TopicPartition>> newPartitions = new ArrayList<>(1);
	newPartitions.add(newPartition1);

	final ClosableBlockingQueue<KafkaTopicPartitionState<TopicPartition>> unassignedPartitionsQueue =
			new ClosableBlockingQueue<>();

	for (KafkaTopicPartitionState<TopicPartition> newPartition : newPartitions) {
		unassignedPartitionsQueue.add(newPartition);
	}

	// --- ratelimiting properties ---
	StreamingRuntimeContext mockRuntimeContext = mock(StreamingRuntimeContext.class);
	when(mockRuntimeContext.getNumberOfParallelSubtasks()).thenReturn(1);
	Properties properties = new Properties();
	KafkaConsumerCallBridge09 mockBridge = mock(KafkaConsumerCallBridge09.class);

	// -- mock Handover and logger ---
	Handover mockHandover = PowerMockito.mock(Handover.class);
	doNothing().when(mockHandover).produce(any());
	Logger mockLogger = mock(Logger.class);

	MetricGroup metricGroup = new UnregisteredMetricsGroup();
	FlinkConnectorRateLimiter rateLimiter = new GuavaFlinkConnectorRateLimiter();
	rateLimiter.setRate(1L);
	rateLimiter.open(mockRuntimeContext);

	// -- Test Kafka Consumer thread ---

	KafkaConsumerThread testThread = new TestKafkaConsumerThreadRateLimit(
			mockLogger,
			mockHandover,
			properties,
			unassignedPartitionsQueue,
			mockBridge,
			"test",
			30L,
			false,
			metricGroup,
			metricGroup,
			mockConsumer,
			rateLimiter
	);

	testThread.start();
	// Wait for 4 seconds to ensure atleast 2 calls to consumer.poll()
	testThread.join(5000);
	assertNotNull(testThread.getRateLimiter());
	assertEquals(testThread.getRateLimiter().getRate(), 1, 0);

	// In a period of 5 seconds, no more than 3 calls to poll should be made.
	// The expected rate is 1 byte / second and we read 4 bytes in every consumer.poll()
	// call. The rate limiter should thus slow down the call by 4 seconds when the rate takes
	// effect.
	verify(mockConsumer, times(3)).poll(anyLong());
	testThread.shutdown();

}
 

/**
 * Tests reassignment works correctly in the case when:
 *  - the consumer has no initial assignments
 *  - new unassigned partitions have undefined offsets
 *  - the consumer was woken up prior to the reassignment
 *
 * <p>In this case, reassignment should not have occurred at all, and the consumer retains the original assignment.
 *
 * <p>Setting a timeout because the test will not finish if there is logic error with
 * the reassignment flow.
 */
@SuppressWarnings("unchecked")
@Test(timeout = 10000)
public void testReassignPartitionsDefinedOffsetsWithoutInitialAssignmentsWhenEarlyWakeup() throws Exception {
	final String testTopic = "test-topic";

	// -------- new partitions with defined offsets --------

	KafkaTopicPartitionState<TopicPartition> newPartition1 = new KafkaTopicPartitionState<>(
		new KafkaTopicPartition(testTopic, 0), new TopicPartition(testTopic, 0));
	newPartition1.setOffset(KafkaTopicPartitionStateSentinel.EARLIEST_OFFSET);

	KafkaTopicPartitionState<TopicPartition> newPartition2 = new KafkaTopicPartitionState<>(
		new KafkaTopicPartition(testTopic, 1), new TopicPartition(testTopic, 1));
	newPartition2.setOffset(KafkaTopicPartitionStateSentinel.EARLIEST_OFFSET);

	List<KafkaTopicPartitionState<TopicPartition>> newPartitions = new ArrayList<>(2);
	newPartitions.add(newPartition1);
	newPartitions.add(newPartition2);

	// -------- setup mock KafkaConsumer --------

	// no initial assignments
	final Map<TopicPartition, Long> mockConsumerAssignmentsAndPositions = new LinkedHashMap<>();

	// mock retrieved values that should replace the EARLIEST_OFFSET sentinels
	final Map<TopicPartition, Long> mockRetrievedPositions = new HashMap<>();
	mockRetrievedPositions.put(newPartition1.getKafkaPartitionHandle(), 23L);
	mockRetrievedPositions.put(newPartition2.getKafkaPartitionHandle(), 32L);

	final KafkaConsumer<byte[], byte[]> mockConsumer = createMockConsumer(
			mockConsumerAssignmentsAndPositions,
			mockRetrievedPositions,
			true,
			null,
			null);

	// -------- setup new partitions to be polled from the unassigned partitions queue --------

	final ClosableBlockingQueue<KafkaTopicPartitionState<TopicPartition>> unassignedPartitionsQueue =
		new ClosableBlockingQueue<>();

	for (KafkaTopicPartitionState<TopicPartition> newPartition : newPartitions) {
		unassignedPartitionsQueue.add(newPartition);
	}

	// -------- start test --------

	final TestKafkaConsumerThread testThread =
		new TestKafkaConsumerThread(mockConsumer, unassignedPartitionsQueue, new Handover());
	testThread.start();

	// pause just before the reassignment so we can inject the wakeup
	testThread.waitPartitionReassignmentInvoked();

	testThread.setOffsetsToCommit(new HashMap<TopicPartition, OffsetAndMetadata>(), mock(KafkaCommitCallback.class));

	// make sure the consumer was actually woken up
	verify(mockConsumer, times(1)).wakeup();

	testThread.startPartitionReassignment();
	testThread.waitPartitionReassignmentComplete();

	// the consumer's assignment should have remained untouched (in this case, empty)
	assertEquals(0, mockConsumerAssignmentsAndPositions.size());

	// the new partitions should have been re-added to the unassigned partitions queue
	assertEquals(2, unassignedPartitionsQueue.size());
}
 

/**
 * Tests reassignment works correctly in the case when:
 *  - the consumer already have some assignments
 *  - new unassigned partitions have undefined offsets (e.g. EARLIEST_OFFSET sentinel value)
 *
 * <p>Setting a timeout because the test will not finish if there is logic error with
 * the reassignment flow.
 */
@SuppressWarnings("unchecked")
@Test(timeout = 10000)
public void testReassigningPartitionsWithoutDefinedOffsets() throws Exception {
	final String testTopic = "test-topic";

	// -------- old partitions --------

	KafkaTopicPartitionState<TopicPartition> oldPartition1 = new KafkaTopicPartitionState<>(
		new KafkaTopicPartition(testTopic, 0), new TopicPartition(testTopic, 0));
	oldPartition1.setOffset(23L);

	KafkaTopicPartitionState<TopicPartition> oldPartition2 = new KafkaTopicPartitionState<>(
		new KafkaTopicPartition(testTopic, 1), new TopicPartition(testTopic, 1));
	oldPartition2.setOffset(32L);

	List<KafkaTopicPartitionState<TopicPartition>> oldPartitions = new ArrayList<>(2);
	oldPartitions.add(oldPartition1);
	oldPartitions.add(oldPartition2);

	// -------- new partitions with undefined offsets --------

	KafkaTopicPartitionState<TopicPartition> newPartition = new KafkaTopicPartitionState<>(
		new KafkaTopicPartition(testTopic, 2), new TopicPartition(testTopic, 2));
	newPartition.setOffset(KafkaTopicPartitionStateSentinel.EARLIEST_OFFSET);

	List<KafkaTopicPartitionState<TopicPartition>> totalPartitions = new ArrayList<>(3);
	totalPartitions.add(oldPartition1);
	totalPartitions.add(oldPartition2);
	totalPartitions.add(newPartition);

	// -------- setup mock KafkaConsumer --------

	// has initial assignments
	final Map<TopicPartition, Long> mockConsumerAssignmentsAndPositions = new HashMap<>();
	for (KafkaTopicPartitionState<TopicPartition> oldPartition : oldPartitions) {
		mockConsumerAssignmentsAndPositions.put(oldPartition.getKafkaPartitionHandle(), oldPartition.getOffset() + 1);
	}

	// mock retrieved values that should replace the EARLIEST_OFFSET sentinels
	final Map<TopicPartition, Long> mockRetrievedPositions = new HashMap<>();
	mockRetrievedPositions.put(newPartition.getKafkaPartitionHandle(), 30L);

	final KafkaConsumer<byte[], byte[]> mockConsumer = createMockConsumer(
			mockConsumerAssignmentsAndPositions,
			mockRetrievedPositions,
			false,
			null,
			null);

	// -------- setup new partitions to be polled from the unassigned partitions queue --------

	final ClosableBlockingQueue<KafkaTopicPartitionState<TopicPartition>> unassignedPartitionsQueue =
		new ClosableBlockingQueue<>();

	unassignedPartitionsQueue.add(newPartition);

	// -------- start test --------

	final TestKafkaConsumerThread testThread =
		new TestKafkaConsumerThread(mockConsumer, unassignedPartitionsQueue, new Handover());
	testThread.start();

	testThread.startPartitionReassignment();
	testThread.waitPartitionReassignmentComplete();

	// the sentinel offset states should have been replaced with defined values according to the retrieved positions
	assertEquals(mockRetrievedPositions.get(newPartition.getKafkaPartitionHandle()) - 1, newPartition.getOffset());

	// verify that the consumer called assign() with all new partitions, and that positions are correctly advanced

	assertEquals(totalPartitions.size(), mockConsumerAssignmentsAndPositions.size());

	// old partitions should be re-seeked to their previous positions
	for (KafkaTopicPartitionState<TopicPartition> partition : totalPartitions) {
		assertTrue(mockConsumerAssignmentsAndPositions.containsKey(partition.getKafkaPartitionHandle()));

		// should be seeked to (offset in state + 1) because offsets in state represent the last processed record
		assertEquals(
			partition.getOffset() + 1,
			mockConsumerAssignmentsAndPositions.get(partition.getKafkaPartitionHandle()).longValue());
	}

	assertEquals(0, unassignedPartitionsQueue.size());
}