Examples with ConsumerFactory org.apache.storm.kafka.spout.internal.ConsumerFactory used on opensource projects

Example 1 with ConsumerFactory

use of org.apache.storm.kafka.spout.internal.ConsumerFactory 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;
}