org.apache.flink.streaming.connectors.kinesis.testutils.KinesisShardIdGenerator Java Examples

@Test
public void testGetShardWithNoNewShards() throws Exception {
	// given
	String fakeStreamName = "fake-stream";

	AmazonKinesis mockClient = mock(AmazonKinesis.class);
	KinesisProxy kinesisProxy = getProxy(mockClient);

	Mockito.when(mockClient.listShards(
		new ListShardsRequest()
		.withStreamName(fakeStreamName)
		.withExclusiveStartShardId(KinesisShardIdGenerator.generateFromShardOrder(1))
	)).thenReturn(new ListShardsResult().withShards(Collections.emptyList()));

	HashMap<String, String> streamHashMap = new HashMap<>();
	streamHashMap.put(fakeStreamName, KinesisShardIdGenerator.generateFromShardOrder(1));

	// when
	GetShardListResult shardListResult = kinesisProxy.getShardList(streamHashMap);

	// then
	Assert.assertFalse(shardListResult.hasRetrievedShards());
}
 

@Test
public void testGetShardWithNoNewShards() throws Exception {
	// given
	String fakeStreamName = "fake-stream";

	AmazonKinesis mockClient = mock(AmazonKinesis.class);
	KinesisProxy kinesisProxy = getProxy(mockClient);

	Mockito.when(mockClient.listShards(
		new ListShardsRequest()
		.withStreamName(fakeStreamName)
		.withExclusiveStartShardId(KinesisShardIdGenerator.generateFromShardOrder(1))
	)).thenReturn(new ListShardsResult().withShards(Collections.emptyList()));

	HashMap<String, String> streamHashMap = new HashMap<>();
	streamHashMap.put(fakeStreamName, KinesisShardIdGenerator.generateFromShardOrder(1));

	// when
	GetShardListResult shardListResult = kinesisProxy.getShardList(streamHashMap);

	// then
	Assert.assertFalse(shardListResult.hasRetrievedShards());
}
 

@Test
public void testGetShardWithNoNewShards() throws Exception {
	// given
	String fakeStreamName = "fake-stream";

	AmazonKinesis mockClient = mock(AmazonKinesis.class);
	KinesisProxy kinesisProxy = getProxy(mockClient);

	Mockito.when(mockClient.listShards(
		new ListShardsRequest()
		.withStreamName(fakeStreamName)
		.withExclusiveStartShardId(KinesisShardIdGenerator.generateFromShardOrder(1))
	)).thenReturn(new ListShardsResult().withShards(Collections.emptyList()));

	HashMap<String, String> streamHashMap = new HashMap<>();
	streamHashMap.put(fakeStreamName, KinesisShardIdGenerator.generateFromShardOrder(1));

	// when
	GetShardListResult shardListResult = kinesisProxy.getShardList(streamHashMap);

	// then
	Assert.assertFalse(shardListResult.hasRetrievedShards());
}
 

private HashMap<StreamShardHandle, SequenceNumber> getFakeRestoredStore(String streamName) {
	HashMap<StreamShardHandle, SequenceNumber> fakeRestoredState = new HashMap<>();

	if (streamName.equals("fakeStream1") || streamName.equals("all")) {
		fakeRestoredState.put(
			new StreamShardHandle("fakeStream1",
				new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(0))),
			new SequenceNumber(UUID.randomUUID().toString()));
		fakeRestoredState.put(
			new StreamShardHandle("fakeStream1",
				new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(1))),
			new SequenceNumber(UUID.randomUUID().toString()));
		fakeRestoredState.put(
			new StreamShardHandle("fakeStream1",
				new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(2))),
			new SequenceNumber(UUID.randomUUID().toString()));
	}

	if (streamName.equals("fakeStream2") || streamName.equals("all")) {
		fakeRestoredState.put(
			new StreamShardHandle("fakeStream2",
				new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(0))),
			new SequenceNumber(UUID.randomUUID().toString()));
		fakeRestoredState.put(
			new StreamShardHandle("fakeStream2",
				new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(1))),
			new SequenceNumber(UUID.randomUUID().toString()));
	}

	return fakeRestoredState;
}
 

private static StreamShardHandle getMockStreamShard(String streamName, int shardId) {
	return new StreamShardHandle(
		streamName,
		new Shard()
			.withShardId(KinesisShardIdGenerator.generateFromShardOrder(shardId))
			.withHashKeyRange(
				new HashKeyRange()
					.withStartingHashKey("0")
					.withEndingHashKey(new BigInteger(StringUtils.repeat("FF", 16), 16).toString())));
}
 

private HashMap<StreamShardHandle, SequenceNumber> getFakeRestoredStore(String streamName) {
	HashMap<StreamShardHandle, SequenceNumber> fakeRestoredState = new HashMap<>();

	if (streamName.equals("fakeStream1") || streamName.equals("all")) {
		fakeRestoredState.put(
			new StreamShardHandle("fakeStream1",
				new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(0))),
			new SequenceNumber(UUID.randomUUID().toString()));
		fakeRestoredState.put(
			new StreamShardHandle("fakeStream1",
				new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(1))),
			new SequenceNumber(UUID.randomUUID().toString()));
		fakeRestoredState.put(
			new StreamShardHandle("fakeStream1",
				new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(2))),
			new SequenceNumber(UUID.randomUUID().toString()));
	}

	if (streamName.equals("fakeStream2") || streamName.equals("all")) {
		fakeRestoredState.put(
			new StreamShardHandle("fakeStream2",
				new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(0))),
			new SequenceNumber(UUID.randomUUID().toString()));
		fakeRestoredState.put(
			new StreamShardHandle("fakeStream2",
				new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(1))),
			new SequenceNumber(UUID.randomUUID().toString()));
	}

	return fakeRestoredState;
}
 

private static StreamShardHandle getMockStreamShard(String streamName, int shardId) {
	return new StreamShardHandle(
		streamName,
		new Shard()
			.withShardId(KinesisShardIdGenerator.generateFromShardOrder(shardId))
			.withHashKeyRange(
				new HashKeyRange()
					.withStartingHashKey("0")
					.withEndingHashKey(new BigInteger(StringUtils.repeat("FF", 16), 16).toString())));
}
 

private static StreamShardHandle getMockStreamShard(String streamName, int shardId) {
	return new StreamShardHandle(
		streamName,
		new Shard()
			.withShardId(KinesisShardIdGenerator.generateFromShardOrder(shardId))
			.withHashKeyRange(
				new HashKeyRange()
					.withStartingHashKey("0")
					.withEndingHashKey(new BigInteger(StringUtils.repeat("FF", 16), 16).toString())));
}
 

private HashMap<StreamShardHandle, SequenceNumber> getFakeRestoredStore(String streamName) {
	HashMap<StreamShardHandle, SequenceNumber> fakeRestoredState = new HashMap<>();

	if (streamName.equals("fakeStream1") || streamName.equals("all")) {
		fakeRestoredState.put(
			new StreamShardHandle("fakeStream1",
				new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(0))),
			new SequenceNumber(UUID.randomUUID().toString()));
		fakeRestoredState.put(
			new StreamShardHandle("fakeStream1",
				new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(1))),
			new SequenceNumber(UUID.randomUUID().toString()));
		fakeRestoredState.put(
			new StreamShardHandle("fakeStream1",
				new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(2))),
			new SequenceNumber(UUID.randomUUID().toString()));
	}

	if (streamName.equals("fakeStream2") || streamName.equals("all")) {
		fakeRestoredState.put(
			new StreamShardHandle("fakeStream2",
				new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(0))),
			new SequenceNumber(UUID.randomUUID().toString()));
		fakeRestoredState.put(
			new StreamShardHandle("fakeStream2",
				new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(1))),
			new SequenceNumber(UUID.randomUUID().toString()));
	}

	return fakeRestoredState;
}
 

@Test
public void testStreamToLastSeenShardStateIsCorrectlySetWhenNewShardsFoundSinceRestoredCheckpointAndSomeStreamsDoNotExist() throws Exception {
	List<String> fakeStreams = new LinkedList<>();
	fakeStreams.add("fakeStream1");
	fakeStreams.add("fakeStream2");
	fakeStreams.add("fakeStream3"); // fakeStream3 will not have any shards
	fakeStreams.add("fakeStream4"); // fakeStream4 will not have any shards

	Map<StreamShardHandle, String> restoredStateUnderTest = new HashMap<>();

	// fakeStream1 has 3 shards before restore
	restoredStateUnderTest.put(
		new StreamShardHandle(
			"fakeStream1",
			new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(0))),
		UUID.randomUUID().toString());
	restoredStateUnderTest.put(
		new StreamShardHandle(
			"fakeStream1",
			new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(1))),
		UUID.randomUUID().toString());
	restoredStateUnderTest.put(
		new StreamShardHandle(
			"fakeStream1",
			new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(2))),
		UUID.randomUUID().toString());

	// fakeStream2 has 2 shards before restore
	restoredStateUnderTest.put(
		new StreamShardHandle(
			"fakeStream2",
			new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(0))),
		UUID.randomUUID().toString());
	restoredStateUnderTest.put(
		new StreamShardHandle(
			"fakeStream2",
			new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(1))),
		UUID.randomUUID().toString());

	Map<String, Integer> streamToShardCount = new HashMap<>();
	streamToShardCount.put("fakeStream1", 3 + 1); // fakeStream1 had 3 shards before & 1 new shard after restore
	streamToShardCount.put("fakeStream2", 2 + 3); // fakeStream2 had 2 shards before & 2 new shard after restore
	streamToShardCount.put("fakeStream3", 0); // no shards can be found for fakeStream3
	streamToShardCount.put("fakeStream4", 0); // no shards can be found for fakeStream4

	HashMap<String, String> subscribedStreamsToLastSeenShardIdsUnderTest =
		KinesisDataFetcher.createInitialSubscribedStreamsToLastDiscoveredShardsState(fakeStreams);

	// using a non-resharded streams kinesis behaviour to represent that Kinesis is not resharded AFTER the restore
	final TestableKinesisDataFetcher<String> fetcher =
		new TestableKinesisDataFetcher<>(
			fakeStreams,
			new TestSourceContext<>(),
			TestUtils.getStandardProperties(),
			new KinesisDeserializationSchemaWrapper<>(new SimpleStringSchema()),
			10,
			2,
			new AtomicReference<>(),
			new LinkedList<>(),
			subscribedStreamsToLastSeenShardIdsUnderTest,
			FakeKinesisBehavioursFactory.nonReshardedStreamsBehaviour(streamToShardCount));

	for (Map.Entry<StreamShardHandle, String> restoredState : restoredStateUnderTest.entrySet()) {
		fetcher.advanceLastDiscoveredShardOfStream(restoredState.getKey().getStreamName(), restoredState.getKey().getShard().getShardId());
		fetcher.registerNewSubscribedShardState(
			new KinesisStreamShardState(KinesisDataFetcher.convertToStreamShardMetadata(restoredState.getKey()),
				restoredState.getKey(), new SequenceNumber(restoredState.getValue())));
	}

	CheckedThread runFetcherThread = new CheckedThread() {
		@Override
		public void go() throws Exception {
			fetcher.runFetcher();
		}
	};
	runFetcherThread.start();

	fetcher.waitUntilInitialDiscovery();
	fetcher.shutdownFetcher();
	runFetcherThread.sync();

	// assert that the streams tracked in the state are identical to the subscribed streams
	Set<String> streamsInState = subscribedStreamsToLastSeenShardIdsUnderTest.keySet();
	assertEquals(fakeStreams.size(), streamsInState.size());
	assertTrue(streamsInState.containsAll(fakeStreams));

	// assert that the last seen shards in state is correctly set
	assertEquals(
		KinesisShardIdGenerator.generateFromShardOrder(3),
		subscribedStreamsToLastSeenShardIdsUnderTest.get("fakeStream1"));
	assertEquals(
		KinesisShardIdGenerator.generateFromShardOrder(4),
		subscribedStreamsToLastSeenShardIdsUnderTest.get("fakeStream2"));
	assertNull(subscribedStreamsToLastSeenShardIdsUnderTest.get("fakeStream3"));
	assertNull(subscribedStreamsToLastSeenShardIdsUnderTest.get("fakeStream4"));
}
 

