Java Code Examples for org.apache.flink.streaming.connectors.kinesis.testutils.TestableFlinkKinesisConsumer#run()

@Test
@SuppressWarnings("unchecked")
public void testFetcherShouldNotBeRestoringFromFailureIfNotRestoringFromCheckpoint() throws Exception {
	KinesisDataFetcher mockedFetcher = mockKinesisDataFetcher();

	// 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));
}
 

@Test
@SuppressWarnings("unchecked")
public void testFetcherShouldNotBeRestoringFromFailureIfNotRestoringFromCheckpoint() throws Exception {
	KinesisDataFetcher mockedFetcher = mockKinesisDataFetcher();

	// 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));
}
 

@Test
@SuppressWarnings("unchecked")
public void testFetcherShouldNotBeRestoringFromFailureIfNotRestoringFromCheckpoint() throws Exception {
	KinesisDataFetcher mockedFetcher = mockKinesisDataFetcher();

	// 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));
}
 

@Test
@SuppressWarnings("unchecked")
public void testFetcherShouldBeCorrectlySeededIfRestoringFromCheckpoint() 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());
	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));

	for (Map.Entry<StreamShardHandle, SequenceNumber> restoredShard : fakeRestoredState.entrySet()) {
		Mockito.verify(mockedFetcher).registerNewSubscribedShardState(
			new KinesisStreamShardState(KinesisDataFetcher.convertToStreamShardMetadata(restoredShard.getKey()),
				restoredShard.getKey(), restoredShard.getValue()));
	}
}
 

@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()));
	}
}
 

/**
 * FLINK-8484: ensure that a state change in the StreamShardMetadata other than {@link StreamShardMetadata#shardId} or
 * {@link StreamShardMetadata#streamName} does not result in the shard not being able to be restored.
 * This handles the corner case where the stored shard metadata is open (no ending sequence number), but after the
 * job restore, the shard has been closed (ending number set) due to re-sharding, and we can no longer rely on
 * {@link StreamShardMetadata#equals(Object)} to find back the sequence number in the collection of restored shard metadata.
 * <p></p>
 * Therefore, we will rely on synchronizing the snapshot's state with the Kinesis shard before attempting to find back
 * the sequence number to restore.
 */
@Test
public void testFindSequenceNumberToRestoreFromIfTheShardHasBeenClosedSinceTheStateWasStored() 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<>();

	// create a fake stream shard handle based on the first entry in the restored state
	final StreamShardHandle originalStreamShardHandle = fakeRestoredState.keySet().iterator().next();
	final StreamShardHandle closedStreamShardHandle = new StreamShardHandle(originalStreamShardHandle.getStreamName(), originalStreamShardHandle.getShard());
	// close the shard handle by setting an ending sequence number
	final SequenceNumberRange sequenceNumberRange = new SequenceNumberRange();
	sequenceNumberRange.setEndingSequenceNumber("1293844");
	closedStreamShardHandle.getShard().setSequenceNumberRange(sequenceNumberRange);

	shards.add(closedStreamShardHandle);

	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));

	Mockito.verify(mockedFetcher).registerNewSubscribedShardState(
		new KinesisStreamShardState(KinesisDataFetcher.convertToStreamShardMetadata(closedStreamShardHandle),
			closedStreamShardHandle, fakeRestoredState.get(closedStreamShardHandle)));
}
 

@Test
@SuppressWarnings("unchecked")
public void testFetcherShouldBeCorrectlySeededIfRestoringFromCheckpoint() 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());
	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));

	for (Map.Entry<StreamShardHandle, SequenceNumber> restoredShard : fakeRestoredState.entrySet()) {
		Mockito.verify(mockedFetcher).registerNewSubscribedShardState(
			new KinesisStreamShardState(KinesisDataFetcher.convertToStreamShardMetadata(restoredShard.getKey()),
				restoredShard.getKey(), restoredShard.getValue()));
	}
}
 

@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()));
	}
}
 

