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

public TestKafkaConsumerThread(
		Consumer<byte[], byte[]> mockConsumer,
		ClosableBlockingQueue<KafkaTopicPartitionState<Object, TopicPartition>> unassignedPartitionsQueue,
		Handover handover) {

	super(
			mock(Logger.class),
			handover,
			new Properties(),
			unassignedPartitionsQueue,
			"test-kafka-consumer-thread",
			0,
			false,
			new UnregisteredMetricsGroup(),
			new UnregisteredMetricsGroup(),
			null);

	this.mockConsumer = mockConsumer;
}
 

public TestKafkaConsumerThread(
		KafkaConsumer<byte[], byte[]> mockConsumer,
		ClosableBlockingQueue<KafkaTopicPartitionState<TopicPartition>> unassignedPartitionsQueue,
		Handover handover) {

	super(
			mock(Logger.class),
			handover,
			new Properties(),
			unassignedPartitionsQueue,
			new KafkaConsumerCallBridge09(),
			"test-kafka-consumer-thread",
			0,
			false,
			new UnregisteredMetricsGroup(),
			new UnregisteredMetricsGroup(),
			null);

	this.mockConsumer = mockConsumer;
}
 

public TestKafkaConsumerThread(
		KafkaConsumer<byte[], byte[]> mockConsumer,
		ClosableBlockingQueue<KafkaTopicPartitionState<TopicPartition>> unassignedPartitionsQueue,
		Handover handover) {

	super(
			mock(Logger.class),
			handover,
			new Properties(),
			unassignedPartitionsQueue,
			new KafkaConsumerCallBridge09(),
			"test-kafka-consumer-thread",
			0,
			false,
			new UnregisteredMetricsGroup(),
			new UnregisteredMetricsGroup(),
			null);

	this.mockConsumer = mockConsumer;
}
 

@Test(timeout = 10000)
public void testCloseWithoutAssignedPartitions() throws Exception {
	// no initial assignment
	final KafkaConsumer<byte[], byte[]> mockConsumer = createMockConsumer(
		new LinkedHashMap<TopicPartition, Long>(),
		Collections.<TopicPartition, Long>emptyMap(),
		false,
		null,
		null);

	// setup latch so the test waits until testThread is blocked on getBatchBlocking method
	final MultiShotLatch getBatchBlockingInvoked = new MultiShotLatch();
	final ClosableBlockingQueue<KafkaTopicPartitionState<TopicPartition>> unassignedPartitionsQueue =
		new ClosableBlockingQueue<KafkaTopicPartitionState<TopicPartition>>() {
			@Override
			public List<KafkaTopicPartitionState<TopicPartition>> getBatchBlocking() throws InterruptedException {
				getBatchBlockingInvoked.trigger();
				return super.getBatchBlocking();
			}
		};

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

	testThread.start();
	getBatchBlockingInvoked.await();
	testThread.shutdown();
	testThread.join();
}
 

public TestKafkaConsumerThreadRateLimit(Logger log,
		Handover handover, Properties kafkaProperties,
		ClosableBlockingQueue<KafkaTopicPartitionState<Object, TopicPartition>> unassignedPartitionsQueue,
		String threadName, long pollTimeout,
		boolean useMetrics, MetricGroup consumerMetricGroup,
		MetricGroup subtaskMetricGroup,
		Consumer<byte[], byte[]> mockConsumer,
		FlinkConnectorRateLimiter rateLimiter) {
	super(log, handover, kafkaProperties, unassignedPartitionsQueue,
			threadName,
			pollTimeout, useMetrics, consumerMetricGroup, subtaskMetricGroup,
		rateLimiter);
	this.mockConsumer = mockConsumer;
}
 

