Examples with MockClientSupplier org.apache.kafka.test.MockClientSupplier used on opensource projects

Example 1 with MockClientSupplier

use of org.apache.kafka.test.MockClientSupplier in project kafka by apache.

the class StreamThreadTest method testMaybeCommit.

@Test
public void testMaybeCommit() throws Exception {
    File baseDir = Files.createTempDirectory("test").toFile();
    try {
        final long commitInterval = 1000L;
        Properties props = configProps();
        props.setProperty(StreamsConfig.STATE_DIR_CONFIG, baseDir.getCanonicalPath());
        props.setProperty(StreamsConfig.COMMIT_INTERVAL_MS_CONFIG, Long.toString(commitInterval));
        StreamsConfig config = new StreamsConfig(props);
        final MockTime mockTime = new MockTime();
        TopologyBuilder builder = new TopologyBuilder().setApplicationId("X");
        builder.addSource("source1", "topic1");
        MockClientSupplier mockClientSupplier = new MockClientSupplier();
        StreamThread thread = new StreamThread(builder, config, mockClientSupplier, applicationId, clientId, processId, new Metrics(), mockTime, new StreamsMetadataState(builder, StreamsMetadataState.UNKNOWN_HOST), 0) {

            @Override
            public void maybeCommit(long now) {
                super.maybeCommit(now);
            }

            @Override
            protected StreamTask createStreamTask(TaskId id, Collection<TopicPartition> partitionsForTask) {
                ProcessorTopology topology = builder.build(id.topicGroupId);
                return new TestStreamTask(id, applicationId, partitionsForTask, topology, consumer, producer, restoreConsumer, config, new MockStreamsMetrics(new Metrics()), stateDirectory);
            }
        };
        initPartitionGrouper(config, thread, mockClientSupplier);
        ConsumerRebalanceListener rebalanceListener = thread.rebalanceListener;
        List<TopicPartition> revokedPartitions;
        List<TopicPartition> assignedPartitions;
        //
        // Assign t1p1 and t1p2. This should create Task 1 & 2
        //
        revokedPartitions = Collections.emptyList();
        assignedPartitions = Arrays.asList(t1p1, t1p2);
        rebalanceListener.onPartitionsRevoked(revokedPartitions);
        rebalanceListener.onPartitionsAssigned(assignedPartitions);
        assertEquals(2, thread.tasks().size());
        // no task is committed before the commit interval
        mockTime.sleep(commitInterval - 10L);
        thread.maybeCommit(mockTime.milliseconds());
        for (StreamTask task : thread.tasks().values()) {
            assertFalse(((TestStreamTask) task).committed);
        }
        // all tasks are committed after the commit interval
        mockTime.sleep(11L);
        thread.maybeCommit(mockTime.milliseconds());
        for (StreamTask task : thread.tasks().values()) {
            assertTrue(((TestStreamTask) task).committed);
            ((TestStreamTask) task).committed = false;
        }
        // no task is committed before the commit interval, again
        mockTime.sleep(commitInterval - 10L);
        thread.maybeCommit(mockTime.milliseconds());
        for (StreamTask task : thread.tasks().values()) {
            assertFalse(((TestStreamTask) task).committed);
        }
        // all tasks are committed after the commit interval, again
        mockTime.sleep(11L);
        thread.maybeCommit(mockTime.milliseconds());
        for (StreamTask task : thread.tasks().values()) {
            assertTrue(((TestStreamTask) task).committed);
            ((TestStreamTask) task).committed = false;
        }
    } finally {
        Utils.delete(baseDir);
    }
}

Example 2 with MockClientSupplier

use of org.apache.kafka.test.MockClientSupplier in project kafka by apache.

the class StreamThreadTest method shouldInitializeRestoreConsumerWithOffsetsFromStandbyTasks.

