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

Example 16 with MockConsumer

use of org.apache.kafka.clients.consumer.MockConsumer in project kafka by apache.

the class StreamThreadTest method shouldThrowTaskMigratedExceptionHandlingRevocation.

@Test
public void shouldThrowTaskMigratedExceptionHandlingRevocation() {
    final Set<TopicPartition> assignedPartitions = Collections.singleton(t1p1);
    final TaskManager taskManager = EasyMock.createNiceMock(TaskManager.class);
    expect(taskManager.producerClientIds()).andStubReturn(Collections.emptySet());
    final MockConsumer<byte[], byte[]> consumer = new MockConsumer<>(OffsetResetStrategy.LATEST);
    consumer.assign(assignedPartitions);
    consumer.updateBeginningOffsets(Collections.singletonMap(t1p1, 0L));
    consumer.updateEndOffsets(Collections.singletonMap(t1p1, 10L));
    taskManager.handleRevocation(assignedPartitions);
    EasyMock.expectLastCall().andThrow(new TaskMigratedException("Revocation non fatal exception", new RuntimeException()));
    EasyMock.replay(taskManager);
    final StreamsMetricsImpl streamsMetrics = new StreamsMetricsImpl(metrics, CLIENT_ID, StreamsConfig.METRICS_LATEST, mockTime);
    final TopologyMetadata topologyMetadata = new TopologyMetadata(internalTopologyBuilder, config);
    topologyMetadata.buildAndRewriteTopology();
    final StreamThread thread = buildStreamThread(consumer, taskManager, config, topologyMetadata).updateThreadMetadata(getSharedAdminClientId(CLIENT_ID));
    consumer.schedulePollTask(() -> {
        thread.setState(StreamThread.State.PARTITIONS_REVOKED);
        thread.rebalanceListener().onPartitionsRevoked(assignedPartitions);
    });
    thread.setState(StreamThread.State.STARTING);
    assertThrows(TaskMigratedException.class, thread::runOnce);
}

Example 17 with MockConsumer

use of org.apache.kafka.clients.consumer.MockConsumer in project kafka by apache.

the class StreamThreadTest method shouldRecoverFromInvalidOffsetExceptionOnRestoreAndFinishRestore.

@Test
public void shouldRecoverFromInvalidOffsetExceptionOnRestoreAndFinishRestore() throws Exception {
    internalStreamsBuilder.stream(Collections.singleton("topic"), consumed).groupByKey().count(Materialized.as("count"));
    internalStreamsBuilder.buildAndOptimizeTopology();
    final StreamThread thread = createStreamThread("clientId", config, false);
    final MockConsumer<byte[], byte[]> mockConsumer = (MockConsumer<byte[], byte[]>) thread.mainConsumer();
    final MockConsumer<byte[], byte[]> mockRestoreConsumer = (MockConsumer<byte[], byte[]>) thread.restoreConsumer();
    final MockAdminClient mockAdminClient = (MockAdminClient) thread.adminClient();
    final TopicPartition topicPartition = new TopicPartition("topic", 0);
    final Set<TopicPartition> topicPartitionSet = Collections.singleton(topicPartition);
    final Map<TaskId, Set<TopicPartition>> activeTasks = new HashMap<>();
    final TaskId task0 = new TaskId(0, 0);
    activeTasks.put(task0, topicPartitionSet);
    thread.taskManager().handleAssignment(activeTasks, emptyMap());
    mockConsumer.updatePartitions("topic", Collections.singletonList(new PartitionInfo("topic", 0, null, new Node[0], new Node[0])));
    mockConsumer.updateBeginningOffsets(Collections.singletonMap(topicPartition, 0L));
    mockRestoreConsumer.updatePartitions("stream-thread-test-count-changelog", Collections.singletonList(new PartitionInfo("stream-thread-test-count-changelog", 0, null, new Node[0], new Node[0])));
    final TopicPartition changelogPartition = new TopicPartition("stream-thread-test-count-changelog", 0);
    final Set<TopicPartition> changelogPartitionSet = Collections.singleton(changelogPartition);
    mockRestoreConsumer.updateBeginningOffsets(Collections.singletonMap(changelogPartition, 0L));
    mockAdminClient.updateEndOffsets(Collections.singletonMap(changelogPartition, 2L));
    mockConsumer.schedulePollTask(() -> {
        thread.setState(StreamThread.State.PARTITIONS_REVOKED);
        thread.rebalanceListener().onPartitionsAssigned(topicPartitionSet);
    });
    try {
        thread.start();
        TestUtils.waitForCondition(() -> mockRestoreConsumer.assignment().size() == 1, "Never get the assignment");
        mockRestoreConsumer.addRecord(new ConsumerRecord<>("stream-thread-test-count-changelog", 0, 0L, "K1".getBytes(), "V1".getBytes()));
        TestUtils.waitForCondition(() -> mockRestoreConsumer.position(changelogPartition) == 1L, "Never restore first record");
        mockRestoreConsumer.setPollException(new InvalidOffsetException("Try Again!") {

            @Override
            public Set<TopicPartition> partitions() {
                return changelogPartitionSet;
            }
        });
        // after handling the exception and reviving the task, the position
        // should be reset to the beginning.
        TestUtils.waitForCondition(() -> mockRestoreConsumer.position(changelogPartition) == 0L, "Never restore first record");
        mockRestoreConsumer.addRecord(new ConsumerRecord<>("stream-thread-test-count-changelog", 0, 0L, "K1".getBytes(), "V1".getBytes()));
        mockRestoreConsumer.addRecord(new ConsumerRecord<>("stream-thread-test-count-changelog", 0, 1L, "K2".getBytes(), "V2".getBytes()));
        TestUtils.waitForCondition(() -> {
            mockRestoreConsumer.assign(changelogPartitionSet);
            return mockRestoreConsumer.position(changelogPartition) == 2L;
        }, "Never finished restore");
    } finally {
        thread.shutdown();
        thread.join(10000);
    }
}