@Test
public void testPeriodicWatermark() {
	final MutableLong clock = new MutableLong();
	final MutableBoolean isTemporaryIdle = new MutableBoolean();
	final List<Watermark> watermarks = new ArrayList<>();

	String fakeStream1 = "fakeStream1";
	StreamShardHandle shardHandle =
		new StreamShardHandle(
			fakeStream1,
			new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(0)));

	TestSourceContext<String> sourceContext =
		new TestSourceContext<String>() {
			@Override
			public void emitWatermark(Watermark mark) {
				watermarks.add(mark);
			}

			@Override
			public void markAsTemporarilyIdle() {
				isTemporaryIdle.setTrue();
			}
		};

	HashMap<String, String> subscribedStreamsToLastSeenShardIdsUnderTest = new HashMap<>();

	final KinesisDataFetcher<String> fetcher =
		new TestableKinesisDataFetcher<String>(
			Collections.singletonList(fakeStream1),
			sourceContext,
			new java.util.Properties(),
			new KinesisDeserializationSchemaWrapper<>(new org.apache.flink.streaming.util.serialization.SimpleStringSchema()),
			1,
			1,
			new AtomicReference<>(),
			new LinkedList<>(),
			subscribedStreamsToLastSeenShardIdsUnderTest,
			FakeKinesisBehavioursFactory.nonReshardedStreamsBehaviour(new HashMap<>())) {

			@Override
			protected long getCurrentTimeMillis() {
				return clock.getValue();
			}
		};
	Whitebox.setInternalState(fetcher, "periodicWatermarkAssigner", watermarkAssigner);

	SequenceNumber seq = new SequenceNumber("fakeSequenceNumber");
	// register shards to subsequently emit records
	int shardIndex =
		fetcher.registerNewSubscribedShardState(
			new KinesisStreamShardState(
				KinesisDataFetcher.convertToStreamShardMetadata(shardHandle), shardHandle, seq));

	StreamRecord<String> record1 =
		new StreamRecord<>(String.valueOf(Long.MIN_VALUE), Long.MIN_VALUE);
	fetcher.emitRecordAndUpdateState(record1.getValue(), record1.getTimestamp(), shardIndex, seq);
	Assert.assertEquals(record1, sourceContext.getCollectedOutputs().poll());

	fetcher.emitWatermark();
	Assert.assertTrue("potential watermark equals previous watermark", watermarks.isEmpty());

	StreamRecord<String> record2 = new StreamRecord<>(String.valueOf(1), 1);
	fetcher.emitRecordAndUpdateState(record2.getValue(), record2.getTimestamp(), shardIndex, seq);
	Assert.assertEquals(record2, sourceContext.getCollectedOutputs().poll());

	fetcher.emitWatermark();
	Assert.assertFalse("watermark advanced", watermarks.isEmpty());
	Assert.assertEquals(new Watermark(record2.getTimestamp()), watermarks.remove(0));
	Assert.assertFalse("not idle", isTemporaryIdle.booleanValue());

	// test idle timeout
	long idleTimeout = 10;
	// advance clock idleTimeout
	clock.add(idleTimeout + 1);
	fetcher.emitWatermark();
	Assert.assertFalse("not idle", isTemporaryIdle.booleanValue());
	Assert.assertTrue("not idle, no new watermark", watermarks.isEmpty());

	// activate idle timeout
	Whitebox.setInternalState(fetcher, "shardIdleIntervalMillis", idleTimeout);
	fetcher.emitWatermark();
	Assert.assertTrue("idle", isTemporaryIdle.booleanValue());
	Assert.assertTrue("idle, no watermark", watermarks.isEmpty());
}
 

@Test
public void testStreamToLastSeenShardStateIsCorrectlySetWhenNoNewShardsSinceRestoredCheckpointAndSomeStreamsDoNotExist() throws Exception {
	List<String> fakeStreams = new LinkedList<>();
	fakeStreams.add("fakeStream1");
	fakeStreams.add("fakeStream2");
	fakeStreams.add("fakeStream3"); // fakeStream3 will not have any shards
	fakeStreams.add("fakeStream4"); // fakeStream4 will not have any shards

	Map<StreamShardHandle, String> restoredStateUnderTest = new HashMap<>();

	// fakeStream1 has 3 shards before restore
	restoredStateUnderTest.put(
		new StreamShardHandle(
			"fakeStream1",
			new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(0))),
		UUID.randomUUID().toString());
	restoredStateUnderTest.put(
		new StreamShardHandle(
			"fakeStream1",
			new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(1))),
		UUID.randomUUID().toString());
	restoredStateUnderTest.put(
		new StreamShardHandle(
			"fakeStream1",
			new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(2))),
		UUID.randomUUID().toString());

	// fakeStream2 has 2 shards before restore
	restoredStateUnderTest.put(
		new StreamShardHandle(
			"fakeStream2",
			new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(0))),
		UUID.randomUUID().toString());
	restoredStateUnderTest.put(
		new StreamShardHandle(
			"fakeStream2",
			new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(1))),
		UUID.randomUUID().toString());

	Map<String, Integer> streamToShardCount = new HashMap<>();
	streamToShardCount.put("fakeStream1", 3); // fakeStream1 has fixed 3 shards
	streamToShardCount.put("fakeStream2", 2); // fakeStream2 has fixed 2 shards
	streamToShardCount.put("fakeStream3", 0); // no shards can be found for fakeStream3
	streamToShardCount.put("fakeStream4", 0); // no shards can be found for fakeStream4

	HashMap<String, String> subscribedStreamsToLastSeenShardIdsUnderTest =
		KinesisDataFetcher.createInitialSubscribedStreamsToLastDiscoveredShardsState(fakeStreams);

	// using a non-resharded streams kinesis behaviour to represent that Kinesis is not resharded AFTER the restore
	final TestableKinesisDataFetcher<String> fetcher =
		new TestableKinesisDataFetcher<>(
			fakeStreams,
			new TestSourceContext<>(),
			TestUtils.getStandardProperties(),
			new KinesisDeserializationSchemaWrapper<>(new SimpleStringSchema()),
			10,
			2,
			new AtomicReference<>(),
			new LinkedList<>(),
			subscribedStreamsToLastSeenShardIdsUnderTest,
			FakeKinesisBehavioursFactory.nonReshardedStreamsBehaviour(streamToShardCount));

	for (Map.Entry<StreamShardHandle, String> restoredState : restoredStateUnderTest.entrySet()) {
		fetcher.advanceLastDiscoveredShardOfStream(restoredState.getKey().getStreamName(), restoredState.getKey().getShard().getShardId());
		fetcher.registerNewSubscribedShardState(
			new KinesisStreamShardState(KinesisDataFetcher.convertToStreamShardMetadata(restoredState.getKey()),
				restoredState.getKey(), new SequenceNumber(restoredState.getValue())));
	}

	CheckedThread runFetcherThread = new CheckedThread() {
		@Override
		public void go() throws Exception {
			fetcher.runFetcher();
		}
	};
	runFetcherThread.start();

	fetcher.waitUntilInitialDiscovery();
	fetcher.shutdownFetcher();
	runFetcherThread.sync();

	// assert that the streams tracked in the state are identical to the subscribed streams
	Set<String> streamsInState = subscribedStreamsToLastSeenShardIdsUnderTest.keySet();
	assertEquals(fakeStreams.size(), streamsInState.size());
	assertTrue(streamsInState.containsAll(fakeStreams));

	// assert that the last seen shards in state is correctly set
	assertEquals(
		KinesisShardIdGenerator.generateFromShardOrder(2),
		subscribedStreamsToLastSeenShardIdsUnderTest.get("fakeStream1"));
	assertEquals(
		KinesisShardIdGenerator.generateFromShardOrder(1),
		subscribedStreamsToLastSeenShardIdsUnderTest.get("fakeStream2"));
	assertNull(subscribedStreamsToLastSeenShardIdsUnderTest.get("fakeStream3"));
	assertNull(subscribedStreamsToLastSeenShardIdsUnderTest.get("fakeStream4"));
}
 