@Test
public void shouldInitializeRestoreConsumerWithOffsetsFromStandbyTasks() throws Exception {
    final KStreamBuilder builder = new KStreamBuilder();
    builder.setApplicationId(applicationId);
    builder.stream("t1").groupByKey().count("count-one");
    builder.stream("t2").groupByKey().count("count-two");
    final StreamsConfig config = new StreamsConfig(configProps());
    final MockClientSupplier clientSupplier = new MockClientSupplier();
    final StreamThread thread = new StreamThread(builder, config, clientSupplier, applicationId, clientId, processId, new Metrics(), new MockTime(), new StreamsMetadataState(builder, StreamsMetadataState.UNKNOWN_HOST), 0);
    final MockConsumer<byte[], byte[]> restoreConsumer = clientSupplier.restoreConsumer;
    restoreConsumer.updatePartitions("stream-thread-test-count-one-changelog", Collections.singletonList(new PartitionInfo("stream-thread-test-count-one-changelog", 0, null, new Node[0], new Node[0])));
    restoreConsumer.updatePartitions("stream-thread-test-count-two-changelog", Collections.singletonList(new PartitionInfo("stream-thread-test-count-two-changelog", 0, null, new Node[0], new Node[0])));
    final Map<TaskId, Set<TopicPartition>> standbyTasks = new HashMap<>();
    final TopicPartition t1 = new TopicPartition("t1", 0);
    standbyTasks.put(new TaskId(0, 0), Utils.mkSet(t1));
    thread.partitionAssignor(new StreamPartitionAssignor() {

        @Override
        Map<TaskId, Set<TopicPartition>> standbyTasks() {
            return standbyTasks;
        }
    });
    thread.rebalanceListener.onPartitionsRevoked(Collections.<TopicPartition>emptyList());
    thread.rebalanceListener.onPartitionsAssigned(Collections.<TopicPartition>emptyList());
    assertThat(restoreConsumer.assignment(), equalTo(Utils.mkSet(new TopicPartition("stream-thread-test-count-one-changelog", 0))));
    // assign an existing standby plus a new one
    standbyTasks.put(new TaskId(1, 0), Utils.mkSet(new TopicPartition("t2", 0)));
    thread.rebalanceListener.onPartitionsRevoked(Collections.<TopicPartition>emptyList());
    thread.rebalanceListener.onPartitionsAssigned(Collections.<TopicPartition>emptyList());
    assertThat(restoreConsumer.assignment(), equalTo(Utils.mkSet(new TopicPartition("stream-thread-test-count-one-changelog", 0), new TopicPartition("stream-thread-test-count-two-changelog", 0))));
}

Example 3 with MockClientSupplier

use of org.apache.kafka.test.MockClientSupplier in project kafka by apache.

the class StreamThreadTest method shouldNotViolateAtLeastOnceWhenExceptionOccursDuringCloseTopologyWhenSuspendingState.

@Test
public void shouldNotViolateAtLeastOnceWhenExceptionOccursDuringCloseTopologyWhenSuspendingState() throws Exception {
    final KStreamBuilder builder = new KStreamBuilder();
    builder.setApplicationId(applicationId);
    builder.stream("t1").groupByKey();
    final StreamsConfig config = new StreamsConfig(configProps());
    final MockClientSupplier clientSupplier = new MockClientSupplier();
    final TestStreamTask testStreamTask = new TestStreamTask(new TaskId(0, 0), applicationId, Utils.mkSet(new TopicPartition("t1", 0)), builder.build(0), clientSupplier.consumer, clientSupplier.producer, clientSupplier.restoreConsumer, config, new MockStreamsMetrics(new Metrics()), new StateDirectory(applicationId, config.getString(StreamsConfig.STATE_DIR_CONFIG), time)) {

        @Override
        public void closeTopology() {
            throw new RuntimeException("KABOOM!");
        }
    };
    final StreamsConfig config1 = new StreamsConfig(configProps());
    final StreamThread thread = new StreamThread(builder, config1, clientSupplier, applicationId, clientId, processId, new Metrics(), new MockTime(), new StreamsMetadataState(builder, StreamsMetadataState.UNKNOWN_HOST), 0) {

        @Override
        protected StreamTask createStreamTask(final TaskId id, final Collection<TopicPartition> partitions) {
            return testStreamTask;
        }
    };
    final Map<TaskId, Set<TopicPartition>> activeTasks = new HashMap<>();
    activeTasks.put(testStreamTask.id, testStreamTask.partitions);
    thread.partitionAssignor(new MockStreamsPartitionAssignor(activeTasks));
    thread.rebalanceListener.onPartitionsRevoked(Collections.<TopicPartition>emptyList());
    thread.rebalanceListener.onPartitionsAssigned(testStreamTask.partitions);
    try {
        thread.rebalanceListener.onPartitionsRevoked(Collections.<TopicPartition>emptyList());
        fail("should have thrown exception");
    } catch (Exception e) {
    // expected
    }
    assertFalse(testStreamTask.committed);
}

Example 4 with MockClientSupplier

use of org.apache.kafka.test.MockClientSupplier in project kafka by apache.

the class StreamThreadTest method testMaybeClean.