@Test(timeout = 10000)
public void testCloseWithoutAssignedPartitions() throws Exception {
	// no initial assignment
	final Consumer<byte[], byte[]> mockConsumer = createMockConsumer(
		new LinkedHashMap<TopicPartition, Long>(),
		Collections.<TopicPartition, Long>emptyMap(),
		false,
		null,
		null);

	// setup latch so the test waits until testThread is blocked on getBatchBlocking method
	final MultiShotLatch getBatchBlockingInvoked = new MultiShotLatch();
	final ClosableBlockingQueue<KafkaTopicPartitionState<Object, TopicPartition>> unassignedPartitionsQueue =
		new ClosableBlockingQueue<KafkaTopicPartitionState<Object, TopicPartition>>() {
			@Override
			public List<KafkaTopicPartitionState<Object, TopicPartition>> getBatchBlocking() throws InterruptedException {
				getBatchBlockingInvoked.trigger();
				return super.getBatchBlocking();
			}
		};

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

	testThread.start();
	getBatchBlockingInvoked.await();
	testThread.shutdown();
	testThread.join();
}
 

public KafkaConsumerThread(
		Logger log,
		Handover handover,
		Properties kafkaProperties,
		ClosableBlockingQueue<KafkaTopicPartitionState<T, TopicPartition>> unassignedPartitionsQueue,
		String threadName,
		long pollTimeout,
		boolean useMetrics,
		MetricGroup consumerMetricGroup,
		MetricGroup subtaskMetricGroup,
		FlinkConnectorRateLimiter rateLimiter) {

	super(threadName);
	setDaemon(true);

	this.log = checkNotNull(log);
	this.handover = checkNotNull(handover);
	this.kafkaProperties = checkNotNull(kafkaProperties);
	this.consumerMetricGroup = checkNotNull(consumerMetricGroup);
	this.subtaskMetricGroup = checkNotNull(subtaskMetricGroup);

	this.unassignedPartitionsQueue = checkNotNull(unassignedPartitionsQueue);

	this.pollTimeout = pollTimeout;
	this.useMetrics = useMetrics;

	this.consumerReassignmentLock = new Object();
	this.nextOffsetsToCommit = new AtomicReference<>();
	this.running = true;

	if (rateLimiter != null) {
		this.rateLimiter = rateLimiter;
	}
}
 

public KafkaConsumerThread(
		Logger log,
		Handover handover,
		Properties kafkaProperties,
		ClosableBlockingQueue<KafkaTopicPartitionState<T, TopicPartition>> unassignedPartitionsQueue,
		String threadName,
		long pollTimeout,
		boolean useMetrics,
		MetricGroup consumerMetricGroup,
		MetricGroup subtaskMetricGroup) {

	super(threadName);
	setDaemon(true);

	this.log = checkNotNull(log);
	this.handover = checkNotNull(handover);
	this.kafkaProperties = checkNotNull(kafkaProperties);
	this.consumerMetricGroup = checkNotNull(consumerMetricGroup);
	this.subtaskMetricGroup = checkNotNull(subtaskMetricGroup);

	this.unassignedPartitionsQueue = checkNotNull(unassignedPartitionsQueue);

	this.pollTimeout = pollTimeout;
	this.useMetrics = useMetrics;

	this.consumerReassignmentLock = new Object();
	this.nextOffsetsToCommit = new AtomicReference<>();
	this.running = true;
}
 

public KafkaConsumerThread(
		Logger log,
		Handover handover,
		Properties kafkaProperties,
		ClosableBlockingQueue<KafkaTopicPartitionState<TopicPartition>> unassignedPartitionsQueue,
		String threadName,
		long pollTimeout,
		boolean useMetrics,
		MetricGroup consumerMetricGroup,
		MetricGroup subtaskMetricGroup) {

	super(threadName);
	setDaemon(true);

	this.log = checkNotNull(log);
	this.handover = checkNotNull(handover);
	this.kafkaProperties = checkNotNull(kafkaProperties);
	this.consumerMetricGroup = checkNotNull(consumerMetricGroup);
	this.subtaskMetricGroup = checkNotNull(subtaskMetricGroup);

	this.unassignedPartitionsQueue = checkNotNull(unassignedPartitionsQueue);

	this.pollTimeout = pollTimeout;
	this.useMetrics = useMetrics;

	this.consumerReassignmentLock = new Object();
	this.nextOffsetsToCommit = new AtomicReference<>();
	this.running = true;
}
 

