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

Example 6 with MockClientSupplier

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

the class StreamThreadTest method testPartitionAssignmentChange.

@SuppressWarnings("unchecked")
@Test
public void testPartitionAssignmentChange() throws Exception {
    StreamsConfig config = new StreamsConfig(configProps());
    StateListenerStub stateListener = new StateListenerStub();
    TopologyBuilder builder = new TopologyBuilder().setApplicationId("X");
    builder.addSource("source1", "topic1");
    builder.addSource("source2", "topic2");
    builder.addSource("source3", "topic3");
    builder.addProcessor("processor", new MockProcessorSupplier(), "source2", "source3");
    MockClientSupplier mockClientSupplier = new MockClientSupplier();
    StreamThread thread = new StreamThread(builder, config, mockClientSupplier, applicationId, clientId, processId, new Metrics(), Time.SYSTEM, new StreamsMetadataState(builder, StreamsMetadataState.UNKNOWN_HOST), 0) {

        @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);
        }
    };
    thread.setStateListener(stateListener);
    assertEquals(thread.state(), StreamThread.State.RUNNING);
    initPartitionGrouper(config, thread, mockClientSupplier);
    ConsumerRebalanceListener rebalanceListener = thread.rebalanceListener;
    assertTrue(thread.tasks().isEmpty());
    List<TopicPartition> revokedPartitions;
    List<TopicPartition> assignedPartitions;
    Set<TopicPartition> expectedGroup1;
    Set<TopicPartition> expectedGroup2;
    revokedPartitions = Collections.emptyList();
    assignedPartitions = Collections.singletonList(t1p1);
    expectedGroup1 = new HashSet<>(Arrays.asList(t1p1));
    rebalanceListener.onPartitionsRevoked(revokedPartitions);
    assertEquals(thread.state(), StreamThread.State.PARTITIONS_REVOKED);
    Assert.assertEquals(stateListener.numChanges, 1);
    Assert.assertEquals(stateListener.oldState, StreamThread.State.RUNNING);
    Assert.assertEquals(stateListener.newState, StreamThread.State.PARTITIONS_REVOKED);
    rebalanceListener.onPartitionsAssigned(assignedPartitions);
    assertEquals(thread.state(), StreamThread.State.RUNNING);
    Assert.assertEquals(stateListener.numChanges, 3);
    Assert.assertEquals(stateListener.oldState, StreamThread.State.ASSIGNING_PARTITIONS);
    Assert.assertEquals(stateListener.newState, StreamThread.State.RUNNING);
    assertTrue(thread.tasks().containsKey(task1));
    assertEquals(expectedGroup1, thread.tasks().get(task1).partitions());
    assertEquals(1, thread.tasks().size());
    revokedPartitions = assignedPartitions;
    assignedPartitions = Collections.singletonList(t1p2);
    expectedGroup2 = new HashSet<>(Arrays.asList(t1p2));
    rebalanceListener.onPartitionsRevoked(revokedPartitions);
    assertFalse(thread.tasks().containsKey(task1));
    assertEquals(0, thread.tasks().size());
    rebalanceListener.onPartitionsAssigned(assignedPartitions);
    assertTrue(thread.tasks().containsKey(task2));
    assertEquals(expectedGroup2, thread.tasks().get(task2).partitions());
    assertEquals(1, thread.tasks().size());
    revokedPartitions = assignedPartitions;
    assignedPartitions = Arrays.asList(t1p1, t1p2);
    expectedGroup1 = new HashSet<>(Collections.singleton(t1p1));
    expectedGroup2 = new HashSet<>(Collections.singleton(t1p2));
    rebalanceListener.onPartitionsRevoked(revokedPartitions);
    rebalanceListener.onPartitionsAssigned(assignedPartitions);
    assertTrue(thread.tasks().containsKey(task1));
    assertTrue(thread.tasks().containsKey(task2));
    assertEquals(expectedGroup1, thread.tasks().get(task1).partitions());
    assertEquals(expectedGroup2, thread.tasks().get(task2).partitions());
    assertEquals(2, thread.tasks().size());
    revokedPartitions = assignedPartitions;
    assignedPartitions = Arrays.asList(t2p1, t2p2, t3p1, t3p2);
    expectedGroup1 = new HashSet<>(Arrays.asList(t2p1, t3p1));
    expectedGroup2 = new HashSet<>(Arrays.asList(t2p2, t3p2));
    rebalanceListener.onPartitionsRevoked(revokedPartitions);
    rebalanceListener.onPartitionsAssigned(assignedPartitions);
    assertTrue(thread.tasks().containsKey(task4));
    assertTrue(thread.tasks().containsKey(task5));
    assertEquals(expectedGroup1, thread.tasks().get(task4).partitions());
    assertEquals(expectedGroup2, thread.tasks().get(task5).partitions());
    assertEquals(2, thread.tasks().size());
    revokedPartitions = assignedPartitions;
    assignedPartitions = Arrays.asList(t1p1, t2p1, t3p1);
    expectedGroup1 = new HashSet<>(Arrays.asList(t1p1));
    expectedGroup2 = new HashSet<>(Arrays.asList(t2p1, t3p1));
    rebalanceListener.onPartitionsRevoked(revokedPartitions);
    rebalanceListener.onPartitionsAssigned(assignedPartitions);
    assertTrue(thread.tasks().containsKey(task1));
    assertTrue(thread.tasks().containsKey(task4));
    assertEquals(expectedGroup1, thread.tasks().get(task1).partitions());
    assertEquals(expectedGroup2, thread.tasks().get(task4).partitions());
    assertEquals(2, thread.tasks().size());
    revokedPartitions = assignedPartitions;
    assignedPartitions = Arrays.asList(t1p1, t2p1, t3p1);
    expectedGroup1 = new HashSet<>(Arrays.asList(t1p1));
    expectedGroup2 = new HashSet<>(Arrays.asList(t2p1, t3p1));
    rebalanceListener.onPartitionsRevoked(revokedPartitions);
    rebalanceListener.onPartitionsAssigned(assignedPartitions);
    assertTrue(thread.tasks().containsKey(task1));
    assertTrue(thread.tasks().containsKey(task4));
    assertEquals(expectedGroup1, thread.tasks().get(task1).partitions());
    assertEquals(expectedGroup2, thread.tasks().get(task4).partitions());
    assertEquals(2, thread.tasks().size());
    revokedPartitions = assignedPartitions;
    assignedPartitions = Collections.emptyList();
    rebalanceListener.onPartitionsRevoked(revokedPartitions);
    rebalanceListener.onPartitionsAssigned(assignedPartitions);
    assertTrue(thread.tasks().isEmpty());
    thread.close();
    assertTrue((thread.state() == StreamThread.State.PENDING_SHUTDOWN) || (thread.state() == StreamThread.State.NOT_RUNNING));
}