@Test
public void testStreamToLastSeenShardStateIsCorrectlySetWhenNewShardsFoundSinceRestoredCheckpoint() throws Exception {
	List<String> fakeStreams = new LinkedList<>();
	fakeStreams.add("fakeStream1");
	fakeStreams.add("fakeStream2");

	Map<StreamShardHandle, String> restoredStateUnderTest = new HashMap<>();

	// fakeStream1 has 3 shards before restore
	restoredStateUnderTest.put(
		new StreamShardHandle(
			"fakeStream1",
			new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(0))),
		UUID.randomUUID().toString());
	restoredStateUnderTest.put(
		new StreamShardHandle(
			"fakeStream1",
			new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(1))),
		UUID.randomUUID().toString());
	restoredStateUnderTest.put(
		new StreamShardHandle(
			"fakeStream1",
			new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(2))),
		UUID.randomUUID().toString());

	// fakeStream2 has 2 shards before restore
	restoredStateUnderTest.put(
		new StreamShardHandle(
			"fakeStream2",
			new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(0))),
		UUID.randomUUID().toString());
	restoredStateUnderTest.put(
		new StreamShardHandle(
			"fakeStream2",
			new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(1))),
		UUID.randomUUID().toString());

	Map<String, Integer> streamToShardCount = new HashMap<>();
	streamToShardCount.put("fakeStream1", 3 + 1); // fakeStream1 had 3 shards before & 1 new shard after restore
	streamToShardCount.put("fakeStream2", 2 + 3); // fakeStream2 had 2 shards before & 3 new shard after restore

	HashMap<String, String> subscribedStreamsToLastSeenShardIdsUnderTest =
		KinesisDataFetcher.createInitialSubscribedStreamsToLastDiscoveredShardsState(fakeStreams);

	// using a non-resharded streams kinesis behaviour to represent that Kinesis is not resharded AFTER the restore
	final TestableKinesisDataFetcher<String> fetcher =
		new TestableKinesisDataFetcher<>(
			fakeStreams,
			new TestSourceContext<>(),
			TestUtils.getStandardProperties(),
			new KinesisDeserializationSchemaWrapper<>(new SimpleStringSchema()),
			10,
			2,
			new AtomicReference<>(),
			new LinkedList<>(),
			subscribedStreamsToLastSeenShardIdsUnderTest,
			FakeKinesisBehavioursFactory.nonReshardedStreamsBehaviour(streamToShardCount));

	for (Map.Entry<StreamShardHandle, String> restoredState : restoredStateUnderTest.entrySet()) {
		fetcher.advanceLastDiscoveredShardOfStream(restoredState.getKey().getStreamName(), restoredState.getKey().getShard().getShardId());
		fetcher.registerNewSubscribedShardState(
			new KinesisStreamShardState(KinesisDataFetcher.convertToStreamShardMetadata(restoredState.getKey()),
				restoredState.getKey(), new SequenceNumber(restoredState.getValue())));
	}

	CheckedThread runFetcherThread = new CheckedThread() {
		@Override
		public void go() throws Exception {
			fetcher.runFetcher();
		}
	};
	runFetcherThread.start();

	fetcher.waitUntilInitialDiscovery();
	fetcher.shutdownFetcher();
	runFetcherThread.sync();

	// assert that the streams tracked in the state are identical to the subscribed streams
	Set<String> streamsInState = subscribedStreamsToLastSeenShardIdsUnderTest.keySet();
	assertEquals(fakeStreams.size(), streamsInState.size());
	assertTrue(streamsInState.containsAll(fakeStreams));

	// assert that the last seen shards in state is correctly set
	for (Map.Entry<String, String> streamToLastSeenShard : subscribedStreamsToLastSeenShardIdsUnderTest.entrySet()) {
		assertEquals(
			KinesisShardIdGenerator.generateFromShardOrder(streamToShardCount.get(streamToLastSeenShard.getKey()) - 1),
			streamToLastSeenShard.getValue());
	}
}
 

@Test
public void testRestoreWithReshardedStream() throws Exception {
	final List<StreamShardHandle> initialDiscoveryShards = new ArrayList<>(TEST_STATE.size());
	for (StreamShardMetadata shardMetadata : TEST_STATE.keySet()) {
		// setup the closed shard
		Shard closedShard = new Shard();
		closedShard.setShardId(shardMetadata.getShardId());

		SequenceNumberRange closedSequenceNumberRange = new SequenceNumberRange();
		closedSequenceNumberRange.withStartingSequenceNumber("1");
		closedSequenceNumberRange.withEndingSequenceNumber("1087654321"); // this represents a closed shard
		closedShard.setSequenceNumberRange(closedSequenceNumberRange);

		initialDiscoveryShards.add(new StreamShardHandle(shardMetadata.getStreamName(), closedShard));

		// setup the new shards
		Shard newSplitShard1 = new Shard();
		newSplitShard1.setShardId(KinesisShardIdGenerator.generateFromShardOrder(1));

		SequenceNumberRange newSequenceNumberRange1 = new SequenceNumberRange();
		newSequenceNumberRange1.withStartingSequenceNumber("1087654322");
		newSplitShard1.setSequenceNumberRange(newSequenceNumberRange1);

		newSplitShard1.setParentShardId(TEST_SHARD_ID);

		Shard newSplitShard2 = new Shard();
		newSplitShard2.setShardId(KinesisShardIdGenerator.generateFromShardOrder(2));

		SequenceNumberRange newSequenceNumberRange2 = new SequenceNumberRange();
		newSequenceNumberRange2.withStartingSequenceNumber("2087654322");
		newSplitShard2.setSequenceNumberRange(newSequenceNumberRange2);

		newSplitShard2.setParentShardId(TEST_SHARD_ID);

		initialDiscoveryShards.add(new StreamShardHandle(shardMetadata.getStreamName(), newSplitShard1));
		initialDiscoveryShards.add(new StreamShardHandle(shardMetadata.getStreamName(), newSplitShard2));
	}

	final TestFetcher<String> fetcher = new TestFetcher<>(
		Collections.singletonList(TEST_STREAM_NAME),
		new TestSourceContext<>(),
		new TestRuntimeContext(true, 1, 0),
		TestUtils.getStandardProperties(),
		new KinesisDeserializationSchemaWrapper<>(new SimpleStringSchema()),
		null,
		initialDiscoveryShards);

	final DummyFlinkKinesisConsumer<String> consumerFunction = new DummyFlinkKinesisConsumer<>(
		fetcher, new KinesisDeserializationSchemaWrapper<>(new SimpleStringSchema()));

	StreamSource<String, DummyFlinkKinesisConsumer<String>> consumerOperator =
		new StreamSource<>(consumerFunction);

	final AbstractStreamOperatorTestHarness<String> testHarness =
		new AbstractStreamOperatorTestHarness<>(consumerOperator, 1, 1, 0);

	testHarness.setup();
	testHarness.initializeState(
		OperatorSnapshotUtil.getResourceFilename(
			"kinesis-consumer-migration-test-flink" + testMigrateVersion + "-snapshot"));
	testHarness.open();

	consumerFunction.run(new TestSourceContext<>());

	// assert that state is correctly restored
	assertNotEquals(null, consumerFunction.getRestoredState());
	assertEquals(1, consumerFunction.getRestoredState().size());
	assertEquals(TEST_STATE, removeEquivalenceWrappers(consumerFunction.getRestoredState()));

	// assert that the fetcher is registered with all shards, including new shards
	assertEquals(3, fetcher.getSubscribedShardsState().size());

	KinesisStreamShardState restoredClosedShardState = fetcher.getSubscribedShardsState().get(0);
	assertEquals(TEST_STREAM_NAME, restoredClosedShardState.getStreamShardHandle().getStreamName());
	assertEquals(TEST_SHARD_ID, restoredClosedShardState.getStreamShardHandle().getShard().getShardId());
	assertTrue(restoredClosedShardState.getStreamShardHandle().isClosed());
	assertEquals(TEST_SEQUENCE_NUMBER, restoredClosedShardState.getLastProcessedSequenceNum());

	KinesisStreamShardState restoredNewSplitShard1 = fetcher.getSubscribedShardsState().get(1);
	assertEquals(TEST_STREAM_NAME, restoredNewSplitShard1.getStreamShardHandle().getStreamName());
	assertEquals(KinesisShardIdGenerator.generateFromShardOrder(1), restoredNewSplitShard1.getStreamShardHandle().getShard().getShardId());
	assertFalse(restoredNewSplitShard1.getStreamShardHandle().isClosed());
	// new shards should be consumed from the beginning
	assertEquals(SentinelSequenceNumber.SENTINEL_EARLIEST_SEQUENCE_NUM.get(), restoredNewSplitShard1.getLastProcessedSequenceNum());

	KinesisStreamShardState restoredNewSplitShard2 = fetcher.getSubscribedShardsState().get(2);
	assertEquals(TEST_STREAM_NAME, restoredNewSplitShard2.getStreamShardHandle().getStreamName());
	assertEquals(KinesisShardIdGenerator.generateFromShardOrder(2), restoredNewSplitShard2.getStreamShardHandle().getShard().getShardId());
	assertFalse(restoredNewSplitShard2.getStreamShardHandle().isClosed());
	// new shards should be consumed from the beginning
	assertEquals(SentinelSequenceNumber.SENTINEL_EARLIEST_SEQUENCE_NUM.get(), restoredNewSplitShard2.getLastProcessedSequenceNum());

	consumerOperator.close();
	consumerOperator.cancel();
}
 

