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

Example 1 with TestableFlinkKinesisConsumer

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

the class FlinkKinesisConsumerTest method testFetcherShouldNotBeRestoringFromFailureIfNotRestoringFromCheckpoint.

// ----------------------------------------------------------------------
// Tests related to fetcher initialization
// ----------------------------------------------------------------------
@Test
@SuppressWarnings("unchecked")
public void testFetcherShouldNotBeRestoringFromFailureIfNotRestoringFromCheckpoint() throws Exception {
    KinesisDataFetcher mockedFetcher = Mockito.mock(KinesisDataFetcher.class);
    PowerMockito.whenNew(KinesisDataFetcher.class).withAnyArguments().thenReturn(mockedFetcher);
    // assume the given config is correct
    PowerMockito.mockStatic(KinesisConfigUtil.class);
    PowerMockito.doNothing().when(KinesisConfigUtil.class);
    TestableFlinkKinesisConsumer consumer = new TestableFlinkKinesisConsumer("fakeStream", new Properties(), 10, 2);
    consumer.open(new Configuration());
    consumer.run(Mockito.mock(SourceFunction.SourceContext.class));
    Mockito.verify(mockedFetcher).setIsRestoringFromFailure(false);
}

Example 2 with TestableFlinkKinesisConsumer

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

the class FlinkKinesisConsumerTest method testFetcherShouldBeCorrectlySeededIfRestoringFromCheckpoint.

@Test
@SuppressWarnings("unchecked")
public void testFetcherShouldBeCorrectlySeededIfRestoringFromCheckpoint() throws Exception {
    KinesisDataFetcher mockedFetcher = Mockito.mock(KinesisDataFetcher.class);
    PowerMockito.whenNew(KinesisDataFetcher.class).withAnyArguments().thenReturn(mockedFetcher);
    // assume the given config is correct
    PowerMockito.mockStatic(KinesisConfigUtil.class);
    PowerMockito.doNothing().when(KinesisConfigUtil.class);
    HashMap<KinesisStreamShard, SequenceNumber> fakeRestoredState = new HashMap<>();
    fakeRestoredState.put(new KinesisStreamShard("fakeStream1", new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(0))), new SequenceNumber(UUID.randomUUID().toString()));
    fakeRestoredState.put(new KinesisStreamShard("fakeStream1", new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(1))), new SequenceNumber(UUID.randomUUID().toString()));
    fakeRestoredState.put(new KinesisStreamShard("fakeStream1", new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(2))), new SequenceNumber(UUID.randomUUID().toString()));
    fakeRestoredState.put(new KinesisStreamShard("fakeStream2", new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(0))), new SequenceNumber(UUID.randomUUID().toString()));
    fakeRestoredState.put(new KinesisStreamShard("fakeStream2", new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(1))), new SequenceNumber(UUID.randomUUID().toString()));
    TestableFlinkKinesisConsumer consumer = new TestableFlinkKinesisConsumer("fakeStream", new Properties(), 10, 2);
    consumer.restoreState(fakeRestoredState);
    consumer.open(new Configuration());
    consumer.run(Mockito.mock(SourceFunction.SourceContext.class));
    Mockito.verify(mockedFetcher).setIsRestoringFromFailure(true);
    for (Map.Entry<KinesisStreamShard, SequenceNumber> restoredShard : fakeRestoredState.entrySet()) {
        Mockito.verify(mockedFetcher).advanceLastDiscoveredShardOfStream(restoredShard.getKey().getStreamName(), restoredShard.getKey().getShard().getShardId());
        Mockito.verify(mockedFetcher).registerNewSubscribedShardState(new KinesisStreamShardState(restoredShard.getKey(), restoredShard.getValue()));
    }
}