Examples with KafkaPartitioner org.apache.flink.streaming.connectors.kafka.partitioner.KafkaPartitioner used on opensource projects

Example 1 with KafkaPartitioner

use of org.apache.flink.streaming.connectors.kafka.partitioner.KafkaPartitioner in project flink by apache.

the class FlinkKafkaProducerBaseTest method testPartitionerOpenedWithDeterminatePartitionList.

/**
	 * Tests that partitions list is determinate and correctly provided to custom partitioner
	 */
@Test
public void testPartitionerOpenedWithDeterminatePartitionList() throws Exception {
    KafkaPartitioner mockPartitioner = mock(KafkaPartitioner.class);
    RuntimeContext mockRuntimeContext = mock(RuntimeContext.class);
    when(mockRuntimeContext.getIndexOfThisSubtask()).thenReturn(0);
    when(mockRuntimeContext.getNumberOfParallelSubtasks()).thenReturn(1);
    // out-of-order list of 4 partitions
    List<PartitionInfo> mockPartitionsList = new ArrayList<>(4);
    mockPartitionsList.add(new PartitionInfo(DummyFlinkKafkaProducer.DUMMY_TOPIC, 3, null, null, null));
    mockPartitionsList.add(new PartitionInfo(DummyFlinkKafkaProducer.DUMMY_TOPIC, 1, null, null, null));
    mockPartitionsList.add(new PartitionInfo(DummyFlinkKafkaProducer.DUMMY_TOPIC, 0, null, null, null));
    mockPartitionsList.add(new PartitionInfo(DummyFlinkKafkaProducer.DUMMY_TOPIC, 2, null, null, null));
    final DummyFlinkKafkaProducer producer = new DummyFlinkKafkaProducer(FakeStandardProducerConfig.get(), mockPartitioner);
    producer.setRuntimeContext(mockRuntimeContext);
    final KafkaProducer mockProducer = producer.getMockKafkaProducer();
    when(mockProducer.partitionsFor(anyString())).thenReturn(mockPartitionsList);
    when(mockProducer.metrics()).thenReturn(null);
    producer.open(new Configuration());
    // the out-of-order partitions list should be sorted before provided to the custom partitioner's open() method
    int[] correctPartitionList = { 0, 1, 2, 3 };
    verify(mockPartitioner).open(0, 1, correctPartitionList);
}