@Test
public void testGetShardList() throws Exception {
	List<String> shardIds =
			Arrays.asList(
					"shardId-000000000000",
					"shardId-000000000001",
					"shardId-000000000002",
					"shardId-000000000003");
	String nextToken = "NextToken";
	String fakeStreamName = "fake-stream";
	List<Shard> shards = shardIds
					.stream()
					.map(shardId -> new Shard().withShardId(shardId))
					.collect(Collectors.toList());
	AmazonKinesis mockClient = mock(AmazonKinesis.class);
	KinesisProxy kinesisProxy = getProxy(mockClient);

	ListShardsResult responseWithMoreData =
			new ListShardsResult().withShards(shards.subList(0, 2)).withNextToken(nextToken);
	ListShardsResult responseFinal =
			new ListShardsResult().withShards(shards.subList(2, shards.size())).withNextToken(null);
	doReturn(responseWithMoreData)
			.when(mockClient)
			.listShards(argThat(initialListShardsRequestMatcher()));
	doReturn(responseFinal).
					when(mockClient).
					listShards(argThat(listShardsNextToken(nextToken)));
	HashMap<String, String> streamHashMap =
			createInitialSubscribedStreamsToLastDiscoveredShardsState(Arrays.asList(fakeStreamName));
	GetShardListResult shardListResult = kinesisProxy.getShardList(streamHashMap);

	Assert.assertEquals(shardListResult.hasRetrievedShards(), true);

	Set<String> expectedStreams = new HashSet<>();
	expectedStreams.add(fakeStreamName);
	Assert.assertEquals(shardListResult.getStreamsWithRetrievedShards(), expectedStreams);
	List<StreamShardHandle> actualShardList =
			shardListResult.getRetrievedShardListOfStream(fakeStreamName);
	List<StreamShardHandle> expectedStreamShard = new ArrayList<>();
	assertThat(actualShardList, hasSize(4));
	for (int i = 0; i < 4; i++) {
		StreamShardHandle shardHandle =
				new StreamShardHandle(
						fakeStreamName,
						new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(i)));
		expectedStreamShard.add(shardHandle);
	}

	Assert.assertThat(
			actualShardList,
			containsInAnyOrder(
					expectedStreamShard.toArray(new StreamShardHandle[actualShardList.size()])));
}
 

@Test
public void testStreamToLastSeenShardStateIsCorrectlySetWhenNoNewShardsSinceRestoredCheckpoint() throws Exception {
	List<String> fakeStreams = new LinkedList<>();
	fakeStreams.add("fakeStream1");
	fakeStreams.add("fakeStream2");

	Map<StreamShardHandle, String> restoredStateUnderTest = new HashMap<>();

	// fakeStream1 has 3 shards before restore
	restoredStateUnderTest.put(
		new StreamShardHandle(
			"fakeStream1",
			new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(0))),
		UUID.randomUUID().toString());
	restoredStateUnderTest.put(
		new StreamShardHandle(
			"fakeStream1",
			new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(1))),
		UUID.randomUUID().toString());
	restoredStateUnderTest.put(
		new StreamShardHandle(
			"fakeStream1",
			new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(2))),
		UUID.randomUUID().toString());

	// fakeStream2 has 2 shards before restore
	restoredStateUnderTest.put(
		new StreamShardHandle(
			"fakeStream2",
			new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(0))),
		UUID.randomUUID().toString());
	restoredStateUnderTest.put(
		new StreamShardHandle(
			"fakeStream2",
			new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(1))),
		UUID.randomUUID().toString());

	Map<String, Integer> streamToShardCount = new HashMap<>();
	streamToShardCount.put("fakeStream1", 3); // fakeStream1 will still have 3 shards after restore
	streamToShardCount.put("fakeStream2", 2); // fakeStream2 will still have 2 shards after restore

	HashMap<String, String> subscribedStreamsToLastSeenShardIdsUnderTest =
		KinesisDataFetcher.createInitialSubscribedStreamsToLastDiscoveredShardsState(fakeStreams);

	final TestableKinesisDataFetcher<String> fetcher =
		new TestableKinesisDataFetcher<>(
			fakeStreams,
			new TestSourceContext<>(),
			TestUtils.getStandardProperties(),
			new KinesisDeserializationSchemaWrapper<>(new SimpleStringSchema()),
			10,
			2,
			new AtomicReference<>(),
			new LinkedList<>(),
			subscribedStreamsToLastSeenShardIdsUnderTest,
			FakeKinesisBehavioursFactory.nonReshardedStreamsBehaviour(streamToShardCount));

	for (Map.Entry<StreamShardHandle, String> restoredState : restoredStateUnderTest.entrySet()) {
		fetcher.advanceLastDiscoveredShardOfStream(restoredState.getKey().getStreamName(), restoredState.getKey().getShard().getShardId());
		fetcher.registerNewSubscribedShardState(
			new KinesisStreamShardState(KinesisDataFetcher.convertToStreamShardMetadata(restoredState.getKey()),
				restoredState.getKey(), new SequenceNumber(restoredState.getValue())));
	}

	CheckedThread runFetcherThread = new CheckedThread() {
		@Override
		public void go() throws Exception {
			fetcher.runFetcher();
		}
	};
	runFetcherThread.start();

	fetcher.waitUntilInitialDiscovery();
	fetcher.shutdownFetcher();
	runFetcherThread.sync();

	// assert that the streams tracked in the state are identical to the subscribed streams
	Set<String> streamsInState = subscribedStreamsToLastSeenShardIdsUnderTest.keySet();
	assertEquals(fakeStreams.size(), streamsInState.size());
	assertTrue(streamsInState.containsAll(fakeStreams));

	// assert that the last seen shards in state is correctly set
	for (Map.Entry<String, String> streamToLastSeenShard : subscribedStreamsToLastSeenShardIdsUnderTest.entrySet()) {
		assertEquals(
			KinesisShardIdGenerator.generateFromShardOrder(streamToShardCount.get(streamToLastSeenShard.getKey()) - 1),
			streamToLastSeenShard.getValue());
	}
}
 

@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<>(),
			TestUtils.getStandardProperties(),
			new KinesisDeserializationSchemaWrapper<>(new SimpleStringSchema()),
			10,
			2,
			new AtomicReference<>(),
			new LinkedList<>(),
			subscribedStreamsToLastSeenShardIdsUnderTest,
			FakeKinesisBehavioursFactory.nonReshardedStreamsBehaviour(streamToShardCount));

	final DummyFlinkKinesisConsumer<String> consumer = new DummyFlinkKinesisConsumer<>(
			TestUtils.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();

	// assert that the streams tracked in the state are identical to the subscribed streams
	Set<String> streamsInState = subscribedStreamsToLastSeenShardIdsUnderTest.keySet();
	assertEquals(fakeStreams.size(), streamsInState.size());
	assertTrue(streamsInState.containsAll(fakeStreams));

	// assert that the last seen shards in state is correctly set
	for (Map.Entry<String, String> streamToLastSeenShard : subscribedStreamsToLastSeenShardIdsUnderTest.entrySet()) {
		assertEquals(
			KinesisShardIdGenerator.generateFromShardOrder(streamToShardCount.get(streamToLastSeenShard.getKey()) - 1),
			streamToLastSeenShard.getValue());
	}
}
 

@Test
public void testGetShardListWithNewShardsOnSecondRun() throws Exception {
	// given
	List<String> shardIds =
		Arrays.asList(
			KinesisShardIdGenerator.generateFromShardOrder(0),
			KinesisShardIdGenerator.generateFromShardOrder(1)
		);
	String fakeStreamName = "fake-stream";
	List<Shard> shards = shardIds
		.stream()
		.map(shardId -> new Shard().withShardId(shardId))
		.collect(Collectors.toList());

	AmazonKinesis mockClient = mock(AmazonKinesis.class);
	KinesisProxy kinesisProxy = getProxy(mockClient);

	ListShardsResult responseFirst =
		new ListShardsResult().withShards(shards).withNextToken(null);
	doReturn(responseFirst)
		.when(mockClient)
		.listShards(argThat(initialListShardsRequestMatcher()));
	HashMap<String, String> streamHashMap =
		createInitialSubscribedStreamsToLastDiscoveredShardsState(Collections.singletonList(fakeStreamName));

	// when
	GetShardListResult shardListResult = kinesisProxy.getShardList(streamHashMap);

	// then
	Assert.assertTrue(shardListResult.hasRetrievedShards());

	Set<String> expectedStreams = new HashSet<>();
	expectedStreams.add(fakeStreamName);
	Assert.assertEquals(shardListResult.getStreamsWithRetrievedShards(), expectedStreams);

	List<StreamShardHandle> actualShardList =
		shardListResult.getRetrievedShardListOfStream(fakeStreamName);
	Assert.assertThat(actualShardList, hasSize(2));

	List<StreamShardHandle> expectedStreamShard = IntStream.range(0, actualShardList.size())
		.mapToObj(i ->
			new StreamShardHandle(
				fakeStreamName,
				new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(i)))
		).collect(Collectors.toList());

	Assert.assertThat(actualShardList, containsInAnyOrder(
		expectedStreamShard.toArray(new StreamShardHandle[actualShardList.size()])));

	// given new shards
	ListShardsResult responseSecond =
		new ListShardsResult().withShards(new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(2)))
			.withNextToken(null);
	doReturn(responseSecond)
		.when(mockClient)
		.listShards(argThat(initialListShardsRequestMatcher()));

	// when new shards
	GetShardListResult newShardListResult = kinesisProxy.getShardList(streamHashMap);

	// then new shards
	Assert.assertTrue(newShardListResult.hasRetrievedShards());
	Assert.assertEquals(newShardListResult.getStreamsWithRetrievedShards(), expectedStreams);

	List<StreamShardHandle> newActualShardList =
		newShardListResult.getRetrievedShardListOfStream(fakeStreamName);
	Assert.assertThat(newActualShardList, hasSize(1));

	List<StreamShardHandle> newExpectedStreamShard = Collections.singletonList(
		new StreamShardHandle(
			fakeStreamName,
			new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(2)))
	);

	Assert.assertThat(newActualShardList, containsInAnyOrder(
		newExpectedStreamShard.toArray(new StreamShardHandle[newActualShardList.size()])));
}
 

