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

The following examples show how to use org.apache.flink.streaming.connectors.kinesis.testutils.KinesisShardIdGenerator. These examples are extracted from open source projects. You can vote up the ones you like or vote down the ones you don't like, and go to the original project or source file by following the links above each example. You may check out the related API usage on the sidebar.
Example 1
Source Project: Flink-CEPplus   Source File: KinesisProxyTest.java    License: Apache License 2.0 6 votes vote down vote up
@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());
}
 
Example 2
Source Project: flink   Source File: KinesisProxyTest.java    License: Apache License 2.0 6 votes vote down vote up
@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());
}
 
Example 3
Source Project: flink   Source File: KinesisProxyTest.java    License: Apache License 2.0 6 votes vote down vote up
@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());
}
 
Example 4
Source Project: Flink-CEPplus   Source File: ShardConsumerTest.java    License: Apache License 2.0 5 votes vote down vote up
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())));
}
 
Example 5
Source Project: Flink-CEPplus   Source File: FlinkKinesisConsumerTest.java    License: Apache License 2.0 5 votes vote down vote up
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;
}
 
Example 6
Source Project: flink   Source File: ShardConsumerTest.java    License: Apache License 2.0 5 votes vote down vote up
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())));
}
 
Example 7
Source Project: flink   Source File: FlinkKinesisConsumerTest.java    License: Apache License 2.0 5 votes vote down vote up
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;
}
 
Example 8
Source Project: flink   Source File: ShardConsumerTest.java    License: Apache License 2.0 5 votes vote down vote up
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())));
}
 
Example 9
Source Project: flink   Source File: FlinkKinesisConsumerTest.java    License: Apache License 2.0 5 votes vote down vote up
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;
}
 
Example 10
Source Project: Flink-CEPplus   Source File: KinesisDataFetcherTest.java    License: Apache License 2.0 4 votes vote down vote up
@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());
	}
}
 
Example 11
Source Project: Flink-CEPplus   Source File: KinesisDataFetcherTest.java    License: Apache License 2.0 4 votes vote down vote up
@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());
	}
}
 
Example 12
Source Project: Flink-CEPplus   Source File: KinesisDataFetcherTest.java    License: Apache License 2.0 4 votes vote down vote up
@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());
	}
}
 
Example 13
Source Project: Flink-CEPplus   Source File: KinesisDataFetcherTest.java    License: Apache License 2.0 4 votes vote down vote up
@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"));
}
 
Example 14
Source Project: Flink-CEPplus   Source File: KinesisDataFetcherTest.java    License: Apache License 2.0 4 votes vote down vote up
@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"));
}
 
Example 15
Source Project: Flink-CEPplus   Source File: KinesisDataFetcherTest.java    License: Apache License 2.0 4 votes vote down vote up
@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());
}
 
Example 16
@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();
}
 
Example 17
Source Project: Flink-CEPplus   Source File: KinesisProxyTest.java    License: Apache License 2.0 4 votes vote down vote up
@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()])));
}
 
Example 18
Source Project: Flink-CEPplus   Source File: KinesisProxyTest.java    License: Apache License 2.0 4 votes vote down vote up
@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()])));
}
 
Example 19
Source Project: Flink-CEPplus   Source File: FlinkKinesisConsumerTest.java    License: Apache License 2.0 4 votes vote down vote up
@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)));
}
 
Example 20
Source Project: Flink-CEPplus   Source File: FlinkKinesisConsumerTest.java    License: Apache License 2.0 4 votes vote down vote up
@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()));
	}
}
 
Example 21
Source Project: flink   Source File: KinesisDataFetcherTest.java    License: Apache License 2.0 4 votes vote down vote up
@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());
	}
}
 
Example 22
Source Project: flink   Source File: KinesisDataFetcherTest.java    License: Apache License 2.0 4 votes vote down vote up
@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());
	}
}
 
Example 23
Source Project: flink   Source File: KinesisDataFetcherTest.java    License: Apache License 2.0 4 votes vote down vote up
@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());
	}
}
 
Example 24
Source Project: flink   Source File: KinesisDataFetcherTest.java    License: Apache License 2.0 4 votes vote down vote up
@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"));
}
 
Example 25
Source Project: flink   Source File: KinesisDataFetcherTest.java    License: Apache License 2.0 4 votes vote down vote up
@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"));
}
 
Example 26
Source Project: flink   Source File: KinesisDataFetcherTest.java    License: Apache License 2.0 4 votes vote down vote up
@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());
}
 
Example 27
Source Project: flink   Source File: FlinkKinesisConsumerMigrationTest.java    License: Apache License 2.0 4 votes vote down vote up
@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();
}
 
Example 28
Source Project: flink   Source File: KinesisProxyTest.java    License: Apache License 2.0 4 votes vote down vote up
@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()])));
}
 
Example 29
Source Project: flink   Source File: KinesisProxyTest.java    License: Apache License 2.0 4 votes vote down vote up
@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()])));
}
 
Example 30
Source Project: flink   Source File: FlinkKinesisConsumerTest.java    License: Apache License 2.0 4 votes vote down vote up
@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)));
}