Examples with KinesisShardAssigner org.apache.flink.streaming.connectors.kinesis.KinesisShardAssigner used on opensource projects

Example 1 with KinesisShardAssigner

use of org.apache.flink.streaming.connectors.kinesis.KinesisShardAssigner in project flink by apache.

the class KinesisDataFetcherTest method testShardToSubtaskMappingWithCustomHashFunction.

// ----------------------------------------------------------------------
// Tests shard distribution with custom hash function
// ----------------------------------------------------------------------
@Test
public void testShardToSubtaskMappingWithCustomHashFunction() throws Exception {
    int totalCountOfSubtasks = 10;
    int shardCount = 3;
    for (int i = 0; i < 2; i++) {
        final int hash = i;
        final KinesisShardAssigner allShardsSingleSubtaskFn = (shard, subtasks) -> hash;
        Map<String, Integer> streamToShardCount = new HashMap<>();
        List<String> fakeStreams = new LinkedList<>();
        fakeStreams.add("fakeStream");
        streamToShardCount.put("fakeStream", shardCount);
        for (int j = 0; j < totalCountOfSubtasks; j++) {
            int subtaskIndex = j;
            // subscribe with default hashing
            final TestableKinesisDataFetcher fetcher = new TestableKinesisDataFetcher(fakeStreams, new TestSourceContext<>(), new Properties(), new KinesisDeserializationSchemaWrapper<>(new SimpleStringSchema()), totalCountOfSubtasks, subtaskIndex, new AtomicReference<>(), new LinkedList<>(), KinesisDataFetcher.createInitialSubscribedStreamsToLastDiscoveredShardsState(fakeStreams), FakeKinesisBehavioursFactory.nonReshardedStreamsBehaviour(streamToShardCount));
            Whitebox.setInternalState(fetcher, "shardAssigner", // override hashing
            allShardsSingleSubtaskFn);
            List<StreamShardHandle> shards = fetcher.discoverNewShardsToSubscribe();
            fetcher.shutdownFetcher();
            String msg = String.format("for hash=%d, subtask=%d", hash, subtaskIndex);
            if (j == i) {
                assertEquals(msg, shardCount, shards.size());
            } else {
                assertEquals(msg, 0, shards.size());
            }
        }
    }
}