@Test
public void testMaybeClean() throws Exception {
    File baseDir = Files.createTempDirectory("test").toFile();
    try {
        final long cleanupDelay = 1000L;
        Properties props = configProps();
        props.setProperty(StreamsConfig.STATE_CLEANUP_DELAY_MS_CONFIG, Long.toString(cleanupDelay));
        props.setProperty(StreamsConfig.STATE_DIR_CONFIG, baseDir.getCanonicalPath());
        StreamsConfig config = new StreamsConfig(props);
        File applicationDir = new File(baseDir, applicationId);
        applicationDir.mkdir();
        File stateDir1 = new File(applicationDir, task1.toString());
        File stateDir2 = new File(applicationDir, task2.toString());
        File stateDir3 = new File(applicationDir, task3.toString());
        File extraDir = new File(applicationDir, "X");
        stateDir1.mkdir();
        stateDir2.mkdir();
        stateDir3.mkdir();
        extraDir.mkdir();
        final MockTime mockTime = new MockTime();
        TopologyBuilder builder = new TopologyBuilder().setApplicationId("X");
        builder.addSource("source1", "topic1");
        MockClientSupplier mockClientSupplier = new MockClientSupplier();
        StreamThread thread = new StreamThread(builder, config, mockClientSupplier, applicationId, clientId, processId, new Metrics(), mockTime, new StreamsMetadataState(builder, StreamsMetadataState.UNKNOWN_HOST), 0) {

            @Override
            public void maybeClean(long now) {
                super.maybeClean(now);
            }

            @Override
            protected StreamTask createStreamTask(TaskId id, Collection<TopicPartition> partitionsForTask) {
                ProcessorTopology topology = builder.build(id.topicGroupId);
                return new TestStreamTask(id, applicationId, partitionsForTask, topology, consumer, producer, restoreConsumer, config, new MockStreamsMetrics(new Metrics()), stateDirectory);
            }
        };
        initPartitionGrouper(config, thread, mockClientSupplier);
        ConsumerRebalanceListener rebalanceListener = thread.rebalanceListener;
        assertTrue(thread.tasks().isEmpty());
        mockTime.sleep(cleanupDelay);
        // all directories exist since an assignment didn't happen
        assertTrue(stateDir1.exists());
        assertTrue(stateDir2.exists());
        assertTrue(stateDir3.exists());
        assertTrue(extraDir.exists());
        List<TopicPartition> revokedPartitions;
        List<TopicPartition> assignedPartitions;
        Map<TaskId, StreamTask> prevTasks;
        //
        // Assign t1p1 and t1p2. This should create task1 & task2
        //
        final Map<TaskId, Set<TopicPartition>> activeTasks = new HashMap<>();
        activeTasks.put(task1, Collections.singleton(t1p1));
        activeTasks.put(task2, Collections.singleton(t1p2));
        thread.partitionAssignor(new MockStreamsPartitionAssignor(activeTasks));
        revokedPartitions = Collections.emptyList();
        assignedPartitions = Arrays.asList(t1p1, t1p2);
        prevTasks = new HashMap<>(thread.tasks());
        rebalanceListener.onPartitionsRevoked(revokedPartitions);
        rebalanceListener.onPartitionsAssigned(assignedPartitions);
        // there shouldn't be any previous task
        assertTrue(prevTasks.isEmpty());
        // task 1 & 2 are created
        assertEquals(2, thread.tasks().size());
        // all directories should still exit before the cleanup delay time
        mockTime.sleep(cleanupDelay - 10L);
        thread.maybeClean(mockTime.milliseconds());
        assertTrue(stateDir1.exists());
        assertTrue(stateDir2.exists());
        assertTrue(stateDir3.exists());
        assertTrue(extraDir.exists());
        // all state directories except for task task2 & task3 will be removed. the extra directory should still exists
        mockTime.sleep(11L);
        thread.maybeClean(mockTime.milliseconds());
        assertTrue(stateDir1.exists());
        assertTrue(stateDir2.exists());
        assertFalse(stateDir3.exists());
        assertTrue(extraDir.exists());
        //
        // Revoke t1p1 and t1p2. This should remove task1 & task2
        //
        activeTasks.clear();
        revokedPartitions = assignedPartitions;
        assignedPartitions = Collections.emptyList();
        prevTasks = new HashMap<>(thread.tasks());
        rebalanceListener.onPartitionsRevoked(revokedPartitions);
        rebalanceListener.onPartitionsAssigned(assignedPartitions);
        // previous tasks should be committed
        assertEquals(2, prevTasks.size());
        for (StreamTask task : prevTasks.values()) {
            assertTrue(((TestStreamTask) task).committed);
            ((TestStreamTask) task).committed = false;
        }
        // no task
        assertTrue(thread.tasks().isEmpty());
        // all state directories for task task1 & task2 still exist before the cleanup delay time
        mockTime.sleep(cleanupDelay - 10L);
        thread.maybeClean(mockTime.milliseconds());
        assertTrue(stateDir1.exists());
        assertTrue(stateDir2.exists());
        assertFalse(stateDir3.exists());
        assertTrue(extraDir.exists());
        // all state directories for task task1 & task2 are removed
        mockTime.sleep(11L);
        thread.maybeClean(mockTime.milliseconds());
        assertFalse(stateDir1.exists());
        assertFalse(stateDir2.exists());
        assertFalse(stateDir3.exists());
        assertTrue(extraDir.exists());
    } finally {
        Utils.delete(baseDir);
    }
}

