KinesisDataFetcher { protected static HashMap<String, String> createInitialSubscribedStreamsToLastDiscoveredShardsState(List<String> streams) { HashMap<String, String> initial = new HashMap<>(); for (String stream : streams) { initial.put(stream, null); } return initial; } KinesisDataFetcher(List<String> streams,
							SourceFunction.SourceContext<T> sourceContext,
							RuntimeContext runtimeContext,
							Properties configProps,
							KinesisDeserializationSchema<T> deserializationSchema); @VisibleForTesting protected  KinesisDataFetcher(List<String> streams,
								SourceFunction.SourceContext<T> sourceContext,
								Object checkpointLock,
								RuntimeContext runtimeContext,
								Properties configProps,
								KinesisDeserializationSchema<T> deserializationSchema,
								AtomicReference<Throwable> error,
								List<KinesisStreamShardState> subscribedShardsState,
								HashMap<String, String> subscribedStreamsToLastDiscoveredShardIds,
								KinesisProxyInterface kinesis); void runFetcher(); HashMap<StreamShardMetadata, SequenceNumber> snapshotState(); void shutdownFetcher(); void awaitTermination(); void advanceLastDiscoveredShardOfStream(String stream, String shardId); List<StreamShardHandle> discoverNewShardsToSubscribe(); int registerNewSubscribedShardState(KinesisStreamShardState newSubscribedShardState); static boolean isThisSubtaskShouldSubscribeTo(StreamShardHandle shard,
														int totalNumberOfConsumerSubtasks,
														int indexOfThisConsumerSubtask); @VisibleForTesting List<KinesisStreamShardState> getSubscribedShardsState(); static StreamShardMetadata convertToStreamShardMetadata(StreamShardHandle streamShardHandle); static StreamShardHandle convertToStreamShardHandle(StreamShardMetadata streamShardMetadata);  }
@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<String> fetcher = new TestableKinesisDataFetcher<>( fakeStreams, new TestSourceContext<>(), getStandardProperties(), new KinesisDeserializationSchemaWrapper<>(new SimpleStringSchema()), 10, 2, new AtomicReference<>(), new LinkedList<>(), subscribedStreamsToLastSeenShardIdsUnderTest, FakeKinesisBehavioursFactory.nonReshardedStreamsBehaviour(streamToShardCount)); final DummyFlinkKafkaConsumer<String> consumer = new DummyFlinkKafkaConsumer<>( getStandardProperties(), fetcher, 1, 0); CheckedThread consumerThread = new CheckedThread() { @Override public void go() throws Exception { consumer.run(new TestSourceContext<>()); } }; consumerThread.start(); fetcher.waitUntilRun(); consumer.cancel(); consumerThread.sync(); Set<String> streamsInState = subscribedStreamsToLastSeenShardIdsUnderTest.keySet(); assertEquals(fakeStreams.size(), streamsInState.size()); assertTrue(streamsInState.containsAll(fakeStreams)); for (Map.Entry<String, String> streamToLastSeenShard : subscribedStreamsToLastSeenShardIdsUnderTest.entrySet()) { assertEquals( KinesisShardIdGenerator.generateFromShardOrder(streamToShardCount.get(streamToLastSeenShard.getKey()) - 1), streamToLastSeenShard.getValue()); } }