public TestKafkaConsumerThreadRateLimit(Logger log,
		Handover handover, Properties kafkaProperties,
		ClosableBlockingQueue<KafkaTopicPartitionState<TopicPartition>> unassignedPartitionsQueue,
		KafkaConsumerCallBridge09 consumerCallBridge, String threadName, long pollTimeout,
		boolean useMetrics, MetricGroup consumerMetricGroup,
		MetricGroup subtaskMetricGroup,
		KafkaConsumer mockConsumer,
		FlinkConnectorRateLimiter rateLimiter) {
	super(log, handover, kafkaProperties, unassignedPartitionsQueue, consumerCallBridge,
			threadName,
			pollTimeout, useMetrics, consumerMetricGroup, subtaskMetricGroup,
		rateLimiter);
	this.mockConsumer = mockConsumer;
}
 

public KafkaConsumerThread(
		Logger log,
		Handover handover,
		Properties kafkaProperties,
		ClosableBlockingQueue<KafkaTopicPartitionState<TopicPartition>> unassignedPartitionsQueue,
		KafkaConsumerCallBridge09 consumerCallBridge,
		String threadName,
		long pollTimeout,
		boolean useMetrics,
		MetricGroup consumerMetricGroup,
		MetricGroup subtaskMetricGroup,
		FlinkConnectorRateLimiter rateLimiter) {

	super(threadName);
	setDaemon(true);

	this.log = checkNotNull(log);
	this.handover = checkNotNull(handover);
	this.kafkaProperties = checkNotNull(kafkaProperties);
	this.consumerMetricGroup = checkNotNull(consumerMetricGroup);
	this.subtaskMetricGroup = checkNotNull(subtaskMetricGroup);
	this.consumerCallBridge = checkNotNull(consumerCallBridge);

	this.unassignedPartitionsQueue = checkNotNull(unassignedPartitionsQueue);

	this.pollTimeout = pollTimeout;
	this.useMetrics = useMetrics;

	this.consumerReassignmentLock = new Object();
	this.nextOffsetsToCommit = new AtomicReference<>();
	this.running = true;

	if (rateLimiter != null) {
		this.rateLimiter = rateLimiter;
	}
}
 

public KafkaConsumerThread(
		Logger log,
		Handover handover,
		Properties kafkaProperties,
		ClosableBlockingQueue<KafkaTopicPartitionState<TopicPartition>> unassignedPartitionsQueue,
		KafkaConsumerCallBridge09 consumerCallBridge,
		String threadName,
		long pollTimeout,
		boolean useMetrics,
		MetricGroup consumerMetricGroup,
		MetricGroup subtaskMetricGroup,
		FlinkConnectorRateLimiter rateLimiter) {

	super(threadName);
	setDaemon(true);

	this.log = checkNotNull(log);
	this.handover = checkNotNull(handover);
	this.kafkaProperties = checkNotNull(kafkaProperties);
	this.consumerMetricGroup = checkNotNull(consumerMetricGroup);
	this.subtaskMetricGroup = checkNotNull(subtaskMetricGroup);
	this.consumerCallBridge = checkNotNull(consumerCallBridge);

	this.unassignedPartitionsQueue = checkNotNull(unassignedPartitionsQueue);

	this.pollTimeout = pollTimeout;
	this.useMetrics = useMetrics;

	this.consumerReassignmentLock = new Object();
	this.nextOffsetsToCommit = new AtomicReference<>();
	this.running = true;

	if (rateLimiter != null) {
		this.rateLimiter = rateLimiter;
	}
}
 