@Test
public void testUseRestoredStateForSnapshotIfFetcherNotInitialized() throws Exception {
	Properties config = TestUtils.getStandardProperties();

	List<Tuple2<StreamShardMetadata, SequenceNumber>> globalUnionState = new ArrayList<>(4);
	globalUnionState.add(Tuple2.of(
		KinesisDataFetcher.convertToStreamShardMetadata(new StreamShardHandle("fakeStream",
			new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(0)))),
		new SequenceNumber("1")));
	globalUnionState.add(Tuple2.of(
		KinesisDataFetcher.convertToStreamShardMetadata(new StreamShardHandle("fakeStream",
			new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(1)))),
		new SequenceNumber("1")));
	globalUnionState.add(Tuple2.of(
		KinesisDataFetcher.convertToStreamShardMetadata(new StreamShardHandle("fakeStream",
			new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(2)))),
		new SequenceNumber("1")));
	globalUnionState.add(Tuple2.of(
		KinesisDataFetcher.convertToStreamShardMetadata(new StreamShardHandle("fakeStream",
			new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(3)))),
		new SequenceNumber("1")));

	TestingListState<Tuple2<StreamShardMetadata, SequenceNumber>> listState = new TestingListState<>();
	for (Tuple2<StreamShardMetadata, SequenceNumber> state : globalUnionState) {
		listState.add(state);
	}

	FlinkKinesisConsumer<String> consumer = new FlinkKinesisConsumer<>("fakeStream", new SimpleStringSchema(), config);
	RuntimeContext context = mock(RuntimeContext.class);
	when(context.getIndexOfThisSubtask()).thenReturn(0);
	when(context.getNumberOfParallelSubtasks()).thenReturn(2);
	consumer.setRuntimeContext(context);

	OperatorStateStore operatorStateStore = mock(OperatorStateStore.class);
	when(operatorStateStore.getUnionListState(Matchers.any(ListStateDescriptor.class))).thenReturn(listState);

	StateInitializationContext initializationContext = mock(StateInitializationContext.class);
	when(initializationContext.getOperatorStateStore()).thenReturn(operatorStateStore);
	when(initializationContext.isRestored()).thenReturn(true);

	consumer.initializeState(initializationContext);

	// only opened, not run
	consumer.open(new Configuration());

	// arbitrary checkpoint id and timestamp
	consumer.snapshotState(new StateSnapshotContextSynchronousImpl(123, 123));

	assertTrue(listState.isClearCalled());

	// the checkpointed list state should contain only the shards that it should subscribe to
	assertEquals(globalUnionState.size() / 2, listState.getList().size());
	assertTrue(listState.getList().contains(globalUnionState.get(0)));
	assertTrue(listState.getList().contains(globalUnionState.get(2)));
}
 

@Test
@SuppressWarnings("unchecked")
public void testFetcherShouldBeCorrectlySeededWithNewDiscoveredKinesisStreamShard() throws Exception {

	// ----------------------------------------------------------------------
	// setup initial state
	// ----------------------------------------------------------------------

	HashMap<StreamShardHandle, SequenceNumber> fakeRestoredState = getFakeRestoredStore("all");

	// ----------------------------------------------------------------------
	// mock operator state backend and initial state for initializeState()
	// ----------------------------------------------------------------------

	TestingListState<Tuple2<StreamShardMetadata, SequenceNumber>> listState = new TestingListState<>();
	for (Map.Entry<StreamShardHandle, SequenceNumber> state : fakeRestoredState.entrySet()) {
		listState.add(Tuple2.of(KinesisDataFetcher.convertToStreamShardMetadata(state.getKey()), state.getValue()));
	}

	OperatorStateStore operatorStateStore = mock(OperatorStateStore.class);
	when(operatorStateStore.getUnionListState(Matchers.any(ListStateDescriptor.class))).thenReturn(listState);

	StateInitializationContext initializationContext = mock(StateInitializationContext.class);
	when(initializationContext.getOperatorStateStore()).thenReturn(operatorStateStore);
	when(initializationContext.isRestored()).thenReturn(true);

	// ----------------------------------------------------------------------
	// mock fetcher
	// ----------------------------------------------------------------------

	KinesisDataFetcher mockedFetcher = mockKinesisDataFetcher();
	List<StreamShardHandle> shards = new ArrayList<>();
	shards.addAll(fakeRestoredState.keySet());
	shards.add(new StreamShardHandle("fakeStream2",
		new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(2))));
	when(mockedFetcher.discoverNewShardsToSubscribe()).thenReturn(shards);

	// assume the given config is correct
	PowerMockito.mockStatic(KinesisConfigUtil.class);
	PowerMockito.doNothing().when(KinesisConfigUtil.class);

	// ----------------------------------------------------------------------
	// start to test fetcher's initial state seeding
	// ----------------------------------------------------------------------

	TestableFlinkKinesisConsumer consumer = new TestableFlinkKinesisConsumer(
		"fakeStream", new Properties(), 10, 2);
	consumer.initializeState(initializationContext);
	consumer.open(new Configuration());
	consumer.run(Mockito.mock(SourceFunction.SourceContext.class));

	fakeRestoredState.put(new StreamShardHandle("fakeStream2",
			new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(2))),
		SentinelSequenceNumber.SENTINEL_EARLIEST_SEQUENCE_NUM.get());
	for (Map.Entry<StreamShardHandle, SequenceNumber> restoredShard : fakeRestoredState.entrySet()) {
		Mockito.verify(mockedFetcher).registerNewSubscribedShardState(
			new KinesisStreamShardState(KinesisDataFetcher.convertToStreamShardMetadata(restoredShard.getKey()),
				restoredShard.getKey(), restoredShard.getValue()));
	}
}
 

@Test
public void testUseRestoredStateForSnapshotIfFetcherNotInitialized() throws Exception {
	Properties config = TestUtils.getStandardProperties();

	List<Tuple2<StreamShardMetadata, SequenceNumber>> globalUnionState = new ArrayList<>(4);
	globalUnionState.add(Tuple2.of(
		KinesisDataFetcher.convertToStreamShardMetadata(new StreamShardHandle("fakeStream",
			new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(0)))),
		new SequenceNumber("1")));
	globalUnionState.add(Tuple2.of(
		KinesisDataFetcher.convertToStreamShardMetadata(new StreamShardHandle("fakeStream",
			new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(1)))),
		new SequenceNumber("1")));
	globalUnionState.add(Tuple2.of(
		KinesisDataFetcher.convertToStreamShardMetadata(new StreamShardHandle("fakeStream",
			new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(2)))),
		new SequenceNumber("1")));
	globalUnionState.add(Tuple2.of(
		KinesisDataFetcher.convertToStreamShardMetadata(new StreamShardHandle("fakeStream",
			new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(3)))),
		new SequenceNumber("1")));

	TestingListState<Tuple2<StreamShardMetadata, SequenceNumber>> listState = new TestingListState<>();
	for (Tuple2<StreamShardMetadata, SequenceNumber> state : globalUnionState) {
		listState.add(state);
	}

	FlinkKinesisConsumer<String> consumer = new FlinkKinesisConsumer<>("fakeStream", new SimpleStringSchema(), config);
	RuntimeContext context = mock(RuntimeContext.class);
	when(context.getIndexOfThisSubtask()).thenReturn(0);
	when(context.getNumberOfParallelSubtasks()).thenReturn(2);
	consumer.setRuntimeContext(context);

	OperatorStateStore operatorStateStore = mock(OperatorStateStore.class);
	when(operatorStateStore.getUnionListState(Matchers.any(ListStateDescriptor.class))).thenReturn(listState);

	StateInitializationContext initializationContext = mock(StateInitializationContext.class);
	when(initializationContext.getOperatorStateStore()).thenReturn(operatorStateStore);
	when(initializationContext.isRestored()).thenReturn(true);

	consumer.initializeState(initializationContext);

	// only opened, not run
	consumer.open(new Configuration());

	// arbitrary checkpoint id and timestamp
	consumer.snapshotState(new StateSnapshotContextSynchronousImpl(123, 123));

	assertTrue(listState.isClearCalled());

	// the checkpointed list state should contain only the shards that it should subscribe to
	assertEquals(globalUnionState.size() / 2, listState.getList().size());
	assertTrue(listState.getList().contains(globalUnionState.get(0)));
	assertTrue(listState.getList().contains(globalUnionState.get(2)));
}
 

@Test
@SuppressWarnings("unchecked")
public void testFetcherShouldBeCorrectlySeededWithNewDiscoveredKinesisStreamShard() throws Exception {

	// ----------------------------------------------------------------------
	// setup initial state
	// ----------------------------------------------------------------------

	HashMap<StreamShardHandle, SequenceNumber> fakeRestoredState = getFakeRestoredStore("all");

	// ----------------------------------------------------------------------
	// mock operator state backend and initial state for initializeState()
	// ----------------------------------------------------------------------

	TestingListState<Tuple2<StreamShardMetadata, SequenceNumber>> listState = new TestingListState<>();
	for (Map.Entry<StreamShardHandle, SequenceNumber> state : fakeRestoredState.entrySet()) {
		listState.add(Tuple2.of(KinesisDataFetcher.convertToStreamShardMetadata(state.getKey()), state.getValue()));
	}

	OperatorStateStore operatorStateStore = mock(OperatorStateStore.class);
	when(operatorStateStore.getUnionListState(Matchers.any(ListStateDescriptor.class))).thenReturn(listState);

	StateInitializationContext initializationContext = mock(StateInitializationContext.class);
	when(initializationContext.getOperatorStateStore()).thenReturn(operatorStateStore);
	when(initializationContext.isRestored()).thenReturn(true);

	// ----------------------------------------------------------------------
	// mock fetcher
	// ----------------------------------------------------------------------

	KinesisDataFetcher mockedFetcher = mockKinesisDataFetcher();
	List<StreamShardHandle> shards = new ArrayList<>();
	shards.addAll(fakeRestoredState.keySet());
	shards.add(new StreamShardHandle("fakeStream2",
		new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(2))));
	when(mockedFetcher.discoverNewShardsToSubscribe()).thenReturn(shards);

	// assume the given config is correct
	PowerMockito.mockStatic(KinesisConfigUtil.class);
	PowerMockito.doNothing().when(KinesisConfigUtil.class);

	// ----------------------------------------------------------------------
	// start to test fetcher's initial state seeding
	// ----------------------------------------------------------------------

	TestableFlinkKinesisConsumer consumer = new TestableFlinkKinesisConsumer(
		"fakeStream", new Properties(), 10, 2);
	consumer.initializeState(initializationContext);
	consumer.open(new Configuration());
	consumer.run(Mockito.mock(SourceFunction.SourceContext.class));

	fakeRestoredState.put(new StreamShardHandle("fakeStream2",
			new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(2))),
		SentinelSequenceNumber.SENTINEL_EARLIEST_SEQUENCE_NUM.get());
	for (Map.Entry<StreamShardHandle, SequenceNumber> restoredShard : fakeRestoredState.entrySet()) {
		Mockito.verify(mockedFetcher).registerNewSubscribedShardState(
			new KinesisStreamShardState(KinesisDataFetcher.convertToStreamShardMetadata(restoredShard.getKey()),
				restoredShard.getKey(), restoredShard.getValue()));
	}
}
 

