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

Example 1 with ManualPartitioner

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

the class KafkaTridentSpoutEmitterPartitioningTest method testGetPartitionsForTask.

@Test
public void testGetPartitionsForTask() {
    // Verify correct wrapping/unwrapping of partition and delegation of partition assignment
    ManualPartitioner partitionerMock = mock(ManualPartitioner.class);
    when(partitionerMock.getPartitionsForThisTask(any(), any())).thenAnswer(invocation -> {
        List<TopicPartition> partitions = new ArrayList<>(invocation.getArgument(0));
        partitions.remove(0);
        return new HashSet<>(partitions);
    });
    KafkaTridentSpoutEmitter<String, String> emitter = new KafkaTridentSpoutEmitter<>(SingleTopicKafkaTridentSpoutConfiguration.createKafkaSpoutConfigBuilder(mock(TopicFilter.class), partitionerMock, -1).build(), topologyContextMock, config -> consumer, new TopicAssigner());
    List<KafkaTridentSpoutTopicPartition> allPartitions = new ArrayList<>();
    for (int i = 0; i < 10; i++) {
        allPartitions.add(new KafkaTridentSpoutTopicPartition(SingleTopicKafkaSpoutConfiguration.TOPIC, i));
    }
    List<TopicPartition> unwrappedPartitions = allPartitions.stream().map(kttp -> kttp.getTopicPartition()).collect(Collectors.toList());
    List<KafkaTridentSpoutTopicPartition> partitionsForTask = emitter.getPartitionsForTask(0, 2, allPartitions);
    verify(partitionerMock).getPartitionsForThisTask(eq(unwrappedPartitions), any(TopologyContext.class));
    allPartitions.remove(0);
    assertThat("Should have assigned all except the first partition to this task", new HashSet<>(partitionsForTask), is(new HashSet<>(allPartitions)));
}

Example 2 with ManualPartitioner

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