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

@Test(expected = IllegalStateException.class)
public void shouldThrowExceptionWhenSuccessorsNotFoundForExpiredShard() throws Exception {
  // given
  StartingPoint latestStartingPoint = new StartingPoint(InitialPositionInStream.LATEST);
  Shard closedShard10 =
      createClosedShard("0010").withParentShardId("0008").withAdjacentParentShardId("0005");
  List<Shard> shards =
      ImmutableList.of(
          shard00,
          shard01,
          shard02,
          shard03,
          shard04,
          shard05,
          shard06,
          shard07,
          shard08,
          shard09,
          closedShard10);

  when(kinesis.listShards(STREAM_NAME)).thenReturn(shards);

  // when
  underTest.findShardsAtStartingPoint(kinesis, STREAM_NAME, latestStartingPoint);
}
 

@Override
public ConnectorSplitSource getSplits(ConnectorTransactionHandle transactionHandle, ConnectorSession session, ConnectorTableHandle table, ConnectorSplitManager.SplitSchedulingStrategy splitSchedulingStrategy)
{
    KinesisTableHandle kinesisTableHandle = (KinesisTableHandle) table;

    InternalStreamDescription description = this.getStreamDescription(kinesisTableHandle.getStreamName());

    ImmutableList.Builder<ConnectorSplit> builder = ImmutableList.builder();
    for (Shard shard : description.getShards()) {
        KinesisSplit split = new KinesisSplit(
                kinesisTableHandle.getStreamName(),
                kinesisTableHandle.getMessageDataFormat(),
                shard.getShardId(),
                shard.getSequenceNumberRange().getStartingSequenceNumber(),
                shard.getSequenceNumberRange().getEndingSequenceNumber());
        builder.add(split);
    }

    return new FixedSplitSource(builder.build());
}
 

/**
 * Validates the shards at the given startingPoint. Validity is checked by getting an iterator at
 * the startingPoint and then trying to read some records. This action does not affect the records
 * at all. If the shard is valid then it will get read from exactly the same point and these
 * records will be read again.
 */
private Set<Shard> validateShards(
    SimplifiedKinesisClient kinesis,
    Iterable<Shard> rootShards,
    String streamName,
    StartingPoint startingPoint)
    throws TransientKinesisException {
  Set<Shard> validShards = new HashSet<>();
  ShardIteratorType shardIteratorType =
      ShardIteratorType.fromValue(startingPoint.getPositionName());
  for (Shard shard : rootShards) {
    String shardIterator =
        kinesis.getShardIterator(
            streamName,
            shard.getShardId(),
            shardIteratorType,
            null,
            startingPoint.getTimestamp());
    GetKinesisRecordsResult records =
        kinesis.getRecords(shardIterator, streamName, shard.getShardId());
    if (records.getNextShardIterator() != null || !records.getRecords().isEmpty()) {
      validShards.add(shard);
    }
  }
  return validShards;
}
 

@Test
@Ignore("Kinesalite doesn't support updateShardCount. Test only against real AWS Kinesis")
public void testPartitionCountIncreasedIfAutoAddPartitionsSet() {
	KinesisBinderConfigurationProperties configurationProperties = new KinesisBinderConfigurationProperties();

	String stream = "existing" + System.currentTimeMillis();

	AMAZON_KINESIS.createStream(stream, 1);

	List<Shard> shards = describeStream(stream);

	assertThat(shards.size()).isEqualTo(1);

	configurationProperties.setMinShardCount(6);
	configurationProperties.setAutoAddShards(true);
	KinesisTestBinder binder = getBinder(configurationProperties);

	ExtendedConsumerProperties<KinesisConsumerProperties> consumerProperties = createConsumerProperties();
	Binding<?> binding = binder.bindConsumer(stream, "test", new NullChannel(),
			consumerProperties);
	binding.unbind();

	shards = describeStream(stream);

	assertThat(shards.size()).isEqualTo(6);
}
 

/**
 * Utility function to convert {@link StreamShardMetadata} into {@link StreamShardHandle}.
 *
 * @param streamShardMetadata the {@link StreamShardMetadata} to be converted
 * @return a {@link StreamShardHandle} object
 */