Example 18 with MockConsumer

use of org.apache.kafka.clients.consumer.MockConsumer in project kafka by apache.

the class StreamThreadTest method shouldNotCloseTaskAndRemoveFromTaskManagerIfProducerGotFencedInCommitTransactionWhenSuspendingTasks.

@Test
public void shouldNotCloseTaskAndRemoveFromTaskManagerIfProducerGotFencedInCommitTransactionWhenSuspendingTasks() {
    final StreamThread thread = createStreamThread(CLIENT_ID, new StreamsConfig(configProps(true)), true);
    internalTopologyBuilder.addSource(null, "name", null, null, null, topic1);
    internalTopologyBuilder.addSink("out", "output", null, null, null, "name");
    thread.setState(StreamThread.State.STARTING);
    thread.rebalanceListener().onPartitionsRevoked(Collections.emptySet());
    final Map<TaskId, Set<TopicPartition>> activeTasks = new HashMap<>();
    final List<TopicPartition> assignedPartitions = new ArrayList<>();
    // assign single partition
    assignedPartitions.add(t1p1);
    activeTasks.put(task1, Collections.singleton(t1p1));
    thread.taskManager().handleAssignment(activeTasks, emptyMap());
    final MockConsumer<byte[], byte[]> mockConsumer = (MockConsumer<byte[], byte[]>) thread.mainConsumer();
    mockConsumer.assign(assignedPartitions);
    mockConsumer.updateBeginningOffsets(Collections.singletonMap(t1p1, 0L));
    thread.rebalanceListener().onPartitionsAssigned(assignedPartitions);
    thread.runOnce();
    assertThat(thread.activeTasks().size(), equalTo(1));
    // need to process a record to enable committing
    addRecord(mockConsumer, 0L);
    thread.runOnce();
    clientSupplier.producers.get(0).commitTransactionException = new ProducerFencedException("Producer is fenced");
    assertThrows(TaskMigratedException.class, () -> thread.rebalanceListener().onPartitionsRevoked(assignedPartitions));
    assertFalse(clientSupplier.producers.get(0).transactionCommitted());
    assertFalse(clientSupplier.producers.get(0).closed());
    assertEquals(1, thread.activeTasks().size());
}

Example 19 with MockConsumer

use of org.apache.kafka.clients.consumer.MockConsumer in project kafka by apache.

the class StreamThreadTest method shouldInjectProducerPerThreadUsingClientSupplierOnCreateIfEosV2Enabled.

@Test
public void shouldInjectProducerPerThreadUsingClientSupplierOnCreateIfEosV2Enabled() {
    internalTopologyBuilder.addSource(null, "source1", null, null, null, topic1);
    final Properties props = configProps(true);
    final StreamThread thread = createStreamThread(CLIENT_ID, new StreamsConfig(props), true);
    thread.setState(StreamThread.State.STARTING);
    thread.rebalanceListener().onPartitionsRevoked(Collections.emptyList());
    final Map<TaskId, Set<TopicPartition>> activeTasks = new HashMap<>();
    final List<TopicPartition> assignedPartitions = new ArrayList<>();
    // assign single partition
    assignedPartitions.add(t1p1);
    assignedPartitions.add(t1p2);
    activeTasks.put(task1, Collections.singleton(t1p1));
    activeTasks.put(task2, Collections.singleton(t1p2));
    thread.taskManager().handleAssignment(activeTasks, emptyMap());
    final MockConsumer<byte[], byte[]> mockConsumer = (MockConsumer<byte[], byte[]>) thread.mainConsumer();
    mockConsumer.assign(assignedPartitions);
    final Map<TopicPartition, Long> beginOffsets = new HashMap<>();
    beginOffsets.put(t1p1, 0L);
    beginOffsets.put(t1p2, 0L);
    mockConsumer.updateBeginningOffsets(beginOffsets);
    thread.rebalanceListener().onPartitionsAssigned(new HashSet<>(assignedPartitions));
    thread.runOnce();
    assertThat(clientSupplier.producers.size(), is(1));
    assertSame(clientSupplier.consumer, thread.mainConsumer());
    assertSame(clientSupplier.restoreConsumer, thread.restoreConsumer());
}

Example 20 with MockConsumer

use of org.apache.kafka.clients.consumer.MockConsumer in project kafka by apache.

the class StoreChangelogReaderTest method shouldThrowIfPositionFail.

@Test
public void shouldThrowIfPositionFail() {
    final TaskId taskId = new TaskId(0, 0);
    EasyMock.expect(activeStateManager.taskId()).andReturn(taskId);
    EasyMock.expect(storeMetadata.offset()).andReturn(10L).anyTimes();
    EasyMock.replay(activeStateManager, storeMetadata, store);
    final MockConsumer<byte[], byte[]> consumer = new MockConsumer<byte[], byte[]>(OffsetResetStrategy.EARLIEST) {

        @Override
        public long position(final TopicPartition partition) {
            throw kaboom;
        }
    };
    adminClient.updateEndOffsets(Collections.singletonMap(tp, 10L));
    final StoreChangelogReader changelogReader = new StoreChangelogReader(time, config, logContext, adminClient, consumer, callback);
    changelogReader.register(tp, activeStateManager);
    final StreamsException thrown = assertThrows(StreamsException.class, () -> changelogReader.restore(Collections.singletonMap(taskId, mock(Task.class))));
    assertEquals(kaboom, thrown.getCause());
}