Examples with TestPartitionDiscoverer org.apache.flink.streaming.connectors.kafka.testutils.TestPartitionDiscoverer used on opensource projects

Example 6 with TestPartitionDiscoverer

use of org.apache.flink.streaming.connectors.kafka.testutils.TestPartitionDiscoverer in project flink by apache.

the class AbstractPartitionDiscovererTest method testPartitionsEqualConsumersFixedPartitions.

@Test
public void testPartitionsEqualConsumersFixedPartitions() throws Exception {
    List<KafkaTopicPartition> mockGetAllPartitionsForTopicsReturn = Arrays.asList(new KafkaTopicPartition(TEST_TOPIC, 0), new KafkaTopicPartition(TEST_TOPIC, 1), new KafkaTopicPartition(TEST_TOPIC, 2), new KafkaTopicPartition(TEST_TOPIC, 3));
    int numSubtasks = mockGetAllPartitionsForTopicsReturn.size();
    // get the start index; the assertions below will fail if the assignment logic does not meet
    // correct contracts
    int numConsumers = KafkaTopicPartitionAssigner.assign(mockGetAllPartitionsForTopicsReturn.get(0), numSubtasks);
    for (int subtaskIndex = 0; subtaskIndex < mockGetAllPartitionsForTopicsReturn.size(); subtaskIndex++) {
        TestPartitionDiscoverer partitionDiscoverer = new TestPartitionDiscoverer(topicsDescriptor, subtaskIndex, mockGetAllPartitionsForTopicsReturn.size(), TestPartitionDiscoverer.createMockGetAllTopicsSequenceFromFixedReturn(Collections.singletonList(TEST_TOPIC)), TestPartitionDiscoverer.createMockGetAllPartitionsFromTopicsSequenceFromFixedReturn(mockGetAllPartitionsForTopicsReturn));
        partitionDiscoverer.open();
        List<KafkaTopicPartition> initialDiscovery = partitionDiscoverer.discoverPartitions();
        assertEquals(1, initialDiscovery.size());
        assertTrue(contains(mockGetAllPartitionsForTopicsReturn, initialDiscovery.get(0).getPartition()));
        assertEquals(getExpectedSubtaskIndex(initialDiscovery.get(0), numConsumers, numSubtasks), subtaskIndex);
        // subsequent discoveries should not find anything
        List<KafkaTopicPartition> secondDiscovery = partitionDiscoverer.discoverPartitions();
        List<KafkaTopicPartition> thirdDiscovery = partitionDiscoverer.discoverPartitions();
        assertEquals(0, secondDiscovery.size());
        assertEquals(0, thirdDiscovery.size());
    }
}

Example 7 with TestPartitionDiscoverer

use of org.apache.flink.streaming.connectors.kafka.testutils.TestPartitionDiscoverer in project flink by apache.

the class AbstractPartitionDiscovererTest method testNonContiguousPartitionIdDiscovery.

@Test
public void testNonContiguousPartitionIdDiscovery() throws Exception {
    List<KafkaTopicPartition> mockGetAllPartitionsForTopicsReturn1 = Arrays.asList(new KafkaTopicPartition("test-topic", 1), new KafkaTopicPartition("test-topic", 4));
    List<KafkaTopicPartition> mockGetAllPartitionsForTopicsReturn2 = Arrays.asList(new KafkaTopicPartition("test-topic", 0), new KafkaTopicPartition("test-topic", 1), new KafkaTopicPartition("test-topic", 2), new KafkaTopicPartition("test-topic", 3), new KafkaTopicPartition("test-topic", 4));
    TestPartitionDiscoverer partitionDiscoverer = new TestPartitionDiscoverer(topicsDescriptor, 0, 1, TestPartitionDiscoverer.createMockGetAllTopicsSequenceFromFixedReturn(Collections.singletonList("test-topic")), // Iterable.remove() is used.
    Arrays.asList(new ArrayList<>(mockGetAllPartitionsForTopicsReturn1), new ArrayList<>(mockGetAllPartitionsForTopicsReturn2)));
    partitionDiscoverer.open();
    List<KafkaTopicPartition> discoveredPartitions1 = partitionDiscoverer.discoverPartitions();
    assertEquals(2, discoveredPartitions1.size());
    assertTrue(discoveredPartitions1.contains(new KafkaTopicPartition("test-topic", 1)));
    assertTrue(discoveredPartitions1.contains(new KafkaTopicPartition("test-topic", 4)));
    List<KafkaTopicPartition> discoveredPartitions2 = partitionDiscoverer.discoverPartitions();
    assertEquals(3, discoveredPartitions2.size());
    assertTrue(discoveredPartitions2.contains(new KafkaTopicPartition("test-topic", 0)));
    assertTrue(discoveredPartitions2.contains(new KafkaTopicPartition("test-topic", 2)));
    assertTrue(discoveredPartitions2.contains(new KafkaTopicPartition("test-topic", 3)));
}