Example 7 with MockClientSupplier

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

the class StreamPartitionAssignorTest method testSubscription.

@SuppressWarnings("unchecked")
@Test
public void testSubscription() throws Exception {
    builder.addSource("source1", "topic1");
    builder.addSource("source2", "topic2");
    builder.addProcessor("processor", new MockProcessorSupplier(), "source1", "source2");
    final Set<TaskId> prevTasks = Utils.mkSet(new TaskId(0, 1), new TaskId(1, 1), new TaskId(2, 1));
    final Set<TaskId> cachedTasks = Utils.mkSet(new TaskId(0, 1), new TaskId(1, 1), new TaskId(2, 1), new TaskId(0, 2), new TaskId(1, 2), new TaskId(2, 2));
    String clientId = "client-id";
    UUID processId = UUID.randomUUID();
    StreamThread thread = new StreamThread(builder, config, new MockClientSupplier(), "test", clientId, processId, new Metrics(), Time.SYSTEM, new StreamsMetadataState(builder, StreamsMetadataState.UNKNOWN_HOST), 0) {

        @Override
        public Set<TaskId> prevActiveTasks() {
            return prevTasks;
        }

        @Override
        public Set<TaskId> cachedTasks() {
            return cachedTasks;
        }
    };
    partitionAssignor.configure(config.getConsumerConfigs(thread, "test", clientId));
    PartitionAssignor.Subscription subscription = partitionAssignor.subscription(Utils.mkSet("topic1", "topic2"));
    Collections.sort(subscription.topics());
    assertEquals(Utils.mkList("topic1", "topic2"), subscription.topics());
    Set<TaskId> standbyTasks = new HashSet<>(cachedTasks);
    standbyTasks.removeAll(prevTasks);
    SubscriptionInfo info = new SubscriptionInfo(processId, prevTasks, standbyTasks, null);
    assertEquals(info.encode(), subscription.userData());
}