public static StreamShardHandle convertToStreamShardHandle(StreamShardMetadata streamShardMetadata) {
	Shard shard = new Shard();
	shard.withShardId(streamShardMetadata.getShardId());
	shard.withParentShardId(streamShardMetadata.getParentShardId());
	shard.withAdjacentParentShardId(streamShardMetadata.getAdjacentParentShardId());

	HashKeyRange hashKeyRange = new HashKeyRange();
	hashKeyRange.withStartingHashKey(streamShardMetadata.getStartingHashKey());
	hashKeyRange.withEndingHashKey(streamShardMetadata.getEndingHashKey());
	shard.withHashKeyRange(hashKeyRange);

	SequenceNumberRange sequenceNumberRange = new SequenceNumberRange();
	sequenceNumberRange.withStartingSequenceNumber(streamShardMetadata.getStartingSequenceNumber());
	sequenceNumberRange.withEndingSequenceNumber(streamShardMetadata.getEndingSequenceNumber());
	shard.withSequenceNumberRange(sequenceNumberRange);

	return new StreamShardHandle(streamShardMetadata.getStreamName(), shard);
}
 

/**
 * Utility function to convert {@link StreamShardMetadata} into {@link StreamShardHandle}.
 *
 * @param streamShardMetadata the {@link StreamShardMetadata} to be converted
 * @return a {@link StreamShardHandle} object
 */
public static StreamShardHandle convertToStreamShardHandle(StreamShardMetadata streamShardMetadata) {
	Shard shard = new Shard();
	shard.withShardId(streamShardMetadata.getShardId());
	shard.withParentShardId(streamShardMetadata.getParentShardId());
	shard.withAdjacentParentShardId(streamShardMetadata.getAdjacentParentShardId());

	HashKeyRange hashKeyRange = new HashKeyRange();
	hashKeyRange.withStartingHashKey(streamShardMetadata.getStartingHashKey());
	hashKeyRange.withEndingHashKey(streamShardMetadata.getEndingHashKey());
	shard.withHashKeyRange(hashKeyRange);

	SequenceNumberRange sequenceNumberRange = new SequenceNumberRange();
	sequenceNumberRange.withStartingSequenceNumber(streamShardMetadata.getStartingSequenceNumber());
	sequenceNumberRange.withEndingSequenceNumber(streamShardMetadata.getEndingSequenceNumber());
	shard.withSequenceNumberRange(sequenceNumberRange);

	return new StreamShardHandle(streamShardMetadata.getStreamName(), shard);
}
 

public BlockingQueueKinesis(Map<String, List<BlockingQueue<String>>> streamsToShardCount) {
	for (Map.Entry<String, List<BlockingQueue<String>>> streamToShardQueues : streamsToShardCount.entrySet()) {
		String streamName = streamToShardQueues.getKey();
		int shardCount = streamToShardQueues.getValue().size();

		if (shardCount == 0) {
			// don't do anything
		} else {
			List<StreamShardHandle> shardsOfStream = new ArrayList<>(shardCount);
			for (int i = 0; i < shardCount; i++) {
				StreamShardHandle shardHandle = new StreamShardHandle(
					streamName,
					new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(i))
						.withSequenceNumberRange(new SequenceNumberRange().withStartingSequenceNumber("0"))
						.withHashKeyRange(new HashKeyRange().withStartingHashKey("0").withEndingHashKey("0")));
				shardsOfStream.add(shardHandle);
				shardIteratorToQueueMap.put(getShardIterator(shardHandle), streamToShardQueues.getValue().get(i));
			}
			streamsWithListOfShards.put(streamName, shardsOfStream);
		}
	}
}
 

/**
 * Utility function to convert {@link StreamShardMetadata} into {@link StreamShardHandle}.
 *
 * @param streamShardMetadata the {@link StreamShardMetadata} to be converted
 * @return a {@link StreamShardHandle} object
 */
