Examples with TopicAssigner org.apache.storm.kafka.spout.subscription.TopicAssigner used on opensource projects

Example 6 with TopicAssigner

use of org.apache.storm.kafka.spout.subscription.TopicAssigner in project storm by apache.

the class SpoutWithMockedConsumerSetupHelper method setupSpout.

/**
 * Creates, opens and activates a KafkaSpout using a mocked consumer. The TopicFilter and ManualPartitioner should be mock objects,
 * since this method shortcircuits the TopicPartition assignment process and just calls onPartitionsAssigned on the rebalance listener.
 *
 * @param <K> The Kafka key type
 * @param <V> The Kafka value type
 * @param spoutConfig The spout config to use
 * @param topoConf The topo conf to pass to the spout
 * @param contextMock The topo context to pass to the spout
 * @param collectorMock The mocked collector to pass to the spout
 * @param consumerMock The mocked consumer
 * @param assignedPartitions The partitions to assign to this spout. The consumer will act like these partitions are assigned to it.
 * @return The spout
 */
public static <K, V> KafkaSpout<K, V> setupSpout(KafkaSpoutConfig<K, V> spoutConfig, Map<String, Object> topoConf, TopologyContext contextMock, SpoutOutputCollector collectorMock, KafkaConsumer<K, V> consumerMock, TopicPartition... assignedPartitions) {
    TopicFilter topicFilter = spoutConfig.getTopicFilter();
    ManualPartitioner topicPartitioner = spoutConfig.getTopicPartitioner();
    if (!mockingDetails(topicFilter).isMock() || !mockingDetails(topicPartitioner).isMock()) {
        throw new IllegalStateException("Use a mocked TopicFilter and a mocked ManualPartitioner when using this method, it helps avoid complex stubbing");
    }
    Set<TopicPartition> assignedPartitionsSet = new HashSet<>(Arrays.asList(assignedPartitions));
    TopicAssigner assigner = mock(TopicAssigner.class);
    doAnswer(invocation -> {
        ConsumerRebalanceListener listener = invocation.getArgument(2);
        listener.onPartitionsAssigned(assignedPartitionsSet);
        return null;
    }).when(assigner).assignPartitions(any(), any(), any());
    when(consumerMock.assignment()).thenReturn(assignedPartitionsSet);
    ConsumerFactory<K, V> consumerFactory = (kafkaSpoutConfig) -> consumerMock;
    KafkaSpout<K, V> spout = new KafkaSpout<>(spoutConfig, consumerFactory, assigner);
    spout.open(topoConf, contextMock, collectorMock);
    spout.activate();
    return spout;
}

Example 7 with TopicAssigner

use of org.apache.storm.kafka.spout.subscription.TopicAssigner in project storm by apache.

the class KafkaSpoutRebalanceTest method spoutMustIgnoreAcksForTuplesItIsNotAssignedAfterRebalance.

@Test
public void spoutMustIgnoreAcksForTuplesItIsNotAssignedAfterRebalance() throws Exception {
    // Acking tuples for partitions that are no longer assigned is useless since the spout will not be allowed to commit them
    try (SimulatedTime simulatedTime = new SimulatedTime()) {
        TopicAssigner assignerMock = mock(TopicAssigner.class);
        KafkaSpout<String, String> spout = new KafkaSpout<>(createKafkaSpoutConfigBuilder(topicFilterMock, partitionerMock, -1).setOffsetCommitPeriodMs(offsetCommitPeriodMs).build(), consumerFactory, assignerMock);
        String topic = SingleTopicKafkaSpoutConfiguration.TOPIC;
        TopicPartition partitionThatWillBeRevoked = new TopicPartition(topic, 1);
        TopicPartition assignedPartition = new TopicPartition(topic, 2);
        // Emit a message on each partition and revoke the first partition
        List<KafkaSpoutMessageId> emittedMessageIds = emitOneMessagePerPartitionThenRevokeOnePartition(spout, partitionThatWillBeRevoked, assignedPartition, assignerMock);
        // Ack both emitted tuples
        spout.ack(emittedMessageIds.get(0));
        spout.ack(emittedMessageIds.get(1));
        // Ensure the commit timer has expired
        Time.advanceTime(offsetCommitPeriodMs + KafkaSpout.TIMER_DELAY_MS);
        // Make the spout commit any acked tuples
        spout.nextTuple();
        // Verify that it only committed the message on the assigned partition
        verify(consumerMock, times(1)).commitSync(commitCapture.capture());
        Map<TopicPartition, OffsetAndMetadata> commitCaptureMap = commitCapture.getValue();
        assertThat(commitCaptureMap, hasKey(assignedPartition));
        assertThat(commitCaptureMap, not(hasKey(partitionThatWillBeRevoked)));
    }
}