Example 8 with MockClientSupplier

use of org.apache.kafka.test.MockClientSupplier in project apache-kafka-on-k8s by banzaicloud.

the class KafkaStreamsTest method shouldCleanupResourcesOnCloseWithoutPreviousStart.

@Test
public void shouldCleanupResourcesOnCloseWithoutPreviousStart() throws Exception {
    final StreamsBuilder builder = new StreamsBuilder();
    builder.globalTable("anyTopic");
    List<Node> nodes = Arrays.asList(new Node(0, "localhost", 8121));
    Cluster cluster = new Cluster("mockClusterId", nodes, Collections.<PartitionInfo>emptySet(), Collections.<String>emptySet(), Collections.<String>emptySet(), nodes.get(0));
    MockClientSupplier clientSupplier = new MockClientSupplier();
    clientSupplier.setClusterForAdminClient(cluster);
    final KafkaStreams streams = new KafkaStreams(builder.build(), props, clientSupplier);
    streams.close();
    TestUtils.waitForCondition(new TestCondition() {

        @Override
        public boolean conditionMet() {
            return streams.state() == KafkaStreams.State.NOT_RUNNING;
        }
    }, 10 * 1000, "Streams never stopped.");
    // Ensure that any created clients are closed
    assertTrue(clientSupplier.consumer.closed());
    assertTrue(clientSupplier.restoreConsumer.closed());
    for (MockProducer p : clientSupplier.producers) {
        assertTrue(p.closed());
    }
}

Example 9 with MockClientSupplier

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

the class KafkaStreamsTest method shouldReturnEmptyLocalStorePartitionLags.

@Test
public void shouldReturnEmptyLocalStorePartitionLags() {
    // Mock all calls made to compute the offset lags,
    final ListOffsetsResult result = EasyMock.mock(ListOffsetsResult.class);
    final KafkaFutureImpl<Map<TopicPartition, ListOffsetsResultInfo>> allFuture = new KafkaFutureImpl<>();
    allFuture.complete(Collections.emptyMap());
    EasyMock.expect(result.all()).andReturn(allFuture);
    final MockAdminClient mockAdminClient = EasyMock.partialMockBuilder(MockAdminClient.class).addMockedMethod("listOffsets", Map.class).createMock();
    EasyMock.expect(mockAdminClient.listOffsets(anyObject())).andStubReturn(result);
    final MockClientSupplier mockClientSupplier = EasyMock.partialMockBuilder(MockClientSupplier.class).addMockedMethod("getAdmin").createMock();
    EasyMock.expect(mockClientSupplier.getAdmin(anyObject())).andReturn(mockAdminClient);
    EasyMock.replay(result, mockAdminClient, mockClientSupplier);
    try (final KafkaStreams streams = new KafkaStreams(getBuilderWithSource().build(), props, mockClientSupplier, time)) {
        streams.start();
        assertEquals(0, streams.allLocalStorePartitionLags().size());
    }
}

Example 10 with MockClientSupplier

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

the class StreamsProducerTest method shouldThrowTimeoutExceptionOnEosInitTxTimeout.

@Test
public void shouldThrowTimeoutExceptionOnEosInitTxTimeout() {
    // use `nonEosMockProducer` instead of `eosMockProducer` to avoid double Tx-Init
    nonEosMockProducer.initTransactionException = new TimeoutException("KABOOM!");
    final KafkaClientSupplier clientSupplier = new MockClientSupplier() {

        @Override
        public Producer<byte[], byte[]> getProducer(final Map<String, Object> config) {
            return nonEosMockProducer;
        }
    };
    final StreamsProducer streamsProducer = new StreamsProducer(eosAlphaConfig, "threadId", clientSupplier, new TaskId(0, 0), null, logContext, mockTime);
    final TimeoutException thrown = assertThrows(TimeoutException.class, streamsProducer::initTransaction);
    assertThat(thrown.getMessage(), is("KABOOM!"));
}