Examples with MockRebalanceListener org.apache.kafka.clients.consumer.internals.MockRebalanceListener used on opensource projects

Example 1 with MockRebalanceListener

use of org.apache.kafka.clients.consumer.internals.MockRebalanceListener in project kafka by apache.

the class StreamThreadTest method shouldNotReturnDataAfterTaskMigrated.

@Test
public void shouldNotReturnDataAfterTaskMigrated() {
    final TaskManager taskManager = EasyMock.createNiceMock(TaskManager.class);
    expect(taskManager.producerClientIds()).andStubReturn(Collections.emptySet());
    final InternalTopologyBuilder internalTopologyBuilder = EasyMock.createNiceMock(InternalTopologyBuilder.class);
    expect(internalTopologyBuilder.fullSourceTopicNames()).andReturn(Collections.singletonList(topic1)).times(2);
    final MockConsumer<byte[], byte[]> consumer = new MockConsumer<>(OffsetResetStrategy.LATEST);
    consumer.subscribe(Collections.singletonList(topic1), new MockRebalanceListener());
    consumer.rebalance(Collections.singletonList(t1p1));
    consumer.updateEndOffsets(Collections.singletonMap(t1p1, 10L));
    consumer.seekToEnd(Collections.singletonList(t1p1));
    final ChangelogReader changelogReader = new MockChangelogReader() {

        @Override
        public void restore(final Map<TaskId, Task> tasks) {
            consumer.addRecord(new ConsumerRecord<>(topic1, 1, 11, new byte[0], new byte[0]));
            consumer.addRecord(new ConsumerRecord<>(topic1, 1, 12, new byte[1], new byte[0]));
            throw new TaskMigratedException("Changelog restore found task migrated", new RuntimeException("restore task migrated"));
        }
    };
    taskManager.handleLostAll();
    EasyMock.replay(taskManager, internalTopologyBuilder);
    final StreamsMetricsImpl streamsMetrics = new StreamsMetricsImpl(metrics, CLIENT_ID, StreamsConfig.METRICS_LATEST, mockTime);
    final StreamThread thread = new StreamThread(mockTime, config, null, consumer, consumer, changelogReader, null, taskManager, streamsMetrics, new TopologyMetadata(internalTopologyBuilder, config), CLIENT_ID, new LogContext(""), new AtomicInteger(), new AtomicLong(Long.MAX_VALUE), new LinkedList<>(), null, HANDLER, null).updateThreadMetadata(getSharedAdminClientId(CLIENT_ID));
    final StreamsException thrown = assertThrows(StreamsException.class, thread::run);
    verify(taskManager);
    assertThat(thrown.getCause(), isA(IllegalStateException.class));
    // The Mock consumer shall throw as the assignment has been wiped out, but records are assigned.
    assertEquals("No current assignment for partition topic1-1", thrown.getCause().getMessage());
    assertFalse(consumer.shouldRebalance());
}

Example 2 with MockRebalanceListener

use of org.apache.kafka.clients.consumer.internals.MockRebalanceListener in project kafka by apache.

the class KafkaConsumerTest method testEnforceRebalanceTriggersRebalanceOnNextPoll.

@Test
public void testEnforceRebalanceTriggersRebalanceOnNextPoll() {
    Time time = new MockTime(1L);
    ConsumerMetadata metadata = createMetadata(subscription);
    MockClient client = new MockClient(time, metadata);
    KafkaConsumer<String, String> consumer = newConsumer(time, client, subscription, metadata, assignor, true, groupInstanceId);
    MockRebalanceListener countingRebalanceListener = new MockRebalanceListener();
    initMetadata(client, Utils.mkMap(Utils.mkEntry(topic, 1), Utils.mkEntry(topic2, 1), Utils.mkEntry(topic3, 1)));
    consumer.subscribe(Arrays.asList(topic, topic2), countingRebalanceListener);
    Node node = metadata.fetch().nodes().get(0);
    prepareRebalance(client, node, assignor, Arrays.asList(tp0, t2p0), null);
    // a first rebalance to get the assignment, we need two poll calls since we need two round trips to finish join / sync-group
    consumer.poll(Duration.ZERO);
    consumer.poll(Duration.ZERO);
    // onPartitionsRevoked is not invoked when first joining the group
    assertEquals(countingRebalanceListener.revokedCount, 0);
    assertEquals(countingRebalanceListener.assignedCount, 1);
    consumer.enforceRebalance();
    // the next poll should trigger a rebalance
    consumer.poll(Duration.ZERO);
    assertEquals(countingRebalanceListener.revokedCount, 1);
}