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

Example 36 with MockConsumer

use of org.apache.kafka.clients.consumer.MockConsumer in project apache-kafka-on-k8s by banzaicloud.

the class StreamThreadTest method shouldCloseAllTaskProducersOnCloseIfEosEnabled.

@Test
public void shouldCloseAllTaskProducersOnCloseIfEosEnabled() {
    internalTopologyBuilder.addSource(null, "source1", null, null, null, topic1);
    final StreamThread thread = createStreamThread(clientId, new StreamsConfig(configProps(true)), true);
    thread.setState(StreamThread.State.RUNNING);
    thread.rebalanceListener.onPartitionsRevoked(Collections.<TopicPartition>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().setAssignmentMetadata(activeTasks, Collections.<TaskId, Set<TopicPartition>>emptyMap());
    final MockConsumer<byte[], byte[]> mockConsumer = (MockConsumer<byte[], byte[]>) thread.consumer;
    mockConsumer.assign(assignedPartitions);
    Map<TopicPartition, Long> beginOffsets = new HashMap<>();
    beginOffsets.put(t1p1, 0L);
    beginOffsets.put(t1p2, 0L);
    mockConsumer.updateBeginningOffsets(beginOffsets);
    thread.rebalanceListener.onPartitionsAssigned(assignedPartitions);
    thread.shutdown();
    thread.run();
    for (final Task task : thread.tasks().values()) {
        assertTrue(((MockProducer) ((RecordCollectorImpl) ((StreamTask) task).recordCollector()).producer()).closed());
    }
}

Example 37 with MockConsumer

use of org.apache.kafka.clients.consumer.MockConsumer in project apache-kafka-on-k8s by banzaicloud.

the class StreamThreadTest method shouldRecoverFromInvalidOffsetExceptionOnRestoreAndFinishRestore.

@Test
public void shouldRecoverFromInvalidOffsetExceptionOnRestoreAndFinishRestore() throws Exception {
    internalStreamsBuilder.stream(Collections.singleton("topic"), consumed).groupByKey().count(Materialized.<Object, Long, KeyValueStore<Bytes, byte[]>>as("count"));
    final StreamThread thread = createStreamThread("cliendId", config, false);
    final MockConsumer<byte[], byte[]> mockConsumer = (MockConsumer<byte[], byte[]>) thread.consumer;
    final MockConsumer<byte[], byte[]> mockRestoreConsumer = (MockConsumer<byte[], byte[]>) thread.restoreConsumer;
    final TopicPartition topicPartition = new TopicPartition("topic", 0);
    final Set<TopicPartition> topicPartitionSet = Collections.singleton(topicPartition);
    final Map<TaskId, Set<TopicPartition>> activeTasks = new HashMap<>();
    activeTasks.put(new TaskId(0, 0), topicPartitionSet);
    thread.taskManager().setAssignmentMetadata(activeTasks, Collections.<TaskId, Set<TopicPartition>>emptyMap());
    mockConsumer.updatePartitions("topic", new ArrayList<PartitionInfo>() {

        {
            add(new PartitionInfo("topic", 0, null, new Node[0], new Node[0]));
        }
    });
    mockConsumer.updateBeginningOffsets(Collections.singletonMap(topicPartition, 0L));
    mockRestoreConsumer.updatePartitions("stream-thread-test-count-changelog", new ArrayList<PartitionInfo>() {

        {
            add(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));
    mockRestoreConsumer.updateEndOffsets(Collections.singletonMap(changelogPartition, 2L));
    mockConsumer.schedulePollTask(new Runnable() {

        @Override
        public void run() {
            thread.setState(StreamThread.State.PARTITIONS_REVOKED);
            thread.rebalanceListener.onPartitionsAssigned(topicPartitionSet);
        }
    });
    try {
        thread.start();
        TestUtils.waitForCondition(new TestCondition() {

            @Override
            public boolean conditionMet() {
                return mockRestoreConsumer.assignment().size() == 1;
            }
        }, "Never restore first record");
        mockRestoreConsumer.addRecord(new ConsumerRecord<>("stream-thread-test-count-changelog", 0, 0L, "K1".getBytes(), "V1".getBytes()));
        TestUtils.waitForCondition(new TestCondition() {

            @Override
            public boolean conditionMet() {
                return mockRestoreConsumer.position(changelogPartition) == 1L;
            }
        }, "Never restore first record");
        mockRestoreConsumer.setException(new InvalidOffsetException("Try Again!") {

            @Override
            public Set<TopicPartition> partitions() {
                return changelogPartitionSet;
            }
        });
        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(new TestCondition() {

            @Override
            public boolean conditionMet() {
                mockRestoreConsumer.assign(changelogPartitionSet);
                return mockRestoreConsumer.position(changelogPartition) == 2L;
            }
        }, "Never finished restore");
    } finally {
        thread.shutdown();
        thread.join(10000);
    }
}

Example 38 with MockConsumer

use of org.apache.kafka.clients.consumer.MockConsumer in project apache-kafka-on-k8s by banzaicloud.

the class StreamThreadTest method shouldReturnActiveTaskMetadataWhileRunningState.

@Test
public void shouldReturnActiveTaskMetadataWhileRunningState() {
    internalTopologyBuilder.addSource(null, "source", null, null, null, topic1);
    final StreamThread thread = createStreamThread(clientId, config, false);
    thread.setState(StreamThread.State.RUNNING);
    thread.rebalanceListener.onPartitionsRevoked(null);
    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().setAssignmentMetadata(activeTasks, Collections.<TaskId, Set<TopicPartition>>emptyMap());
    final MockConsumer<byte[], byte[]> mockConsumer = (MockConsumer<byte[], byte[]>) thread.consumer;
    mockConsumer.assign(assignedPartitions);
    mockConsumer.updateBeginningOffsets(Collections.singletonMap(t1p1, 0L));
    thread.rebalanceListener.onPartitionsAssigned(assignedPartitions);
    thread.runOnce(-1);
    ThreadMetadata threadMetadata = thread.threadMetadata();
    assertEquals(StreamThread.State.RUNNING.name(), threadMetadata.threadState());
    assertTrue(threadMetadata.activeTasks().contains(new TaskMetadata(task1.toString(), Utils.mkSet(t1p1))));
    assertTrue(threadMetadata.standbyTasks().isEmpty());
}

Example 39 with MockConsumer

use of org.apache.kafka.clients.consumer.MockConsumer in project apache-kafka-on-k8s by banzaicloud.

the class StreamThreadTest method shouldCloseTaskAsZombieAndRemoveFromActiveTasksIfProducerWasFencedWhileProcessing.

@Test
public void shouldCloseTaskAsZombieAndRemoveFromActiveTasksIfProducerWasFencedWhileProcessing() throws Exception {
    internalTopologyBuilder.addSource(null, "source", null, null, null, topic1);
    internalTopologyBuilder.addSink("sink", "dummyTopic", null, null, null, "source");
    final StreamThread thread = createStreamThread(clientId, new StreamsConfig(configProps(true)), true);
    final MockConsumer<byte[], byte[]> consumer = clientSupplier.consumer;
    consumer.updatePartitions(topic1, Collections.singletonList(new PartitionInfo(topic1, 1, null, null, null)));
    thread.setState(StreamThread.State.RUNNING);
    thread.rebalanceListener.onPartitionsRevoked(null);
    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().setAssignmentMetadata(activeTasks, Collections.<TaskId, Set<TopicPartition>>emptyMap());
    final MockConsumer<byte[], byte[]> mockConsumer = (MockConsumer<byte[], byte[]>) thread.consumer;
    mockConsumer.assign(assignedPartitions);
    mockConsumer.updateBeginningOffsets(Collections.singletonMap(t1p1, 0L));
    thread.rebalanceListener.onPartitionsAssigned(assignedPartitions);
    thread.runOnce(-1);
    assertThat(thread.tasks().size(), equalTo(1));
    final MockProducer producer = clientSupplier.producers.get(0);
    // change consumer subscription from "pattern" to "manual" to be able to call .addRecords()
    consumer.updateBeginningOffsets(Collections.singletonMap(assignedPartitions.iterator().next(), 0L));
    consumer.unsubscribe();
    consumer.assign(new HashSet<>(assignedPartitions));
    consumer.addRecord(new ConsumerRecord<>(topic1, 1, 0, new byte[0], new byte[0]));
    mockTime.sleep(config.getLong(StreamsConfig.COMMIT_INTERVAL_MS_CONFIG) + 1);
    thread.runOnce(-1);
    assertThat(producer.history().size(), equalTo(1));
    assertFalse(producer.transactionCommitted());
    mockTime.sleep(config.getLong(StreamsConfig.COMMIT_INTERVAL_MS_CONFIG) + 1L);
    TestUtils.waitForCondition(new TestCondition() {

        @Override
        public boolean conditionMet() {
            return producer.commitCount() == 1;
        }
    }, "StreamsThread did not commit transaction.");
    producer.fenceProducer();
    mockTime.sleep(config.getLong(StreamsConfig.COMMIT_INTERVAL_MS_CONFIG) + 1L);
    consumer.addRecord(new ConsumerRecord<>(topic1, 1, 0, new byte[0], new byte[0]));
    try {
        thread.runOnce(-1);
        fail("Should have thrown TaskMigratedException");
    } catch (final TaskMigratedException expected) {
    /* ignore */
    }
    TestUtils.waitForCondition(new TestCondition() {

        @Override
        public boolean conditionMet() {
            return thread.tasks().isEmpty();
        }
    }, "StreamsThread did not remove fenced zombie task.");
    assertThat(producer.commitCount(), equalTo(1L));
}

Example 40 with MockConsumer

use of org.apache.kafka.clients.consumer.MockConsumer in project apache-kafka-on-k8s by banzaicloud.

the class StreamThreadTest method testPartitionAssignmentChangeForSingleGroup.

@Test
public void testPartitionAssignmentChangeForSingleGroup() {
    internalTopologyBuilder.addSource(null, "source1", null, null, null, topic1);
    final StreamThread thread = createStreamThread(clientId, config, false);
    final StateListenerStub stateListener = new StateListenerStub();
    thread.setStateListener(stateListener);
    assertEquals(thread.state(), StreamThread.State.CREATED);
    final ConsumerRebalanceListener rebalanceListener = thread.rebalanceListener;
    thread.setState(StreamThread.State.RUNNING);
    List<TopicPartition> revokedPartitions;
    List<TopicPartition> assignedPartitions;
    // revoke nothing
    revokedPartitions = Collections.emptyList();
    rebalanceListener.onPartitionsRevoked(revokedPartitions);
    assertEquals(thread.state(), StreamThread.State.PARTITIONS_REVOKED);
    // assign single partition
    assignedPartitions = Collections.singletonList(t1p1);
    thread.taskManager().setAssignmentMetadata(Collections.<TaskId, Set<TopicPartition>>emptyMap(), Collections.<TaskId, Set<TopicPartition>>emptyMap());
    final MockConsumer<byte[], byte[]> mockConsumer = (MockConsumer<byte[], byte[]>) thread.consumer;
    mockConsumer.assign(assignedPartitions);
    mockConsumer.updateBeginningOffsets(Collections.singletonMap(t1p1, 0L));
    rebalanceListener.onPartitionsAssigned(assignedPartitions);
    thread.runOnce(-1);
    assertEquals(thread.state(), StreamThread.State.RUNNING);
    Assert.assertEquals(4, stateListener.numChanges);
    Assert.assertEquals(StreamThread.State.PARTITIONS_ASSIGNED, stateListener.oldState);
    thread.shutdown();
    assertTrue(thread.state() == StreamThread.State.PENDING_SHUTDOWN);
}