/**
 * FLINK-8484: ensure that a state change in the StreamShardMetadata other than {@link StreamShardMetadata#getShardId()} or
 * {@link StreamShardMetadata#getStreamName()} does not result in the shard not being able to be restored.
 * This handles the corner case where the stored shard metadata is open (no ending sequence number), but after the
 * job restore, the shard has been closed (ending number set) due to re-sharding, and we can no longer rely on
 * {@link StreamShardMetadata#equals(Object)} to find back the sequence number in the collection of restored shard metadata.
 * <p></p>
 * Therefore, we will rely on synchronizing the snapshot's state with the Kinesis shard before attempting to find back
 * the sequence number to restore.
 */
@Test
public void testFindSequenceNumberToRestoreFromIfTheShardHasBeenClosedSinceTheStateWasStored() 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<>();

	// create a fake stream shard handle based on the first entry in the restored state
	final StreamShardHandle originalStreamShardHandle = fakeRestoredState.keySet().iterator().next();
	final StreamShardHandle closedStreamShardHandle = new StreamShardHandle(originalStreamShardHandle.getStreamName(), originalStreamShardHandle.getShard());
	// close the shard handle by setting an ending sequence number
	final SequenceNumberRange sequenceNumberRange = new SequenceNumberRange();
	sequenceNumberRange.setEndingSequenceNumber("1293844");
	closedStreamShardHandle.getShard().setSequenceNumberRange(sequenceNumberRange);

	shards.add(closedStreamShardHandle);

	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));

	Mockito.verify(mockedFetcher).registerNewSubscribedShardState(
		new KinesisStreamShardState(KinesisDataFetcher.convertToStreamShardMetadata(closedStreamShardHandle),
			closedStreamShardHandle, fakeRestoredState.get(closedStreamShardHandle)));
}
 

@Test
@SuppressWarnings("unchecked")
public void testFetcherShouldBeCorrectlySeededIfRestoringFromCheckpoint() 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());
	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));

	for (Map.Entry<StreamShardHandle, SequenceNumber> restoredShard : fakeRestoredState.entrySet()) {
		Mockito.verify(mockedFetcher).registerNewSubscribedShardState(
			new KinesisStreamShardState(KinesisDataFetcher.convertToStreamShardMetadata(restoredShard.getKey()),
				restoredShard.getKey(), restoredShard.getValue()));
	}
}
 

@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()));
	}
}
 

/**
 * FLINK-8484: ensure that a state change in the StreamShardMetadata other than {@link StreamShardMetadata#getShardId()} or
 * {@link StreamShardMetadata#getStreamName()} does not result in the shard not being able to be restored.
 * This handles the corner case where the stored shard metadata is open (no ending sequence number), but after the
 * job restore, the shard has been closed (ending number set) due to re-sharding, and we can no longer rely on
 * {@link StreamShardMetadata#equals(Object)} to find back the sequence number in the collection of restored shard metadata.
 * <p></p>
 * Therefore, we will rely on synchronizing the snapshot's state with the Kinesis shard before attempting to find back
 * the sequence number to restore.
 */
@Test
public void testFindSequenceNumberToRestoreFromIfTheShardHasBeenClosedSinceTheStateWasStored() 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<>();

	// create a fake stream shard handle based on the first entry in the restored state
	final StreamShardHandle originalStreamShardHandle = fakeRestoredState.keySet().iterator().next();
	final StreamShardHandle closedStreamShardHandle = new StreamShardHandle(originalStreamShardHandle.getStreamName(), originalStreamShardHandle.getShard());
	// close the shard handle by setting an ending sequence number
	final SequenceNumberRange sequenceNumberRange = new SequenceNumberRange();
	sequenceNumberRange.setEndingSequenceNumber("1293844");
	closedStreamShardHandle.getShard().setSequenceNumberRange(sequenceNumberRange);

	shards.add(closedStreamShardHandle);

	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));

	Mockito.verify(mockedFetcher).registerNewSubscribedShardState(
		new KinesisStreamShardState(KinesisDataFetcher.convertToStreamShardMetadata(closedStreamShardHandle),
			closedStreamShardHandle, fakeRestoredState.get(closedStreamShardHandle)));
}