@Test
public void testRestoreWithReshardedStream() throws Exception {
	final List<StreamShardHandle> initialDiscoveryShards = new ArrayList<>(TEST_STATE.size());
	for (StreamShardMetadata shardMetadata : TEST_STATE.keySet()) {
		// setup the closed shard
		Shard closedShard = new Shard();
		closedShard.setShardId(shardMetadata.getShardId());

		SequenceNumberRange closedSequenceNumberRange = new SequenceNumberRange();
		closedSequenceNumberRange.withStartingSequenceNumber("1");
		closedSequenceNumberRange.withEndingSequenceNumber("1087654321"); // this represents a closed shard
		closedShard.setSequenceNumberRange(closedSequenceNumberRange);

		initialDiscoveryShards.add(new StreamShardHandle(shardMetadata.getStreamName(), closedShard));

		// setup the new shards
		Shard newSplitShard1 = new Shard();
		newSplitShard1.setShardId(KinesisShardIdGenerator.generateFromShardOrder(1));

		SequenceNumberRange newSequenceNumberRange1 = new SequenceNumberRange();
		newSequenceNumberRange1.withStartingSequenceNumber("1087654322");
		newSplitShard1.setSequenceNumberRange(newSequenceNumberRange1);

		newSplitShard1.setParentShardId(TEST_SHARD_ID);

		Shard newSplitShard2 = new Shard();
		newSplitShard2.setShardId(KinesisShardIdGenerator.generateFromShardOrder(2));

		SequenceNumberRange newSequenceNumberRange2 = new SequenceNumberRange();
		newSequenceNumberRange2.withStartingSequenceNumber("2087654322");
		newSplitShard2.setSequenceNumberRange(newSequenceNumberRange2);

		newSplitShard2.setParentShardId(TEST_SHARD_ID);

		initialDiscoveryShards.add(new StreamShardHandle(shardMetadata.getStreamName(), newSplitShard1));
		initialDiscoveryShards.add(new StreamShardHandle(shardMetadata.getStreamName(), newSplitShard2));
	}

	final TestFetcher<String> fetcher = new TestFetcher<>(
		Collections.singletonList(TEST_STREAM_NAME),
		new TestSourceContext<>(),
		new TestRuntimeContext(true, 1, 0),
		TestUtils.getStandardProperties(),
		new KinesisDeserializationSchemaWrapper<>(new SimpleStringSchema()),
		null,
		initialDiscoveryShards);

	final DummyFlinkKinesisConsumer<String> consumerFunction = new DummyFlinkKinesisConsumer<>(
		fetcher, new KinesisDeserializationSchemaWrapper<>(new SimpleStringSchema()));

	StreamSource<String, DummyFlinkKinesisConsumer<String>> consumerOperator =
		new StreamSource<>(consumerFunction);

	final AbstractStreamOperatorTestHarness<String> testHarness =
		new AbstractStreamOperatorTestHarness<>(consumerOperator, 1, 1, 0);

	testHarness.setup();
	testHarness.initializeState(
		OperatorSnapshotUtil.getResourceFilename(
			"kinesis-consumer-migration-test-flink" + testMigrateVersion + "-snapshot"));
	testHarness.open();

	consumerFunction.run(new TestSourceContext<>());

	// assert that state is correctly restored
	assertNotEquals(null, consumerFunction.getRestoredState());
	assertEquals(1, consumerFunction.getRestoredState().size());
	assertEquals(TEST_STATE, removeEquivalenceWrappers(consumerFunction.getRestoredState()));

	// assert that the fetcher is registered with all shards, including new shards
	assertEquals(3, fetcher.getSubscribedShardsState().size());

	KinesisStreamShardState restoredClosedShardState = fetcher.getSubscribedShardsState().get(0);
	assertEquals(TEST_STREAM_NAME, restoredClosedShardState.getStreamShardHandle().getStreamName());
	assertEquals(TEST_SHARD_ID, restoredClosedShardState.getStreamShardHandle().getShard().getShardId());
	assertTrue(restoredClosedShardState.getStreamShardHandle().isClosed());
	assertEquals(TEST_SEQUENCE_NUMBER, restoredClosedShardState.getLastProcessedSequenceNum());

	KinesisStreamShardState restoredNewSplitShard1 = fetcher.getSubscribedShardsState().get(1);
	assertEquals(TEST_STREAM_NAME, restoredNewSplitShard1.getStreamShardHandle().getStreamName());
	assertEquals(KinesisShardIdGenerator.generateFromShardOrder(1), restoredNewSplitShard1.getStreamShardHandle().getShard().getShardId());
	assertFalse(restoredNewSplitShard1.getStreamShardHandle().isClosed());
	// new shards should be consumed from the beginning
	assertEquals(SentinelSequenceNumber.SENTINEL_EARLIEST_SEQUENCE_NUM.get(), restoredNewSplitShard1.getLastProcessedSequenceNum());

	KinesisStreamShardState restoredNewSplitShard2 = fetcher.getSubscribedShardsState().get(2);
	assertEquals(TEST_STREAM_NAME, restoredNewSplitShard2.getStreamShardHandle().getStreamName());
	assertEquals(KinesisShardIdGenerator.generateFromShardOrder(2), restoredNewSplitShard2.getStreamShardHandle().getShard().getShardId());
	assertFalse(restoredNewSplitShard2.getStreamShardHandle().isClosed());
	// new shards should be consumed from the beginning
	assertEquals(SentinelSequenceNumber.SENTINEL_EARLIEST_SEQUENCE_NUM.get(), restoredNewSplitShard2.getLastProcessedSequenceNum());

	consumerOperator.close();
	consumerOperator.cancel();
}
 

@Test
public void testUseRestoredStateForSnapshotIfFetcherNotInitialized() throws Exception {
	Properties config = TestUtils.getStandardProperties();

	List<Tuple2<StreamShardMetadata, SequenceNumber>> globalUnionState = new ArrayList<>(4);
	globalUnionState.add(Tuple2.of(
		KinesisDataFetcher.convertToStreamShardMetadata(new StreamShardHandle("fakeStream",
			new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(0)))),
		new SequenceNumber("1")));
	globalUnionState.add(Tuple2.of(
		KinesisDataFetcher.convertToStreamShardMetadata(new StreamShardHandle("fakeStream",
			new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(1)))),
		new SequenceNumber("1")));
	globalUnionState.add(Tuple2.of(
		KinesisDataFetcher.convertToStreamShardMetadata(new StreamShardHandle("fakeStream",
			new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(2)))),
		new SequenceNumber("1")));
	globalUnionState.add(Tuple2.of(
		KinesisDataFetcher.convertToStreamShardMetadata(new StreamShardHandle("fakeStream",
			new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(3)))),
		new SequenceNumber("1")));

	TestingListState<Tuple2<StreamShardMetadata, SequenceNumber>> listState = new TestingListState<>();
	for (Tuple2<StreamShardMetadata, SequenceNumber> state : globalUnionState) {
		listState.add(state);
	}

	FlinkKinesisConsumer<String> consumer = new FlinkKinesisConsumer<>("fakeStream", new SimpleStringSchema(), config);
	RuntimeContext context = mock(RuntimeContext.class);
	when(context.getIndexOfThisSubtask()).thenReturn(0);
	when(context.getNumberOfParallelSubtasks()).thenReturn(2);
	consumer.setRuntimeContext(context);

	OperatorStateStore operatorStateStore = mock(OperatorStateStore.class);
	when(operatorStateStore.getUnionListState(Matchers.any(ListStateDescriptor.class))).thenReturn(listState);

	StateInitializationContext initializationContext = mock(StateInitializationContext.class);
	when(initializationContext.getOperatorStateStore()).thenReturn(operatorStateStore);
	when(initializationContext.isRestored()).thenReturn(true);

	consumer.initializeState(initializationContext);

	// only opened, not run
	consumer.open(new Configuration());

	// arbitrary checkpoint id and timestamp
	consumer.snapshotState(new StateSnapshotContextSynchronousImpl(123, 123));

	assertTrue(listState.isClearCalled());

	// the checkpointed list state should contain only the shards that it should subscribe to
	assertEquals(globalUnionState.size() / 2, listState.getList().size());
	assertTrue(listState.getList().contains(globalUnionState.get(0)));
	assertTrue(listState.getList().contains(globalUnionState.get(2)));
}
 

@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<>(),
			TestUtils.getStandardProperties(),
			new KinesisDeserializationSchemaWrapper<>(new SimpleStringSchema()),
			10,
			2,
			new AtomicReference<>(),
			new LinkedList<>(),
			subscribedStreamsToLastSeenShardIdsUnderTest,
			FakeKinesisBehavioursFactory.nonReshardedStreamsBehaviour(streamToShardCount));

	final DummyFlinkKinesisConsumer<String> consumer = new DummyFlinkKinesisConsumer<>(
			TestUtils.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();

	// assert that the streams tracked in the state are identical to the subscribed streams
	Set<String> streamsInState = subscribedStreamsToLastSeenShardIdsUnderTest.keySet();
	assertEquals(fakeStreams.size(), streamsInState.size());
	assertTrue(streamsInState.containsAll(fakeStreams));

	// assert that the last seen shards in state is correctly set
	for (Map.Entry<String, String> streamToLastSeenShard : subscribedStreamsToLastSeenShardIdsUnderTest.entrySet()) {
		assertEquals(
			KinesisShardIdGenerator.generateFromShardOrder(streamToShardCount.get(streamToLastSeenShard.getKey()) - 1),
			streamToLastSeenShard.getValue());
	}
}
 