public static StreamShardHandle convertToStreamShardHandle(StreamShardMetadata streamShardMetadata) {
	Shard shard = new Shard();
	shard.withShardId(streamShardMetadata.getShardId());
	shard.withParentShardId(streamShardMetadata.getParentShardId());
	shard.withAdjacentParentShardId(streamShardMetadata.getAdjacentParentShardId());

	HashKeyRange hashKeyRange = new HashKeyRange();
	hashKeyRange.withStartingHashKey(streamShardMetadata.getStartingHashKey());
	hashKeyRange.withEndingHashKey(streamShardMetadata.getEndingHashKey());
	shard.withHashKeyRange(hashKeyRange);

	SequenceNumberRange sequenceNumberRange = new SequenceNumberRange();
	sequenceNumberRange.withStartingSequenceNumber(streamShardMetadata.getStartingSequenceNumber());
	sequenceNumberRange.withEndingSequenceNumber(streamShardMetadata.getEndingSequenceNumber());
	shard.withSequenceNumberRange(sequenceNumberRange);

	return new StreamShardHandle(streamShardMetadata.getStreamName(), shard);
}
 

@Test
public void shouldMapAllValidShardsToCheckpoints() throws Exception {
  when(shard1.getShardId()).thenReturn("shard-01");
  when(shard2.getShardId()).thenReturn("shard-02");
  when(shard3.getShardId()).thenReturn("shard-03");
  String streamName = "stream";
  Set<Shard> shards = Sets.newHashSet(shard1, shard2);
  StartingPoint startingPoint = new StartingPoint(InitialPositionInStream.LATEST);
  when(startingPointShardsFinder.findShardsAtStartingPoint(
          kinesisClient, "stream", startingPoint))
      .thenReturn(shards);

  DynamicCheckpointGenerator underTest =
      new DynamicCheckpointGenerator(streamName, startingPoint, startingPointShardsFinder);

  KinesisReaderCheckpoint checkpoint = underTest.generate(kinesisClient);
  assertThat(checkpoint)
      .hasSize(2)
      .doesNotContain(new ShardCheckpoint(streamName, shard3.getShardId(), startingPoint));
}
 

private List<StreamShardHandle> getShardsOfStream(
		String streamName,
		@Nullable String lastSeenShardId)
		throws InterruptedException {
	List<StreamShardHandle> shardsOfStream = new ArrayList<>();

	DescribeStreamResult describeStreamResult;
	do {
		describeStreamResult = describeStream(streamName, lastSeenShardId);
		List<Shard> shards = describeStreamResult.getStreamDescription().getShards();
		for (Shard shard : shards) {
			shardsOfStream.add(new StreamShardHandle(streamName, shard));
		}

		if (shards.size() != 0) {
			lastSeenShardId = shards.get(shards.size() - 1).getShardId();
		}
	} while (describeStreamResult.getStreamDescription().isHasMoreShards());

	return shardsOfStream;
}
 

@Test
public void shouldFindFirstShardsWhenAllShardsAreValid() throws Exception {
  // given
  Instant timestampAtTheBeginning = new Instant();
  StartingPoint startingPointAtTheBeginning = new StartingPoint(timestampAtTheBeginning);
  for (Shard shard : allShards) {
    activeAtTimestamp(shard, timestampAtTheBeginning);
  }
  when(kinesis.listShards(STREAM_NAME)).thenReturn(allShards);

  // when
  Iterable<Shard> shardsAtStartingPoint =
      underTest.findShardsAtStartingPoint(kinesis, STREAM_NAME, startingPointAtTheBeginning);

  // then
  assertThat(shardsAtStartingPoint).containsExactlyInAnyOrder(shard00, shard01);
}
 