Example 5 with MockClientSupplier

use of org.apache.kafka.test.MockClientSupplier in project kafka by apache.

the class StreamThreadTest method shouldCloseSuspendedTasksThatAreNoLongerAssignedToThisStreamThreadBeforeCreatingNewTasks.

@Test
public void shouldCloseSuspendedTasksThatAreNoLongerAssignedToThisStreamThreadBeforeCreatingNewTasks() throws Exception {
    final KStreamBuilder builder = new KStreamBuilder();
    builder.setApplicationId(applicationId);
    builder.stream("t1").groupByKey().count("count-one");
    builder.stream("t2").groupByKey().count("count-two");
    final StreamsConfig config = new StreamsConfig(configProps());
    final MockClientSupplier clientSupplier = new MockClientSupplier();
    final StreamThread thread = new StreamThread(builder, config, clientSupplier, applicationId, clientId, processId, new Metrics(), new MockTime(), new StreamsMetadataState(builder, StreamsMetadataState.UNKNOWN_HOST), 0);
    final MockConsumer<byte[], byte[]> restoreConsumer = clientSupplier.restoreConsumer;
    restoreConsumer.updatePartitions("stream-thread-test-count-one-changelog", Collections.singletonList(new PartitionInfo("stream-thread-test-count-one-changelog", 0, null, new Node[0], new Node[0])));
    restoreConsumer.updatePartitions("stream-thread-test-count-two-changelog", Collections.singletonList(new PartitionInfo("stream-thread-test-count-two-changelog", 0, null, new Node[0], new Node[0])));
    final HashMap<TopicPartition, Long> offsets = new HashMap<>();
    offsets.put(new TopicPartition("stream-thread-test-count-one-changelog", 0), 0L);
    offsets.put(new TopicPartition("stream-thread-test-count-two-changelog", 0), 0L);
    restoreConsumer.updateEndOffsets(offsets);
    restoreConsumer.updateBeginningOffsets(offsets);
    final Map<TaskId, Set<TopicPartition>> standbyTasks = new HashMap<>();
    final TopicPartition t1 = new TopicPartition("t1", 0);
    standbyTasks.put(new TaskId(0, 0), Utils.mkSet(t1));
    final Map<TaskId, Set<TopicPartition>> activeTasks = new HashMap<>();
    final TopicPartition t2 = new TopicPartition("t2", 0);
    activeTasks.put(new TaskId(1, 0), Utils.mkSet(t2));
    thread.partitionAssignor(new StreamPartitionAssignor() {

        @Override
        Map<TaskId, Set<TopicPartition>> standbyTasks() {
            return standbyTasks;
        }

        @Override
        Map<TaskId, Set<TopicPartition>> activeTasks() {
            return activeTasks;
        }
    });
    thread.rebalanceListener.onPartitionsRevoked(Collections.<TopicPartition>emptyList());
    thread.rebalanceListener.onPartitionsAssigned(Utils.mkSet(t2));
    // swap the assignment around and make sure we don't get any exceptions
    standbyTasks.clear();
    activeTasks.clear();
    standbyTasks.put(new TaskId(1, 0), Utils.mkSet(t2));
    activeTasks.put(new TaskId(0, 0), Utils.mkSet(t1));
    thread.rebalanceListener.onPartitionsRevoked(Collections.<TopicPartition>emptyList());
    thread.rebalanceListener.onPartitionsAssigned(Utils.mkSet(t1));
}