public KafkaConsumerThread(
		Logger log,
		Handover handover,
		Properties kafkaProperties,
		ClosableBlockingQueue<KafkaTopicPartitionState<TopicPartition>> unassignedPartitionsQueue,
		String threadName,
		long pollTimeout,
		boolean useMetrics,
		MetricGroup consumerMetricGroup,
		MetricGroup subtaskMetricGroup) {

	super(threadName);
	setDaemon(true);

	this.log = checkNotNull(log);
	this.handover = checkNotNull(handover);
	this.kafkaProperties = checkNotNull(kafkaProperties);
	this.consumerMetricGroup = checkNotNull(consumerMetricGroup);
	this.subtaskMetricGroup = checkNotNull(subtaskMetricGroup);

	this.unassignedPartitionsQueue = checkNotNull(unassignedPartitionsQueue);

	this.pollTimeout = pollTimeout;
	this.useMetrics = useMetrics;

	this.consumerReassignmentLock = new Object();
	this.nextOffsetsToCommit = new AtomicReference<>();
	this.running = true;
}
 

public TestKafkaConsumerThreadRateLimit(Logger log,
		Handover handover, Properties kafkaProperties,
		ClosableBlockingQueue<KafkaTopicPartitionState<TopicPartition>> unassignedPartitionsQueue,
		KafkaConsumerCallBridge09 consumerCallBridge, String threadName, long pollTimeout,
		boolean useMetrics, MetricGroup consumerMetricGroup,
		MetricGroup subtaskMetricGroup,
		KafkaConsumer mockConsumer,
		FlinkConnectorRateLimiter rateLimiter) {
	super(log, handover, kafkaProperties, unassignedPartitionsQueue, consumerCallBridge,
			threadName,
			pollTimeout, useMetrics, consumerMetricGroup, subtaskMetricGroup,
		rateLimiter);
	this.mockConsumer = mockConsumer;
}
 

@Test(timeout = 10000)
public void testCloseWithoutAssignedPartitions() throws Exception {
	// no initial assignment
	final KafkaConsumer<byte[], byte[]> mockConsumer = createMockConsumer(
		new LinkedHashMap<TopicPartition, Long>(),
		Collections.<TopicPartition, Long>emptyMap(),
		false,
		null,
		null);

	// setup latch so the test waits until testThread is blocked on getBatchBlocking method
	final MultiShotLatch getBatchBlockingInvoked = new MultiShotLatch();
	final ClosableBlockingQueue<KafkaTopicPartitionState<TopicPartition>> unassignedPartitionsQueue =
		new ClosableBlockingQueue<KafkaTopicPartitionState<TopicPartition>>() {
			@Override
			public List<KafkaTopicPartitionState<TopicPartition>> getBatchBlocking() throws InterruptedException {
				getBatchBlockingInvoked.trigger();
				return super.getBatchBlocking();
			}
		};

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

	testThread.start();
	getBatchBlockingInvoked.await();
	testThread.shutdown();
	testThread.join();
}
 

@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 already have some assignments
 *  - new unassigned partitions already have defined offsets
 *
 * <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 testReassigningPartitionsWithDefinedOffsets() 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 defined offsets --------

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

	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);
	}

	final KafkaConsumer<byte[], byte[]> mockConsumer = createMockConsumer(
			mockConsumerAssignmentsAndPositions,
			Collections.<TopicPartition, Long>emptyMap(),
			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();

	// 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(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());
}
 

/**
 * Tests reassignment works correctly in the case when:
 *  - the consumer already have some assignments
 *  - new unassigned partitions already have defined offsets
 *
 * <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 testReassigningPartitionsWithDefinedOffsets() 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 defined offsets --------

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

	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);
	}

	final Consumer<byte[], byte[]> mockConsumer = createMockConsumer(
			mockConsumerAssignmentsAndPositions,
			Collections.<TopicPartition, Long>emptyMap(),
			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();

	// 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());
}
 

/**
 * 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());
}
 

/**
 * 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(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());
}
 

/**
 * Tests reassignment works correctly in the case when:
 *  - the consumer already have some assignments
 *  - new unassigned partitions already have defined offsets
 *
 * <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 testReassigningPartitionsWithDefinedOffsets() 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 defined offsets --------

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

	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);
	}

	final KafkaConsumer<byte[], byte[]> mockConsumer = createMockConsumer(
			mockConsumerAssignmentsAndPositions,
			Collections.<TopicPartition, Long>emptyMap(),
			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();

	// 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());
}