public BlockingQueueKinesis(Map<String, List<BlockingQueue<String>>> streamsToShardCount) {
	for (Map.Entry<String, List<BlockingQueue<String>>> streamToShardQueues : streamsToShardCount.entrySet()) {
		String streamName = streamToShardQueues.getKey();
		int shardCount = streamToShardQueues.getValue().size();

		if (shardCount == 0) {
			// don't do anything
		} else {
			List<StreamShardHandle> shardsOfStream = new ArrayList<>(shardCount);
			for (int i = 0; i < shardCount; i++) {
				StreamShardHandle shardHandle = new StreamShardHandle(
					streamName,
					new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(i))
						.withSequenceNumberRange(new SequenceNumberRange().withStartingSequenceNumber("0"))
						.withHashKeyRange(new HashKeyRange().withStartingHashKey("0").withEndingHashKey("0")));
				shardsOfStream.add(shardHandle);
				shardIteratorToQueueMap.put(getShardIterator(shardHandle), streamToShardQueues.getValue().get(i));
			}
			streamsWithListOfShards.put(streamName, shardsOfStream);
		}
	}
}
 

public NonReshardedStreamsKinesis(Map<String, Integer> streamsToShardCount) {
	for (Map.Entry<String, Integer> streamToShardCount : streamsToShardCount.entrySet()) {
		String streamName = streamToShardCount.getKey();
		int shardCount = streamToShardCount.getValue();

		if (shardCount == 0) {
			// don't do anything
		} else {
			List<StreamShardHandle> shardsOfStream = new ArrayList<>(shardCount);
			for (int i = 0; i < shardCount; i++) {
				shardsOfStream.add(
					new StreamShardHandle(
						streamName,
						new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(i))));
			}
			streamsWithListOfShards.put(streamName, shardsOfStream);
		}
	}
}
 

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

/**
 * Internal method to retrieve the stream description and get the shards from AWS.
 *
 * Gets from the internal cache unless not yet created or too old.
 *
 * @param streamName Kinesis stream name
 * @return Returns kinesis stream description
 */
protected InternalStreamDescription getStreamDescription(String streamName)
{
    InternalStreamDescription desc = this.streamMap.get(streamName);
    if (desc == null || System.currentTimeMillis() - desc.getCreateTimeStamp() >= MAX_CACHE_AGE_MILLIS) {
        desc = new InternalStreamDescription(streamName);

        DescribeStreamRequest describeStreamRequest = clientManager.getDescribeStreamRequest();
        describeStreamRequest.setStreamName(streamName);

        // Collect shards from Kinesis
        String exclusiveStartShardId = null;
        List<Shard> shards = new ArrayList<>();
        do {
            describeStreamRequest.setExclusiveStartShardId(exclusiveStartShardId);
            DescribeStreamResult describeStreamResult = clientManager.getClient().describeStream(describeStreamRequest);

            String streamStatus = describeStreamResult.getStreamDescription().getStreamStatus();
            if (!streamStatus.equals("ACTIVE") && !streamStatus.equals("UPDATING")) {
                throw new ResourceNotFoundException("Stream not Active");
            }

            desc.addAllShards(describeStreamResult.getStreamDescription().getShards());

            if (describeStreamResult.getStreamDescription().getHasMoreShards() && (shards.size() > 0)) {
                exclusiveStartShardId = shards.get(shards.size() - 1).getShardId();
            }
            else {
                exclusiveStartShardId = null;
            }
        } while (exclusiveStartShardId != null);

        this.streamMap.put(streamName, desc);
    }

    return desc;
}
 

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

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

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

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

	return fakeRestoredState;
}
 

@Test
public void shouldFindEarliestShardsWhenTrimHorizonStartingPointRequested() throws Exception {
  // given
  StartingPoint trimHorizonStartingPoint =
      new StartingPoint(InitialPositionInStream.TRIM_HORIZON);
  when(kinesis.listShards(STREAM_NAME)).thenReturn(allShards);

  // when
  Iterable<Shard> shardsAtStartingPoint =
      underTest.findShardsAtStartingPoint(kinesis, STREAM_NAME, trimHorizonStartingPoint);

  // then
  assertThat(shardsAtStartingPoint).containsExactlyInAnyOrder(shard00, shard01);
}
 

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

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

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

	return fakeRestoredState;
}
 