@Test
public void testStreamToLastSeenShardStateIsCorrectlySetWhenNoNewShardsSinceRestoredCheckpoint() throws Exception {
	List<String> fakeStreams = new LinkedList<>();
	fakeStreams.add("fakeStream1");
	fakeStreams.add("fakeStream2");

	Map<StreamShardHandle, String> restoredStateUnderTest = new HashMap<>();

	// fakeStream1 has 3 shards before restore
	restoredStateUnderTest.put(
		new StreamShardHandle(
			"fakeStream1",
			new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(0))),
		UUID.randomUUID().toString());
	restoredStateUnderTest.put(
		new StreamShardHandle(
			"fakeStream1",
			new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(1))),
		UUID.randomUUID().toString());
	restoredStateUnderTest.put(
		new StreamShardHandle(
			"fakeStream1",
			new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(2))),
		UUID.randomUUID().toString());

	// fakeStream2 has 2 shards before restore
	restoredStateUnderTest.put(
		new StreamShardHandle(
			"fakeStream2",
			new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(0))),
		UUID.randomUUID().toString());
	restoredStateUnderTest.put(
		new StreamShardHandle(
			"fakeStream2",
			new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(1))),
		UUID.randomUUID().toString());

	Map<String, Integer> streamToShardCount = new HashMap<>();
	streamToShardCount.put("fakeStream1", 3); // fakeStream1 will still have 3 shards after restore
	streamToShardCount.put("fakeStream2", 2); // fakeStream2 will still have 2 shards after restore

	HashMap<String, String> subscribedStreamsToLastSeenShardIdsUnderTest =
		KinesisDataFetcher.createInitialSubscribedStreamsToLastDiscoveredShardsState(fakeStreams);

	final TestableKinesisDataFetcher<String> fetcher =
		new TestableKinesisDataFetcher<>(
			fakeStreams,
			new TestSourceContext<>(),
			TestUtils.getStandardProperties(),
			new KinesisDeserializationSchemaWrapper<>(new SimpleStringSchema()),
			10,
			2,
			new AtomicReference<>(),
			new LinkedList<>(),
			subscribedStreamsToLastSeenShardIdsUnderTest,
			FakeKinesisBehavioursFactory.nonReshardedStreamsBehaviour(streamToShardCount));

	for (Map.Entry<StreamShardHandle, String> restoredState : restoredStateUnderTest.entrySet()) {
		fetcher.advanceLastDiscoveredShardOfStream(restoredState.getKey().getStreamName(), restoredState.getKey().getShard().getShardId());
		fetcher.registerNewSubscribedShardState(
			new KinesisStreamShardState(KinesisDataFetcher.convertToStreamShardMetadata(restoredState.getKey()),
				restoredState.getKey(), new SequenceNumber(restoredState.getValue())));
	}

	CheckedThread runFetcherThread = new CheckedThread() {
		@Override
		public void go() throws Exception {
			fetcher.runFetcher();
		}
	};
	runFetcherThread.start();

	fetcher.waitUntilInitialDiscovery();
	fetcher.shutdownFetcher();
	runFetcherThread.sync();

	// assert that the streams tracked in the state are identical to the subscribed streams
	Set<String> streamsInState = subscribedStreamsToLastSeenShardIdsUnderTest.keySet();
	assertEquals(fakeStreams.size(), streamsInState.size());
	assertTrue(streamsInState.containsAll(fakeStreams));

	// assert that the last seen shards in state is correctly set
	for (Map.Entry<String, String> streamToLastSeenShard : subscribedStreamsToLastSeenShardIdsUnderTest.entrySet()) {
		assertEquals(
			KinesisShardIdGenerator.generateFromShardOrder(streamToShardCount.get(streamToLastSeenShard.getKey()) - 1),
			streamToLastSeenShard.getValue());
	}
}
 

@Test
public void testStreamToLastSeenShardStateIsCorrectlySetWhenNewShardsFoundSinceRestoredCheckpoint() throws Exception {
	List<String> fakeStreams = new LinkedList<>();
	fakeStreams.add("fakeStream1");
	fakeStreams.add("fakeStream2");

	Map<StreamShardHandle, String> restoredStateUnderTest = new HashMap<>();

	// fakeStream1 has 3 shards before restore
	restoredStateUnderTest.put(
		new StreamShardHandle(
			"fakeStream1",
			new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(0))),
		UUID.randomUUID().toString());
	restoredStateUnderTest.put(
		new StreamShardHandle(
			"fakeStream1",
			new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(1))),
		UUID.randomUUID().toString());
	restoredStateUnderTest.put(
		new StreamShardHandle(
			"fakeStream1",
			new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(2))),
		UUID.randomUUID().toString());

	// fakeStream2 has 2 shards before restore
	restoredStateUnderTest.put(
		new StreamShardHandle(
			"fakeStream2",
			new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(0))),
		UUID.randomUUID().toString());
	restoredStateUnderTest.put(
		new StreamShardHandle(
			"fakeStream2",
			new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(1))),
		UUID.randomUUID().toString());

	Map<String, Integer> streamToShardCount = new HashMap<>();
	streamToShardCount.put("fakeStream1", 3 + 1); // fakeStream1 had 3 shards before & 1 new shard after restore
	streamToShardCount.put("fakeStream2", 2 + 3); // fakeStream2 had 2 shards before & 3 new shard after restore

	HashMap<String, String> subscribedStreamsToLastSeenShardIdsUnderTest =
		KinesisDataFetcher.createInitialSubscribedStreamsToLastDiscoveredShardsState(fakeStreams);

	// using a non-resharded streams kinesis behaviour to represent that Kinesis is not resharded AFTER the restore
	final TestableKinesisDataFetcher<String> fetcher =
		new TestableKinesisDataFetcher<>(
			fakeStreams,
			new TestSourceContext<>(),
			TestUtils.getStandardProperties(),
			new KinesisDeserializationSchemaWrapper<>(new SimpleStringSchema()),
			10,
			2,
			new AtomicReference<>(),
			new LinkedList<>(),
			subscribedStreamsToLastSeenShardIdsUnderTest,
			FakeKinesisBehavioursFactory.nonReshardedStreamsBehaviour(streamToShardCount));

	for (Map.Entry<StreamShardHandle, String> restoredState : restoredStateUnderTest.entrySet()) {
		fetcher.advanceLastDiscoveredShardOfStream(restoredState.getKey().getStreamName(), restoredState.getKey().getShard().getShardId());
		fetcher.registerNewSubscribedShardState(
			new KinesisStreamShardState(KinesisDataFetcher.convertToStreamShardMetadata(restoredState.getKey()),
				restoredState.getKey(), new SequenceNumber(restoredState.getValue())));
	}

	CheckedThread runFetcherThread = new CheckedThread() {
		@Override
		public void go() throws Exception {
			fetcher.runFetcher();
		}
	};
	runFetcherThread.start();

	fetcher.waitUntilInitialDiscovery();
	fetcher.shutdownFetcher();
	runFetcherThread.sync();

	// assert that the streams tracked in the state are identical to the subscribed streams
	Set<String> streamsInState = subscribedStreamsToLastSeenShardIdsUnderTest.keySet();
	assertEquals(fakeStreams.size(), streamsInState.size());
	assertTrue(streamsInState.containsAll(fakeStreams));

	// assert that the last seen shards in state is correctly set
	for (Map.Entry<String, String> streamToLastSeenShard : subscribedStreamsToLastSeenShardIdsUnderTest.entrySet()) {
		assertEquals(
			KinesisShardIdGenerator.generateFromShardOrder(streamToShardCount.get(streamToLastSeenShard.getKey()) - 1),
			streamToLastSeenShard.getValue());
	}
}
 

@Test
public void testStreamToLastSeenShardStateIsCorrectlySetWhenNoNewShardsSinceRestoredCheckpointAndSomeStreamsDoNotExist() throws Exception {
	List<String> fakeStreams = new LinkedList<>();
	fakeStreams.add("fakeStream1");
	fakeStreams.add("fakeStream2");
	fakeStreams.add("fakeStream3"); // fakeStream3 will not have any shards
	fakeStreams.add("fakeStream4"); // fakeStream4 will not have any shards

	Map<StreamShardHandle, String> restoredStateUnderTest = new HashMap<>();

	// fakeStream1 has 3 shards before restore
	restoredStateUnderTest.put(
		new StreamShardHandle(
			"fakeStream1",
			new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(0))),
		UUID.randomUUID().toString());
	restoredStateUnderTest.put(
		new StreamShardHandle(
			"fakeStream1",
			new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(1))),
		UUID.randomUUID().toString());
	restoredStateUnderTest.put(
		new StreamShardHandle(
			"fakeStream1",
			new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(2))),
		UUID.randomUUID().toString());

	// fakeStream2 has 2 shards before restore
	restoredStateUnderTest.put(
		new StreamShardHandle(
			"fakeStream2",
			new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(0))),
		UUID.randomUUID().toString());
	restoredStateUnderTest.put(
		new StreamShardHandle(
			"fakeStream2",
			new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(1))),
		UUID.randomUUID().toString());

	Map<String, Integer> streamToShardCount = new HashMap<>();
	streamToShardCount.put("fakeStream1", 3); // fakeStream1 has fixed 3 shards
	streamToShardCount.put("fakeStream2", 2); // fakeStream2 has fixed 2 shards
	streamToShardCount.put("fakeStream3", 0); // no shards can be found for fakeStream3
	streamToShardCount.put("fakeStream4", 0); // no shards can be found for fakeStream4

	HashMap<String, String> subscribedStreamsToLastSeenShardIdsUnderTest =
		KinesisDataFetcher.createInitialSubscribedStreamsToLastDiscoveredShardsState(fakeStreams);

	// using a non-resharded streams kinesis behaviour to represent that Kinesis is not resharded AFTER the restore
	final TestableKinesisDataFetcher<String> fetcher =
		new TestableKinesisDataFetcher<>(
			fakeStreams,
			new TestSourceContext<>(),
			TestUtils.getStandardProperties(),
			new KinesisDeserializationSchemaWrapper<>(new SimpleStringSchema()),
			10,
			2,
			new AtomicReference<>(),
			new LinkedList<>(),
			subscribedStreamsToLastSeenShardIdsUnderTest,
			FakeKinesisBehavioursFactory.nonReshardedStreamsBehaviour(streamToShardCount));

	for (Map.Entry<StreamShardHandle, String> restoredState : restoredStateUnderTest.entrySet()) {
		fetcher.advanceLastDiscoveredShardOfStream(restoredState.getKey().getStreamName(), restoredState.getKey().getShard().getShardId());
		fetcher.registerNewSubscribedShardState(
			new KinesisStreamShardState(KinesisDataFetcher.convertToStreamShardMetadata(restoredState.getKey()),
				restoredState.getKey(), new SequenceNumber(restoredState.getValue())));
	}

	CheckedThread runFetcherThread = new CheckedThread() {
		@Override
		public void go() throws Exception {
			fetcher.runFetcher();
		}
	};
	runFetcherThread.start();

	fetcher.waitUntilInitialDiscovery();
	fetcher.shutdownFetcher();
	runFetcherThread.sync();

	// assert that the streams tracked in the state are identical to the subscribed streams
	Set<String> streamsInState = subscribedStreamsToLastSeenShardIdsUnderTest.keySet();
	assertEquals(fakeStreams.size(), streamsInState.size());
	assertTrue(streamsInState.containsAll(fakeStreams));

	// assert that the last seen shards in state is correctly set
	assertEquals(
		KinesisShardIdGenerator.generateFromShardOrder(2),
		subscribedStreamsToLastSeenShardIdsUnderTest.get("fakeStream1"));
	assertEquals(
		KinesisShardIdGenerator.generateFromShardOrder(1),
		subscribedStreamsToLastSeenShardIdsUnderTest.get("fakeStream2"));
	assertNull(subscribedStreamsToLastSeenShardIdsUnderTest.get("fakeStream3"));
	assertNull(subscribedStreamsToLastSeenShardIdsUnderTest.get("fakeStream4"));
}
 