Example 8 with TestPartitionDiscoverer

use of org.apache.flink.streaming.connectors.kafka.testutils.TestPartitionDiscoverer in project flink by apache.

the class AbstractPartitionDiscovererTest method testMultiplePartitionsPerConsumersFixedPartitions.

@Test
public void testMultiplePartitionsPerConsumersFixedPartitions() {
    try {
        final int[] partitionIDs = { 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10 };
        final List<KafkaTopicPartition> mockGetAllPartitionsForTopicsReturn = new ArrayList<>();
        final Set<KafkaTopicPartition> allPartitions = new HashSet<>();
        for (int p : partitionIDs) {
            KafkaTopicPartition part = new KafkaTopicPartition(TEST_TOPIC, p);
            mockGetAllPartitionsForTopicsReturn.add(part);
            allPartitions.add(part);
        }
        final int numConsumers = 3;
        final int minPartitionsPerConsumer = mockGetAllPartitionsForTopicsReturn.size() / numConsumers;
        final int maxPartitionsPerConsumer = mockGetAllPartitionsForTopicsReturn.size() / numConsumers + 1;
        // get the start index; the assertions below will fail if the assignment logic does not
        // meet correct contracts
        int startIndex = KafkaTopicPartitionAssigner.assign(mockGetAllPartitionsForTopicsReturn.get(0), numConsumers);
        for (int subtaskIndex = 0; subtaskIndex < numConsumers; subtaskIndex++) {
            TestPartitionDiscoverer partitionDiscoverer = new TestPartitionDiscoverer(topicsDescriptor, subtaskIndex, numConsumers, TestPartitionDiscoverer.createMockGetAllTopicsSequenceFromFixedReturn(Collections.singletonList(TEST_TOPIC)), TestPartitionDiscoverer.createMockGetAllPartitionsFromTopicsSequenceFromFixedReturn(mockGetAllPartitionsForTopicsReturn));
            partitionDiscoverer.open();
            List<KafkaTopicPartition> initialDiscovery = partitionDiscoverer.discoverPartitions();
            assertTrue(initialDiscovery.size() >= minPartitionsPerConsumer);
            assertTrue(initialDiscovery.size() <= maxPartitionsPerConsumer);
            for (KafkaTopicPartition p : initialDiscovery) {
                // check that the element was actually contained
                assertTrue(allPartitions.remove(p));
                assertEquals(getExpectedSubtaskIndex(p, startIndex, numConsumers), subtaskIndex);
            }
            // subsequent discoveries should not find anything
            List<KafkaTopicPartition> secondDiscovery = partitionDiscoverer.discoverPartitions();
            List<KafkaTopicPartition> thirdDiscovery = partitionDiscoverer.discoverPartitions();
            assertEquals(0, secondDiscovery.size());
            assertEquals(0, thirdDiscovery.size());
        }
        // all partitions must have been assigned
        assertTrue(allPartitions.isEmpty());
    } catch (Exception e) {
        e.printStackTrace();
        fail(e.getMessage());
    }
}