@Test
public void shouldFind3StartingShardsInTheMiddle() throws Exception {
  // given
  Instant timestampAfterShards3And4Merge = new Instant();
  StartingPoint startingPointAfterFirstSplitsAndMerge =
      new StartingPoint(timestampAfterShards3And4Merge);

  expiredAtTimestamp(shard00, timestampAfterShards3And4Merge);
  expiredAtTimestamp(shard01, timestampAfterShards3And4Merge);
  activeAtTimestamp(shard02, timestampAfterShards3And4Merge);
  expiredAtTimestamp(shard03, timestampAfterShards3And4Merge);
  expiredAtTimestamp(shard04, timestampAfterShards3And4Merge);
  activeAtTimestamp(shard05, timestampAfterShards3And4Merge);
  activeAtTimestamp(shard06, timestampAfterShards3And4Merge);
  activeAtTimestamp(shard07, timestampAfterShards3And4Merge);
  activeAtTimestamp(shard08, timestampAfterShards3And4Merge);
  activeAtTimestamp(shard09, timestampAfterShards3And4Merge);
  activeAtTimestamp(shard10, timestampAfterShards3And4Merge);

  when(kinesis.listShards(STREAM_NAME)).thenReturn(allShards);

  // when
  Iterable<Shard> shardsAtStartingPoint =
      underTest.findShardsAtStartingPoint(
          kinesis, STREAM_NAME, startingPointAfterFirstSplitsAndMerge);

  // then
  assertThat(shardsAtStartingPoint).containsExactlyInAnyOrder(shard02, shard05, shard06);
}
 

@Override
public DescribeStreamResult describeStream(String streamName, String exclusiveStartShardId) {
  if (rateLimitDescribeStream-- > 0) {
    throw new LimitExceededException("DescribeStream rate limit exceeded");
  }
  int nextShardId = 0;
  if (exclusiveStartShardId != null) {
    nextShardId = parseInt(exclusiveStartShardId) + 1;
  }
  boolean hasMoreShards = nextShardId + 1 < shardedData.size();

  List<Shard> shards = new ArrayList<>();
  if (nextShardId < shardedData.size()) {
    shards.add(new Shard().withShardId(Integer.toString(nextShardId)));
  }

  HttpResponse response = new HttpResponse(null, null);
  response.setStatusCode(200);
  DescribeStreamResult result = new DescribeStreamResult();
  result.setSdkHttpMetadata(SdkHttpMetadata.from(response));
  result.withStreamDescription(
      new StreamDescription()
          .withHasMoreShards(hasMoreShards)
          .withShards(shards)
          .withStreamName(streamName));
  return result;
}
 

@Test
public void testLegacyKinesisStreamShardToStreamShardMetadataConversion() {
	String streamName = "fakeStream1";
	String shardId = "shard-000001";
	String parentShardId = "shard-000002";
	String adjacentParentShardId = "shard-000003";
	String startingHashKey = "key-000001";
	String endingHashKey = "key-000010";
	String startingSequenceNumber = "seq-0000021";
	String endingSequenceNumber = "seq-00000031";

	StreamShardMetadata streamShardMetadata = new StreamShardMetadata();
	streamShardMetadata.setStreamName(streamName);
	streamShardMetadata.setShardId(shardId);
	streamShardMetadata.setParentShardId(parentShardId);
	streamShardMetadata.setAdjacentParentShardId(adjacentParentShardId);
	streamShardMetadata.setStartingHashKey(startingHashKey);
	streamShardMetadata.setEndingHashKey(endingHashKey);
	streamShardMetadata.setStartingSequenceNumber(startingSequenceNumber);
	streamShardMetadata.setEndingSequenceNumber(endingSequenceNumber);

	Shard shard = new Shard()
		.withShardId(shardId)
		.withParentShardId(parentShardId)
		.withAdjacentParentShardId(adjacentParentShardId)
		.withHashKeyRange(new HashKeyRange()
			.withStartingHashKey(startingHashKey)
			.withEndingHashKey(endingHashKey))
		.withSequenceNumberRange(new SequenceNumberRange()
			.withStartingSequenceNumber(startingSequenceNumber)
			.withEndingSequenceNumber(endingSequenceNumber));
	KinesisStreamShard kinesisStreamShard = new KinesisStreamShard(streamName, shard);

	assertEquals(streamShardMetadata, KinesisStreamShard.convertToStreamShardMetadata(kinesisStreamShard));
}
 

