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

Example 6 with TestableKinesisDataFetcher

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

the class KinesisDataFetcherTest method testStreamToLastSeenShardStateIsCorrectlySetWhenNotRestoringFromFailure.

@Test
public void testStreamToLastSeenShardStateIsCorrectlySetWhenNotRestoringFromFailure() throws Exception {
    List<String> fakeStreams = new LinkedList<>();
    fakeStreams.add("fakeStream1");
    fakeStreams.add("fakeStream2");
    fakeStreams.add("fakeStream3");
    fakeStreams.add("fakeStream4");
    HashMap<String, String> subscribedStreamsToLastSeenShardIdsUnderTest = KinesisDataFetcher.createInitialSubscribedStreamsToLastDiscoveredShardsState(fakeStreams);
    Map<String, Integer> streamToShardCount = new HashMap<>();
    Random rand = new Random();
    for (String fakeStream : fakeStreams) {
        streamToShardCount.put(fakeStream, rand.nextInt(5) + 1);
    }
    final TestableKinesisDataFetcher fetcher = new TestableKinesisDataFetcher(fakeStreams, new Properties(), 10, 2, new AtomicReference<Throwable>(), new LinkedList<KinesisStreamShardState>(), subscribedStreamsToLastSeenShardIdsUnderTest, FakeKinesisBehavioursFactory.nonReshardedStreamsBehaviour(streamToShardCount));
    fetcher.setIsRestoringFromFailure(false);
    PowerMockito.whenNew(ShardConsumer.class).withAnyArguments().thenReturn(Mockito.mock(ShardConsumer.class));
    Thread runFetcherThread = new Thread(new Runnable() {

        @Override
        public void run() {
            try {
                fetcher.runFetcher();
            } catch (Exception e) {
            //
            }
        }
    });
    runFetcherThread.start();
    // sleep a while before closing
    Thread.sleep(1000);
    fetcher.shutdownFetcher();
    // assert that the streams tracked in the state are identical to the subscribed streams
    Set<String> streamsInState = subscribedStreamsToLastSeenShardIdsUnderTest.keySet();
    assertTrue(streamsInState.size() == fakeStreams.size());
    assertTrue(streamsInState.containsAll(fakeStreams));
    // assert that the last seen shards in state is correctly set
    for (Map.Entry<String, String> streamToLastSeenShard : subscribedStreamsToLastSeenShardIdsUnderTest.entrySet()) {
        assertTrue(streamToLastSeenShard.getValue().equals(KinesisShardIdGenerator.generateFromShardOrder(streamToShardCount.get(streamToLastSeenShard.getKey()) - 1)));
    }
}

Example 7 with TestableKinesisDataFetcher

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

the class KinesisDataFetcherTest method testIfNoShardsAreFoundShouldThrowException.

@Test(expected = RuntimeException.class)
public void testIfNoShardsAreFoundShouldThrowException() throws Exception {
    List<String> fakeStreams = new LinkedList<>();
    fakeStreams.add("fakeStream1");
    fakeStreams.add("fakeStream2");
    HashMap<String, String> subscribedStreamsToLastSeenShardIdsUnderTest = KinesisDataFetcher.createInitialSubscribedStreamsToLastDiscoveredShardsState(fakeStreams);
    TestableKinesisDataFetcher fetcher = new TestableKinesisDataFetcher(fakeStreams, new Properties(), 10, 2, new AtomicReference<Throwable>(), new LinkedList<KinesisStreamShardState>(), subscribedStreamsToLastSeenShardIdsUnderTest, FakeKinesisBehavioursFactory.noShardsFoundForRequestedStreamsBehaviour());
    // not restoring
    fetcher.setIsRestoringFromFailure(false);
    // this should throw RuntimeException
    fetcher.runFetcher();
}

Example 8 with TestableKinesisDataFetcher

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

the class ShardConsumerTest method testCorrectNumOfCollectedRecordsAndUpdatedState.

@Test
public void testCorrectNumOfCollectedRecordsAndUpdatedState() {
    KinesisStreamShard fakeToBeConsumedShard = new KinesisStreamShard("fakeStream", new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(0)).withHashKeyRange(new HashKeyRange().withStartingHashKey("0").withEndingHashKey(new BigInteger(StringUtils.repeat("FF", 16), 16).toString())));
    LinkedList<KinesisStreamShardState> subscribedShardsStateUnderTest = new LinkedList<>();
    subscribedShardsStateUnderTest.add(new KinesisStreamShardState(fakeToBeConsumedShard, new SequenceNumber("fakeStartingState")));
    TestableKinesisDataFetcher fetcher = new TestableKinesisDataFetcher(Collections.singletonList("fakeStream"), new Properties(), 10, 2, new AtomicReference<Throwable>(), subscribedShardsStateUnderTest, KinesisDataFetcher.createInitialSubscribedStreamsToLastDiscoveredShardsState(Collections.singletonList("fakeStream")), Mockito.mock(KinesisProxyInterface.class));
    new ShardConsumer<>(fetcher, 0, subscribedShardsStateUnderTest.get(0).getKinesisStreamShard(), subscribedShardsStateUnderTest.get(0).getLastProcessedSequenceNum(), FakeKinesisBehavioursFactory.totalNumOfRecordsAfterNumOfGetRecordsCalls(1000, 9)).run();
    assertTrue(fetcher.getNumOfElementsCollected() == 1000);
    assertTrue(subscribedShardsStateUnderTest.get(0).getLastProcessedSequenceNum().equals(SentinelSequenceNumber.SENTINEL_SHARD_ENDING_SEQUENCE_NUM.get()));
}