com.amazonaws.services.kinesis.model.ListShardsResult Java Examples

The following examples show how to use com.amazonaws.services.kinesis.model.ListShardsResult. 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 File: KinesisProxy.java    From Flink-CEPplus with Apache License 2.0 6 votes vote down vote up 
private List<StreamShardHandle> getShardsOfStream(String streamName, @Nullable String lastSeenShardId) throws InterruptedException {
	List<StreamShardHandle> shardsOfStream = new ArrayList<>();

	// List Shards returns just the first 1000 shard entries. In order to read the entire stream,
	// we need to use the returned nextToken to get additional shards.
	ListShardsResult listShardsResult;
	String startShardToken = null;
	do {
		listShardsResult = listShards(streamName, lastSeenShardId, startShardToken);
		if (listShardsResult == null) {
			// In case we have exceptions while retrieving all shards, ensure that incomplete shard list is not returned.
			// Hence clearing the incomplete shard list before returning it.
			shardsOfStream.clear();
			return shardsOfStream;
		}
		List<Shard> shards = listShardsResult.getShards();
		for (Shard shard : shards) {
			shardsOfStream.add(new StreamShardHandle(streamName, shard));
		}
		startShardToken = listShardsResult.getNextToken();
	} while (startShardToken != null);

	return shardsOfStream;
}
Example #2
Source File: KinesisProxyTest.java    From flink with 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 File: KinesisProxyTest.java    From Flink-CEPplus with 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 File: KinesisProxy.java    From flink with Apache License 2.0 6 votes vote down vote up 
private List<StreamShardHandle> getShardsOfStream(String streamName, @Nullable String lastSeenShardId) throws InterruptedException {
	List<StreamShardHandle> shardsOfStream = new ArrayList<>();

	// List Shards returns just the first 1000 shard entries. In order to read the entire stream,
	// we need to use the returned nextToken to get additional shards.
	ListShardsResult listShardsResult;
	String startShardToken = null;
	do {
		listShardsResult = listShards(streamName, lastSeenShardId, startShardToken);
		if (listShardsResult == null) {
			// In case we have exceptions while retrieving all shards, ensure that incomplete shard list is not returned.
			// Hence clearing the incomplete shard list before returning it.
			shardsOfStream.clear();
			return shardsOfStream;
		}
		List<Shard> shards = listShardsResult.getShards();
		for (Shard shard : shards) {
			shardsOfStream.add(new StreamShardHandle(streamName, shard));
		}
		startShardToken = listShardsResult.getNextToken();
	} while (startShardToken != null);

	return shardsOfStream;
}
Example #5
Source File: KinesisProxy.java    From flink with Apache License 2.0 6 votes vote down vote up 
private List<StreamShardHandle> getShardsOfStream(String streamName, @Nullable String lastSeenShardId) throws InterruptedException {
	List<StreamShardHandle> shardsOfStream = new ArrayList<>();

	// List Shards returns just the first 1000 shard entries. In order to read the entire stream,
	// we need to use the returned nextToken to get additional shards.
	ListShardsResult listShardsResult;
	String startShardToken = null;
	do {
		listShardsResult = listShards(streamName, lastSeenShardId, startShardToken);
		if (listShardsResult == null) {
			// In case we have exceptions while retrieving all shards, ensure that incomplete shard list is not returned.
			// Hence clearing the incomplete shard list before returning it.
			shardsOfStream.clear();
			return shardsOfStream;
		}
		List<Shard> shards = listShardsResult.getShards();
		for (Shard shard : shards) {
			shardsOfStream.add(new StreamShardHandle(streamName, shard));
		}
		startShardToken = listShardsResult.getNextToken();
	} while (startShardToken != null);

	return shardsOfStream;
}
Example #6
Source File: KinesisProxyTest.java    From flink with 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 #7
Source File: KinesisStreamProvisioner.java    From spring-cloud-stream-binder-aws-kinesis with Apache License 2.0 5 votes vote down vote up 
private List<Shard> getShardList(String stream, int retryCount) {
	List<Shard> shardList = new ArrayList<>();

	if (retryCount > configurationProperties.getDescribeStreamRetries()) {
		ResourceNotFoundException resourceNotFoundException = new ResourceNotFoundException(
				"The stream [" + stream + "] isn't ACTIVE or doesn't exist.");
		resourceNotFoundException.setServiceName("Kinesis");

		throw new ProvisioningException(
				"Kinesis org.springframework.cloud.stream.binder.kinesis.provisioning error",
				resourceNotFoundException);
	}

	ListShardsRequest listShardsRequest = new ListShardsRequest().withStreamName(stream);

	try {
		ListShardsResult listShardsResult = amazonKinesis.listShards(listShardsRequest);

		shardList.addAll(listShardsResult.getShards());

	}
	catch (LimitExceededException limitExceededException) {
		logger.info("Got LimitExceededException when describing stream [" + stream + "]. " + "Backing off for ["
				+ this.configurationProperties.getDescribeStreamBackoff() + "] millis.");

		try {
			Thread.sleep(this.configurationProperties.getDescribeStreamBackoff());
			getShardList(stream, retryCount++);
		}
		catch (InterruptedException ex) {
			Thread.currentThread().interrupt();
			throw new ProvisioningException(
					"The [describeStream] thread for the stream [" + stream + "] has been interrupted.", ex);
		}
	}

	return shardList;
}
Example #8
Source File: KinesisProxyTest.java    From flink with Apache License 2.0 4 votes vote down vote up 
@Test
public void testGetShardListRetry() throws Exception {
	Properties kinesisConsumerConfig = new Properties();
	kinesisConsumerConfig.setProperty(ConsumerConfigConstants.AWS_REGION, "us-east-1");

	Shard shard = new Shard();
	shard.setShardId("fake-shard-000000000000");
	final ListShardsResult expectedResult = new ListShardsResult();
	expectedResult.withShards(shard);

	MutableInt exceptionCount = new MutableInt();
	final Throwable[] retriableExceptions = new Throwable[]{
		new AmazonKinesisException("attempt1"),
		new AmazonKinesisException("attempt2"),
	};

	AmazonKinesisClient mockClient = mock(AmazonKinesisClient.class);
	Mockito.when(mockClient.listShards(any())).thenAnswer(new Answer<ListShardsResult>() {

		@Override
		public ListShardsResult answer(InvocationOnMock invocation) throws Throwable {
			if (exceptionCount.intValue() < retriableExceptions.length) {
				exceptionCount.increment();
				throw retriableExceptions[exceptionCount.intValue() - 1];
			}
			return expectedResult;
		}
	});

	KinesisProxy kinesisProxy = new KinesisProxy(kinesisConsumerConfig);
	Whitebox.getField(KinesisProxy.class, "kinesisClient").set(kinesisProxy, mockClient);

	HashMap<String, String> streamNames = new HashMap();
	streamNames.put("fake-stream", null);
	GetShardListResult result = kinesisProxy.getShardList(streamNames);
	assertEquals(retriableExceptions.length, exceptionCount.intValue());
	assertEquals(true, result.hasRetrievedShards());
	assertEquals(shard.getShardId(), result.getLastSeenShardOfStream("fake-stream").getShard().getShardId());

	// test max attempt count exceeded
	int maxRetries = 1;
	exceptionCount.setValue(0);
	kinesisConsumerConfig.setProperty(ConsumerConfigConstants.LIST_SHARDS_RETRIES, String.valueOf(maxRetries));
	kinesisProxy = new KinesisProxy(kinesisConsumerConfig);
	Whitebox.getField(KinesisProxy.class, "kinesisClient").set(kinesisProxy, mockClient);
	try {
		kinesisProxy.getShardList(streamNames);
		Assert.fail("exception expected");
	} catch (SdkClientException ex) {
		assertEquals(retriableExceptions[maxRetries], ex);
	}
	assertEquals(maxRetries + 1, exceptionCount.intValue());
}
Example #9
Source File: KinesisProxyTest.java    From flink with 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 #10
Source File: KinesisProxyTest.java    From flink with 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 #11
Source File: AmazonKinesisMock.java    From beam with Apache License 2.0 4 votes vote down vote up 
@Override
public ListShardsResult listShards(ListShardsRequest listShardsRequest) {
  throw new RuntimeException("Not implemented");
}
Example #12
Source File: KinesisProxyTest.java    From flink with Apache License 2.0 4 votes vote down vote up 
@Test
public void testGetShardListRetry() throws Exception {
	Properties kinesisConsumerConfig = new Properties();
	kinesisConsumerConfig.setProperty(ConsumerConfigConstants.AWS_REGION, "us-east-1");

	Shard shard = new Shard();
	shard.setShardId("fake-shard-000000000000");
	final ListShardsResult expectedResult = new ListShardsResult();
	expectedResult.withShards(shard);

	MutableInt exceptionCount = new MutableInt();
	final Throwable[] retriableExceptions = new Throwable[]{
		new AmazonKinesisException("attempt1"),
		new AmazonKinesisException("attempt2"),
	};

	AmazonKinesisClient mockClient = mock(AmazonKinesisClient.class);
	Mockito.when(mockClient.listShards(any())).thenAnswer(new Answer<ListShardsResult>() {

		@Override
		public ListShardsResult answer(InvocationOnMock invocation) throws Throwable {
			if (exceptionCount.intValue() < retriableExceptions.length) {
				exceptionCount.increment();
				throw retriableExceptions[exceptionCount.intValue() - 1];
			}
			return expectedResult;
		}
	});

	KinesisProxy kinesisProxy = new KinesisProxy(kinesisConsumerConfig);
	Whitebox.getField(KinesisProxy.class, "kinesisClient").set(kinesisProxy, mockClient);

	HashMap<String, String> streamNames = new HashMap();
	streamNames.put("fake-stream", null);
	GetShardListResult result = kinesisProxy.getShardList(streamNames);
	assertEquals(retriableExceptions.length, exceptionCount.intValue());
	assertEquals(true, result.hasRetrievedShards());
	assertEquals(shard.getShardId(), result.getLastSeenShardOfStream("fake-stream").getShard().getShardId());

	// test max attempt count exceeded
	int maxRetries = 1;
	exceptionCount.setValue(0);
	kinesisConsumerConfig.setProperty(ConsumerConfigConstants.LIST_SHARDS_RETRIES, String.valueOf(maxRetries));
	kinesisProxy = new KinesisProxy(kinesisConsumerConfig);
	Whitebox.getField(KinesisProxy.class, "kinesisClient").set(kinesisProxy, mockClient);
	try {
		kinesisProxy.getShardList(streamNames);
		Assert.fail("exception expected");
	} catch (SdkClientException ex) {
		assertEquals(retriableExceptions[maxRetries], ex);
	}
	assertEquals(maxRetries + 1, exceptionCount.intValue());
}
Example #13
Source File: KinesisProxyTest.java    From flink with 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 #14
Source File: KinesisProxyTest.java    From flink with 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 #15
Source File: KinesisProxyTest.java    From Flink-CEPplus with Apache License 2.0 4 votes vote down vote up 
@Test
public void testGetShardListRetry() throws Exception {
	Properties kinesisConsumerConfig = new Properties();
	kinesisConsumerConfig.setProperty(ConsumerConfigConstants.AWS_REGION, "us-east-1");

	Shard shard = new Shard();
	shard.setShardId("fake-shard-000000000000");
	final ListShardsResult expectedResult = new ListShardsResult();
	expectedResult.withShards(shard);

	MutableInt exceptionCount = new MutableInt();
	final Throwable[] retriableExceptions = new Throwable[]{
		new AmazonKinesisException("attempt1"),
		new AmazonKinesisException("attempt2"),
	};

	AmazonKinesisClient mockClient = mock(AmazonKinesisClient.class);
	Mockito.when(mockClient.listShards(any())).thenAnswer(new Answer<ListShardsResult>() {

		@Override
		public ListShardsResult answer(InvocationOnMock invocation) throws Throwable {
			if (exceptionCount.intValue() < retriableExceptions.length) {
				exceptionCount.increment();
				throw retriableExceptions[exceptionCount.intValue() - 1];
			}
			return expectedResult;
		}
	});

	KinesisProxy kinesisProxy = new KinesisProxy(kinesisConsumerConfig);
	Whitebox.getField(KinesisProxy.class, "kinesisClient").set(kinesisProxy, mockClient);

	HashMap<String, String> streamNames = new HashMap();
	streamNames.put("fake-stream", null);
	GetShardListResult result = kinesisProxy.getShardList(streamNames);
	assertEquals(retriableExceptions.length, exceptionCount.intValue());
	assertEquals(true, result.hasRetrievedShards());
	assertEquals(shard.getShardId(), result.getLastSeenShardOfStream("fake-stream").getShard().getShardId());

	// test max attempt count exceeded
	int maxRetries = 1;
	exceptionCount.setValue(0);
	kinesisConsumerConfig.setProperty(ConsumerConfigConstants.LIST_SHARDS_RETRIES, String.valueOf(maxRetries));
	kinesisProxy = new KinesisProxy(kinesisConsumerConfig);
	Whitebox.getField(KinesisProxy.class, "kinesisClient").set(kinesisProxy, mockClient);
	try {
		kinesisProxy.getShardList(streamNames);
		Assert.fail("exception expected");
	} catch (SdkClientException ex) {
		assertEquals(retriableExceptions[maxRetries], ex);
	}
	assertEquals(maxRetries + 1, exceptionCount.intValue());
}
Example #16
Source File: KinesisProxyTest.java    From Flink-CEPplus with 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 #17
Source File: KinesisProxyTest.java    From Flink-CEPplus with 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()])));
}