public List<Shard> listShards(final String streamName) throws TransientKinesisException {
  return wrapExceptions(
      () -> {
        List<Shard> shards = Lists.newArrayList();
        String lastShardId = null;

        // DescribeStream has limits that can be hit fairly easily if we are attempting
        // to configure multiple KinesisIO inputs in the same account. Retry up to
        // LIST_SHARDS_DESCRIBE_STREAM_MAX_ATTEMPTS times if we end up hitting that limit.
        //
        // Only pass the wrapped exception up once that limit is reached. Use FluentBackoff
        // to implement the retry policy.
        FluentBackoff retryBackoff =
            FluentBackoff.DEFAULT
                .withMaxRetries(LIST_SHARDS_DESCRIBE_STREAM_MAX_ATTEMPTS)
                .withInitialBackoff(LIST_SHARDS_DESCRIBE_STREAM_INITIAL_BACKOFF);
        StreamDescription description = null;
        do {
          BackOff backoff = retryBackoff.backoff();
          Sleeper sleeper = Sleeper.DEFAULT;
          while (true) {
            try {
              description =
                  kinesis.describeStream(streamName, lastShardId).getStreamDescription();
              break;
            } catch (LimitExceededException exc) {
              if (!BackOffUtils.next(sleeper, backoff)) {
                throw exc;
              }
            }
          }

          shards.addAll(description.getShards());
          lastShardId = shards.get(shards.size() - 1).getShardId();
        } while (description.getHasMoreShards());

        return shards;
      });
}
 

protected ArrayList<Shard> getShards(InternalStream theStream, String fromShardId)
{
    ArrayList<Shard> externalList = new ArrayList<Shard>();
    for (InternalShard intshard : theStream.getShardsFrom(fromShardId)) {
        externalList.add(intshard);
    }

    return externalList;
}
 

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

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

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

	return fakeRestoredState;
}
 

@Test
void testProvisionConsumerExistingStreamUpdateShards() {
	AmazonKinesis amazonKinesisMock = mock(AmazonKinesis.class);
	ArgumentCaptor<UpdateShardCountRequest> updateShardCaptor = ArgumentCaptor
			.forClass(UpdateShardCountRequest.class);
	String name = "test-stream";
	String group = "test-group";
	int targetShardCount = 2;
	KinesisBinderConfigurationProperties binderProperties = new KinesisBinderConfigurationProperties();
	binderProperties.setMinShardCount(targetShardCount);
	binderProperties.setAutoAddShards(true);
	KinesisStreamProvisioner provisioner = new KinesisStreamProvisioner(
			amazonKinesisMock, binderProperties);

	ExtendedConsumerProperties<KinesisConsumerProperties> extendedConsumerProperties =
			new ExtendedConsumerProperties<>(
			new KinesisConsumerProperties());

	DescribeStreamResult describeOriginalStream = describeStreamResultWithShards(
			Collections.singletonList(new Shard()));

	DescribeStreamResult describeUpdatedStream = describeStreamResultWithShards(
			Arrays.asList(new Shard(), new Shard()));

	when(amazonKinesisMock.describeStream(any(DescribeStreamRequest.class)))
			.thenReturn(describeOriginalStream).thenReturn(describeUpdatedStream);

	provisioner.provisionConsumerDestination(name, group, extendedConsumerProperties);

	verify(amazonKinesisMock, times(1)).updateShardCount(updateShardCaptor.capture());

	assertThat(updateShardCaptor.getValue().getStreamName()).isEqualTo(name);
	assertThat(updateShardCaptor.getValue().getScalingType())
			.isEqualTo(ScalingType.UNIFORM_SCALING.name());
	assertThat(updateShardCaptor.getValue().getTargetShardCount())
			.isEqualTo(targetShardCount);
}
 