@Test
public void testStreamToLastSeenShardStateIsCorrectlySetWhenNewShardsFoundSinceRestoredCheckpointAndSomeStreamsDoNotExist() throws Exception {
	List<String> fakeStreams = new LinkedList<>();
	fakeStreams.add("fakeStream1");
	fakeStreams.add("fakeStream2");
	fakeStreams.add("fakeStream3"); // fakeStream3 will not have any shards
	fakeStreams.add("fakeStream4"); // fakeStream4 will not have any shards

	Map<StreamShardHandle, String> restoredStateUnderTest = new HashMap<>();

	// fakeStream1 has 3 shards before restore
	restoredStateUnderTest.put(
		new StreamShardHandle(
			"fakeStream1",
			new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(0))),
		UUID.randomUUID().toString());
	restoredStateUnderTest.put(
		new StreamShardHandle(
			"fakeStream1",
			new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(1))),
		UUID.randomUUID().toString());
	restoredStateUnderTest.put(
		new StreamShardHandle(
			"fakeStream1",
			new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(2))),
		UUID.randomUUID().toString());

	// fakeStream2 has 2 shards before restore
	restoredStateUnderTest.put(
		new StreamShardHandle(
			"fakeStream2",
			new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(0))),
		UUID.randomUUID().toString());
	restoredStateUnderTest.put(
		new StreamShardHandle(
			"fakeStream2",
			new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(1))),
		UUID.randomUUID().toString());

	Map<String, Integer> streamToShardCount = new HashMap<>();
	streamToShardCount.put("fakeStream1", 3 + 1); // fakeStream1 had 3 shards before & 1 new shard after restore
	streamToShardCount.put("fakeStream2", 2 + 3); // fakeStream2 had 2 shards before & 2 new shard after restore
	streamToShardCount.put("fakeStream3", 0); // no shards can be found for fakeStream3
	streamToShardCount.put("fakeStream4", 0); // no shards can be found for fakeStream4

	HashMap<String, String> subscribedStreamsToLastSeenShardIdsUnderTest =
		KinesisDataFetcher.createInitialSubscribedStreamsToLastDiscoveredShardsState(fakeStreams);

	// using a non-resharded streams kinesis behaviour to represent that Kinesis is not resharded AFTER the restore
	final TestableKinesisDataFetcher<String> fetcher =
		new TestableKinesisDataFetcher<>(
			fakeStreams,
			new TestSourceContext<>(),
			TestUtils.getStandardProperties(),
			new KinesisDeserializationSchemaWrapper<>(new SimpleStringSchema()),
			10,
			2,
			new AtomicReference<>(),
			new LinkedList<>(),
			subscribedStreamsToLastSeenShardIdsUnderTest,
			FakeKinesisBehavioursFactory.nonReshardedStreamsBehaviour(streamToShardCount));

	for (Map.Entry<StreamShardHandle, String> restoredState : restoredStateUnderTest.entrySet()) {
		fetcher.advanceLastDiscoveredShardOfStream(restoredState.getKey().getStreamName(), restoredState.getKey().getShard().getShardId());
		fetcher.registerNewSubscribedShardState(
			new KinesisStreamShardState(KinesisDataFetcher.convertToStreamShardMetadata(restoredState.getKey()),
				restoredState.getKey(), new SequenceNumber(restoredState.getValue())));
	}

	CheckedThread runFetcherThread = new CheckedThread() {
		@Override
		public void go() throws Exception {
			fetcher.runFetcher();
		}
	};
	runFetcherThread.start();

	fetcher.waitUntilInitialDiscovery();
	fetcher.shutdownFetcher();
	runFetcherThread.sync();

	// assert that the streams tracked in the state are identical to the subscribed streams
	Set<String> streamsInState = subscribedStreamsToLastSeenShardIdsUnderTest.keySet();
	assertEquals(fakeStreams.size(), streamsInState.size());
	assertTrue(streamsInState.containsAll(fakeStreams));

	// assert that the last seen shards in state is correctly set
	assertEquals(
		KinesisShardIdGenerator.generateFromShardOrder(3),
		subscribedStreamsToLastSeenShardIdsUnderTest.get("fakeStream1"));
	assertEquals(
		KinesisShardIdGenerator.generateFromShardOrder(4),
		subscribedStreamsToLastSeenShardIdsUnderTest.get("fakeStream2"));
	assertNull(subscribedStreamsToLastSeenShardIdsUnderTest.get("fakeStream3"));
	assertNull(subscribedStreamsToLastSeenShardIdsUnderTest.get("fakeStream4"));
}
 

@Test
public void testPeriodicWatermark() {
	final MutableLong clock = new MutableLong();
	final MutableBoolean isTemporaryIdle = new MutableBoolean();
	final List<Watermark> watermarks = new ArrayList<>();

	String fakeStream1 = "fakeStream1";
	StreamShardHandle shardHandle =
		new StreamShardHandle(
			fakeStream1,
			new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(0)));

	TestSourceContext<String> sourceContext =
		new TestSourceContext<String>() {
			@Override
			public void emitWatermark(Watermark mark) {
				watermarks.add(mark);
			}

			@Override
			public void markAsTemporarilyIdle() {
				isTemporaryIdle.setTrue();
			}
		};

	HashMap<String, String> subscribedStreamsToLastSeenShardIdsUnderTest = new HashMap<>();

	final KinesisDataFetcher<String> fetcher =
		new TestableKinesisDataFetcher<String>(
			Collections.singletonList(fakeStream1),
			sourceContext,
			new java.util.Properties(),
			new KinesisDeserializationSchemaWrapper<>(new org.apache.flink.streaming.util.serialization.SimpleStringSchema()),
			1,
			1,
			new AtomicReference<>(),
			new LinkedList<>(),
			subscribedStreamsToLastSeenShardIdsUnderTest,
			FakeKinesisBehavioursFactory.nonReshardedStreamsBehaviour(new HashMap<>())) {

			@Override
			protected long getCurrentTimeMillis() {
				return clock.getValue();
			}
		};
	Whitebox.setInternalState(fetcher, "periodicWatermarkAssigner", watermarkAssigner);

	SequenceNumber seq = new SequenceNumber("fakeSequenceNumber");
	// register shards to subsequently emit records
	int shardIndex =
		fetcher.registerNewSubscribedShardState(
			new KinesisStreamShardState(
				KinesisDataFetcher.convertToStreamShardMetadata(shardHandle), shardHandle, seq));

	StreamRecord<String> record1 =
		new StreamRecord<>(String.valueOf(Long.MIN_VALUE), Long.MIN_VALUE);
	fetcher.emitRecordAndUpdateState(record1.getValue(), record1.getTimestamp(), shardIndex, seq);
	Assert.assertEquals(record1, sourceContext.getCollectedOutputs().poll());

	fetcher.emitWatermark();
	Assert.assertTrue("potential watermark equals previous watermark", watermarks.isEmpty());

	StreamRecord<String> record2 = new StreamRecord<>(String.valueOf(1), 1);
	fetcher.emitRecordAndUpdateState(record2.getValue(), record2.getTimestamp(), shardIndex, seq);
	Assert.assertEquals(record2, sourceContext.getCollectedOutputs().poll());

	fetcher.emitWatermark();
	Assert.assertFalse("watermark advanced", watermarks.isEmpty());
	Assert.assertEquals(new Watermark(record2.getTimestamp()), watermarks.remove(0));
	Assert.assertFalse("not idle", isTemporaryIdle.booleanValue());

	// test idle timeout
	long idleTimeout = 10;
	// advance clock idleTimeout
	clock.add(idleTimeout + 1);
	fetcher.emitWatermark();
	Assert.assertFalse("not idle", isTemporaryIdle.booleanValue());
	Assert.assertTrue("not idle, no new watermark", watermarks.isEmpty());

	// activate idle timeout
	Whitebox.setInternalState(fetcher, "shardIdleIntervalMillis", idleTimeout);
	fetcher.emitWatermark();
	Assert.assertTrue("idle", isTemporaryIdle.booleanValue());
	Assert.assertTrue("idle, no watermark", watermarks.isEmpty());
}