@Test
void testProvisionProducerSuccessfulWithNewStream() {
	AmazonKinesis amazonKinesisMock = mock(AmazonKinesis.class);
	KinesisBinderConfigurationProperties binderProperties = new KinesisBinderConfigurationProperties();
	KinesisStreamProvisioner provisioner = new KinesisStreamProvisioner(
			amazonKinesisMock, binderProperties);
	ExtendedProducerProperties<KinesisProducerProperties> extendedProducerProperties =
			new ExtendedProducerProperties<>(
			new KinesisProducerProperties());

	String name = "test-stream";
	Integer shards = 1;

	DescribeStreamResult describeStreamResult = describeStreamResultWithShards(
			Collections.singletonList(new Shard()));

	when(amazonKinesisMock.describeStream(any(DescribeStreamRequest.class)))
			.thenThrow(new ResourceNotFoundException("I got nothing"))
			.thenReturn(describeStreamResult);

	when(amazonKinesisMock.createStream(name, shards))
			.thenReturn(new CreateStreamResult());

	ProducerDestination destination = provisioner.provisionProducerDestination(name,
			extendedProducerProperties);

	verify(amazonKinesisMock, times(2))
			.describeStream(any(DescribeStreamRequest.class));

	verify(amazonKinesisMock).createStream(name, shards);

	assertThat(destination.getName()).isEqualTo(name);
}
 

@Test
public void testStreamShardMetadataAndHandleConversion() {
	String streamName = "fakeStream1";
	String shardId = "shard-000001";
	String parentShardId = "shard-000002";
	String adjacentParentShardId = "shard-000003";
	String startingHashKey = "key-000001";
	String endingHashKey = "key-000010";
	String startingSequenceNumber = "seq-0000021";
	String endingSequenceNumber = "seq-00000031";

	StreamShardMetadata kinesisStreamShard = new StreamShardMetadata();
	kinesisStreamShard.setStreamName(streamName);
	kinesisStreamShard.setShardId(shardId);
	kinesisStreamShard.setParentShardId(parentShardId);
	kinesisStreamShard.setAdjacentParentShardId(adjacentParentShardId);
	kinesisStreamShard.setStartingHashKey(startingHashKey);
	kinesisStreamShard.setEndingHashKey(endingHashKey);
	kinesisStreamShard.setStartingSequenceNumber(startingSequenceNumber);
	kinesisStreamShard.setEndingSequenceNumber(endingSequenceNumber);

	Shard shard = new Shard()
		.withShardId(shardId)
		.withParentShardId(parentShardId)
		.withAdjacentParentShardId(adjacentParentShardId)
		.withHashKeyRange(new HashKeyRange()
			.withStartingHashKey(startingHashKey)
			.withEndingHashKey(endingHashKey))
		.withSequenceNumberRange(new SequenceNumberRange()
			.withStartingSequenceNumber(startingSequenceNumber)
			.withEndingSequenceNumber(endingSequenceNumber));
	StreamShardHandle streamShardHandle = new StreamShardHandle(streamName, shard);

	assertEquals(kinesisStreamShard, KinesisDataFetcher.convertToStreamShardMetadata(streamShardHandle));
	assertEquals(streamShardHandle, KinesisDataFetcher.convertToStreamShardHandle(kinesisStreamShard));
}
 

private List<Shard> describeStream(String stream) {
	String exclusiveStartShardId = null;

	DescribeStreamRequest describeStreamRequest = new DescribeStreamRequest()
			.withStreamName(stream);

	List<Shard> shardList = new ArrayList<>();

	while (true) {
		DescribeStreamResult describeStreamResult;

		describeStreamRequest.withExclusiveStartShardId(exclusiveStartShardId);
		describeStreamResult = AMAZON_KINESIS.describeStream(describeStreamRequest);
		StreamDescription streamDescription = describeStreamResult
				.getStreamDescription();
		if (StreamStatus.ACTIVE.toString()
				.equals(streamDescription.getStreamStatus())) {
			shardList.addAll(streamDescription.getShards());

			if (streamDescription.getHasMoreShards()) {
				exclusiveStartShardId = shardList.get(shardList.size() - 1)
						.getShardId();
				continue;
			}
			else {
				return shardList;
			}
		}
		try {
			Thread.sleep(100);
		}
		catch (InterruptedException ex) {
			Thread.currentThread().interrupt();
			throw new IllegalStateException(ex);
		}
	}
}
 

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