software.amazon.awssdk.services.kinesis.model.ShardIteratorType Java Examples

static WrappedRequest<GetShardIteratorRequest> wrapped(GetShardIteratorRequest.Builder builder) {
    GetShardIteratorRequest req = builder.build();
    return new WrappedRequest<GetShardIteratorRequest>() {
        @Override
        public ShardIteratorType shardIteratorType() {
            return req.shardIteratorType();
        }

        @Override
        public String sequenceNumber() {
            return req.startingSequenceNumber();
        }

        @Override
        public Instant timestamp() {
            return req.timestamp();
        }

        @Override
        public GetShardIteratorRequest request() {
            return req;
        }
    };
}
 

@Test
public void testNonNullGetRecords() throws Exception {
    final String nextIterator = "TestIterator";
    final ArgumentCaptor<GetShardIteratorRequest> iteratorCaptor = ArgumentCaptor
            .forClass(GetShardIteratorRequest.class);
    final ArgumentCaptor<GetRecordsRequest> recordsCaptor = ArgumentCaptor.forClass(GetRecordsRequest.class);
    final GetShardIteratorRequest expectedIteratorRequest = makeGetShardIteratorRequest(
            ShardIteratorType.LATEST.name());
    final GetRecordsRequest expectedRecordsRequest = makeGetRecordsRequest(nextIterator);

    final CompletableFuture<GetRecordsResponse> future = mock(CompletableFuture.class);

    when(kinesisClient.getShardIterator(iteratorCaptor.capture()))
            .thenReturn(makeGetShardIteratorResonse(nextIterator));
    when(kinesisClient.getRecords(recordsCaptor.capture())).thenReturn(future);
    when(future.get(anyLong(), any(TimeUnit.class))).thenThrow(
            new ExecutionException(ResourceNotFoundException.builder().message("Test Exception").build()));

    kinesisDataFetcher.initialize(SentinelCheckpoint.LATEST.toString(), INITIAL_POSITION_LATEST);
    DataFetcherResult dataFetcherResult = kinesisDataFetcher.getRecords();

    assertNotNull(dataFetcherResult);
    assertEquals(expectedIteratorRequest.startingSequenceNumber(), iteratorCaptor.getValue().startingSequenceNumber());
    assertEquals(expectedRecordsRequest.shardIterator(), recordsCaptor.getValue().shardIterator());
}
 

static WrappedRequest<SubscribeToShardRequest> wrapped(SubscribeToShardRequest.Builder builder) {
    SubscribeToShardRequest req = builder.build();
    return new WrappedRequest<SubscribeToShardRequest>() {
        @Override
        public ShardIteratorType shardIteratorType() {
            return req.startingPosition().type();
        }

        @Override
        public String sequenceNumber() {
            return req.startingPosition().sequenceNumber();
        }

        @Override
        public Instant timestamp() {
            return req.startingPosition().timestamp();
        }

        @Override
        public SubscribeToShardRequest request() {
            return req;
        }
    };
}
 

private ShardPosition sanitizePositionedShardPosition(ShardPosition shardPosition) {
    try {
        StartFrom startFrom = shardPosition.startFrom();

        if (startFrom == StartFrom.AT_POSITION || startFrom == StartFrom.POSITION) {
            ShardIteratorType type = startFrom == StartFrom.POSITION
                    ? ShardIteratorType.AFTER_SEQUENCE_NUMBER
                    : ShardIteratorType.AT_SEQUENCE_NUMBER;
            kinesisClient.getShardIterator(GetShardIteratorRequest.builder()
                    .shardId(shardName)
                    .streamName(channelName)
                    .shardIteratorType(type)
                    .startingSequenceNumber(shardPosition.position())
                    .build())
                    .get();
        }

        return shardPosition;
    } catch (ExecutionException | InterruptedException | RuntimeException e) {
        LOG.warn(marker, "given shardposition {} / {} not accessible, falling back to horizon", shardPosition.shardName(), shardPosition.position());
    }
    return ShardPosition.fromHorizon(shardPosition.shardName());
}
 

public static void main(String[] args) {

        // snippet-start:[kinesis.java2.stream_example.setup]
        Region region = Region.US_EAST_1;
        KinesisAsyncClient client = KinesisAsyncClient.builder()
        .region(region)
        .build();

        SubscribeToShardRequest request = SubscribeToShardRequest.builder()
                .consumerARN(CONSUMER_ARN)
                .shardId("arn:aws:kinesis:us-east-1:814548047983:stream/StockTradeStream")
                .startingPosition(s -> s.type(ShardIteratorType.LATEST)).build();

        // snippet-end:[kinesis.java2.stream_example.setup]
        SubscribeToShardResponseHandler responseHandler = SubscribeToShardResponseHandler
                .builder()
                .onError(t -> System.err.println("Error during stream - " + t.getMessage()))
                .subscriber(MySubscriber::new)
                .build();

        client.subscribeToShard(request, responseHandler);
        client.close();
    }
 

private static <R> R apply(R initial, UpdatingFunction<ShardIteratorType, R> shardIterFunc,
        UpdatingFunction<Instant, R> dateFunc, UpdatingFunction<String, R> sequenceFunction,
        InitialPositionInStreamExtended initialPositionInStreamExtended, String sequenceNumber,
        ShardIteratorType defaultIteratorType) {
    ShardIteratorType iteratorType = SHARD_ITERATOR_MAPPING.getOrDefault(
            sequenceNumber, defaultIteratorType);
    R result = shardIterFunc.apply(initial, iteratorType);
    switch (iteratorType) {
    case AT_TIMESTAMP:
        return dateFunc.apply(result, initialPositionInStreamExtended.getTimestamp().toInstant());
    case AT_SEQUENCE_NUMBER:
    case AFTER_SEQUENCE_NUMBER:
        return sequenceFunction.apply(result, sequenceNumber);
    default:
        return result;
    }
}
 

private void testInitializeAndFetch(final String iteratorType, final String seqNo,
        final InitialPositionInStreamExtended initialPositionInStream) throws Exception {
    final ArgumentCaptor<GetShardIteratorRequest> iteratorCaptor = ArgumentCaptor
            .forClass(GetShardIteratorRequest.class);
    final ArgumentCaptor<GetRecordsRequest> recordsCaptor = ArgumentCaptor.forClass(GetRecordsRequest.class);
    final String iterator = "foo";
    final List<Record> expectedRecords = Collections.emptyList();
    GetShardIteratorRequest expectedIteratorRequest = makeGetShardIteratorRequest(iteratorType);
    if (iteratorType.equals(ShardIteratorType.AT_TIMESTAMP.toString())) {
        expectedIteratorRequest = expectedIteratorRequest.toBuilder()
                .timestamp(initialPositionInStream.getTimestamp().toInstant()).build();
    } else if (iteratorType.equals(ShardIteratorType.AT_SEQUENCE_NUMBER.toString())) {
        expectedIteratorRequest = expectedIteratorRequest.toBuilder().startingSequenceNumber(seqNo).build();
    }
    final GetRecordsRequest expectedRecordsRequest = makeGetRecordsRequest(iterator);

    when(kinesisClient.getShardIterator(iteratorCaptor.capture()))
            .thenReturn(makeGetShardIteratorResonse(iterator));

    when(kinesisClient.getRecords(recordsCaptor.capture()))
            .thenReturn(makeGetRecordsResponse(null, expectedRecords));

    Checkpointer checkpoint = mock(Checkpointer.class);
    when(checkpoint.getCheckpoint(SHARD_ID)).thenReturn(new ExtendedSequenceNumber(seqNo));

    final GetRecordsRetrievalStrategy getRecordsRetrievalStrategy = new SynchronousGetRecordsRetrievalStrategy(
            kinesisDataFetcher);
    kinesisDataFetcher.initialize(seqNo, initialPositionInStream);

    assertEquals(expectedRecords, getRecordsRetrievalStrategy.getRecords(MAX_RECORDS).records());
    verify(kinesisClient, times(1)).getShardIterator(any(GetShardIteratorRequest.class));
    verify(kinesisClient, times(1)).getRecords(any(GetRecordsRequest.class));
}
 

@Test
public void testGetRecordsWithResourceNotFoundException() throws Exception {
    final ArgumentCaptor<GetShardIteratorRequest> iteratorCaptor = ArgumentCaptor
            .forClass(GetShardIteratorRequest.class);
    final ArgumentCaptor<GetRecordsRequest> recordsCaptor = ArgumentCaptor.forClass(GetRecordsRequest.class);
    // Set up arguments used by proxy
    final String nextIterator = "TestShardIterator";

    final GetShardIteratorRequest expectedIteratorRequest = makeGetShardIteratorRequest(
            ShardIteratorType.LATEST.name());
    final GetRecordsRequest expectedRecordsRequest = makeGetRecordsRequest(nextIterator);

    final CompletableFuture<GetRecordsResponse> future = mock(CompletableFuture.class);

    // Set up proxy mock methods
    when(kinesisClient.getShardIterator(iteratorCaptor.capture()))
            .thenReturn(makeGetShardIteratorResonse(nextIterator));
    when(kinesisClient.getRecords(recordsCaptor.capture())).thenReturn(future);
    when(future.get(anyLong(), any(TimeUnit.class))).thenThrow(
            new ExecutionException(ResourceNotFoundException.builder().message("Test Exception").build()));

    // Create data fectcher and initialize it with latest type checkpoint
    kinesisDataFetcher.initialize(SentinelCheckpoint.LATEST.toString(), INITIAL_POSITION_LATEST);
    final GetRecordsRetrievalStrategy getRecordsRetrievalStrategy = new SynchronousGetRecordsRetrievalStrategy(
            kinesisDataFetcher);
    try {
        // Call records of dataFetcher which will throw an exception
        getRecordsRetrievalStrategy.getRecords(MAX_RECORDS);
    } finally {
        // Test shard has reached the end
        assertTrue("Shard should reach the end", kinesisDataFetcher.isShardEndReached());
        assertEquals(expectedIteratorRequest.startingSequenceNumber(), iteratorCaptor.getValue().startingSequenceNumber());
        assertEquals(expectedRecordsRequest.shardIterator(), recordsCaptor.getValue().shardIterator());
    }
}
 

@Test
public void testadvanceIteratorToTrimHorizonLatestAndAtTimestamp() throws InterruptedException, ExecutionException {
    final ArgumentCaptor<GetShardIteratorRequest> requestCaptor = ArgumentCaptor
            .forClass(GetShardIteratorRequest.class);
    final String iteratorHorizon = "TRIM_HORIZON";
    final String iteratorLatest = "LATEST";
    final String iteratorAtTimestamp = "AT_TIMESTAMP";
    final Map<ShardIteratorType, GetShardIteratorRequest> requestsMap = Arrays
            .stream(new String[] { iteratorHorizon, iteratorLatest, iteratorAtTimestamp })
            .map(this::makeGetShardIteratorRequest)
            .collect(Collectors.toMap(r -> ShardIteratorType.valueOf(r.shardIteratorTypeAsString()), r -> r));
    GetShardIteratorRequest tsReq = requestsMap.get(ShardIteratorType.AT_TIMESTAMP);
    requestsMap.put(ShardIteratorType.AT_TIMESTAMP,
            tsReq.toBuilder().timestamp(INITIAL_POSITION_AT_TIMESTAMP.getTimestamp().toInstant()).build());

    when(kinesisClient.getShardIterator(requestCaptor.capture()))
            .thenReturn(makeGetShardIteratorResonse(iteratorHorizon))
            .thenReturn(makeGetShardIteratorResonse(iteratorLatest))
            .thenReturn(makeGetShardIteratorResonse(iteratorAtTimestamp));

    kinesisDataFetcher.advanceIteratorTo(ShardIteratorType.TRIM_HORIZON.toString(), INITIAL_POSITION_TRIM_HORIZON);
    assertEquals(iteratorHorizon, kinesisDataFetcher.getNextIterator());

    kinesisDataFetcher.advanceIteratorTo(ShardIteratorType.LATEST.toString(), INITIAL_POSITION_LATEST);
    assertEquals(iteratorLatest, kinesisDataFetcher.getNextIterator());

    kinesisDataFetcher.advanceIteratorTo(ShardIteratorType.AT_TIMESTAMP.toString(), INITIAL_POSITION_AT_TIMESTAMP);
    assertEquals(iteratorAtTimestamp, kinesisDataFetcher.getNextIterator());

    final List<GetShardIteratorRequest> requests = requestCaptor.getAllValues();
    assertEquals(3, requests.size());
    requests.forEach(request -> {
        final ShardIteratorType type = ShardIteratorType.fromValue(request.shardIteratorTypeAsString());
        assertEquals(requestsMap.get(type).startingSequenceNumber(), request.startingSequenceNumber());
        requestsMap.remove(type);
    });
    assertEquals(0, requestsMap.size());
}
 

@Test
public void testadvanceIteratorTo() throws KinesisClientLibException, InterruptedException, ExecutionException {
    final Checkpointer checkpoint = mock(Checkpointer.class);
    final String iteratorA = "foo";
    final String iteratorB = "bar";
    final String seqA = "123";
    final String seqB = "456";

    ArgumentCaptor<GetShardIteratorRequest> shardIteratorRequestCaptor = ArgumentCaptor
            .forClass(GetShardIteratorRequest.class);

    when(kinesisClient.getShardIterator(shardIteratorRequestCaptor.capture()))
            .thenReturn(makeGetShardIteratorResonse(iteratorA)).thenReturn(makeGetShardIteratorResonse(iteratorA))
            .thenReturn(makeGetShardIteratorResonse(iteratorB));
    when(checkpoint.getCheckpoint(SHARD_ID)).thenReturn(new ExtendedSequenceNumber(seqA));

    kinesisDataFetcher.initialize(seqA, null);
    kinesisDataFetcher.advanceIteratorTo(seqA, null);
    kinesisDataFetcher.advanceIteratorTo(seqB, null);

    final List<GetShardIteratorRequest> shardIteratorRequests = shardIteratorRequestCaptor.getAllValues();
    assertEquals(3, shardIteratorRequests.size());

    int count = 0;
    for (GetShardIteratorRequest request : shardIteratorRequests) {
        assertEquals(STREAM_NAME, request.streamName());
        assertEquals(SHARD_ID, request.shardId());
        assertEquals(ShardIteratorType.AT_SEQUENCE_NUMBER.toString(), request.shardIteratorTypeAsString());
        if (count == 2) {
            assertEquals(seqB, request.startingSequenceNumber());
        } else {
            assertEquals(seqA, request.startingSequenceNumber());
        }
        count++;
    }
}
 

private <T, R> void latestTest(Supplier<T> supplier, Consumer<R> baseVerifier, IteratorApply<T> iteratorRequest,
        Function<T, WrappedRequest<R>> toRequest) {
    String sequenceNumber = SentinelCheckpoint.LATEST.name();
    InitialPositionInStreamExtended initialPosition = InitialPositionInStreamExtended
            .newInitialPosition(InitialPositionInStream.LATEST);
    updateTest(supplier, baseVerifier, iteratorRequest, toRequest, sequenceNumber, initialPosition,
            ShardIteratorType.LATEST, null, null);
}
 

private <T, R> void trimHorizonTest(Supplier<T> supplier, Consumer<R> baseVerifier,
        IteratorApply<T> iteratorRequest, Function<T, WrappedRequest<R>> toRequest) {
    String sequenceNumber = SentinelCheckpoint.TRIM_HORIZON.name();
    InitialPositionInStreamExtended initialPosition = InitialPositionInStreamExtended
            .newInitialPosition(InitialPositionInStream.TRIM_HORIZON);
    updateTest(supplier, baseVerifier, iteratorRequest, toRequest, sequenceNumber, initialPosition,
            ShardIteratorType.TRIM_HORIZON, null, null);
}
 

private <T, R> void sequenceNumber(Supplier<T> supplier, Consumer<R> baseVerifier, IteratorApply<T> iteratorRequest,
        Function<T, WrappedRequest<R>> toRequest, ShardIteratorType shardIteratorType) {
    InitialPositionInStreamExtended initialPosition = InitialPositionInStreamExtended
            .newInitialPosition(InitialPositionInStream.TRIM_HORIZON);
    updateTest(supplier, baseVerifier, iteratorRequest, toRequest, SEQUENCE_NUMBER, initialPosition,
            shardIteratorType, "1234", null);
}
 

private <T, R> void timeStampTest(Supplier<T> supplier, Consumer<R> baseVerifier, IteratorApply<T> iteratorRequest,
        Function<T, WrappedRequest<R>> toRequest) {
    String sequenceNumber = SentinelCheckpoint.AT_TIMESTAMP.name();
    InitialPositionInStreamExtended initialPosition = InitialPositionInStreamExtended
            .newInitialPositionAtTimestamp(new Date(TIMESTAMP.toEpochMilli()));
    updateTest(supplier, baseVerifier, iteratorRequest, toRequest, sequenceNumber, initialPosition,
            ShardIteratorType.AT_TIMESTAMP, null, TIMESTAMP);
}
 

private <T, R> void updateTest(Supplier<T> supplier, Consumer<R> baseVerifier, IteratorApply<T> iteratorRequest,
        Function<T, WrappedRequest<R>> toRequest, String sequenceNumber,
        InitialPositionInStreamExtended initialPositionInStream, ShardIteratorType expectedShardIteratorType,
        String expectedSequenceNumber, Instant expectedTimestamp) {
    T base = supplier.get();
    T updated = iteratorRequest.apply(base, sequenceNumber, initialPositionInStream);
    WrappedRequest<R> request = toRequest.apply(updated);
    baseVerifier.accept(request.request());
    assertThat(request.shardIteratorType(), equalTo(expectedShardIteratorType));
    assertThat(request.sequenceNumber(), equalTo(expectedSequenceNumber));
    assertThat(request.timestamp(), equalTo(expectedTimestamp));

}
 

public static void main(String[] args) {

        KinesisAsyncClient client = KinesisAsyncClient.create();

        SubscribeToShardRequest request = SubscribeToShardRequest.builder()
                .consumerARN(CONSUMER_ARN)
                .shardId("shardId-000000000000")
                .startingPosition(StartingPosition.builder().type(ShardIteratorType.LATEST).build())
                .build();

        responseHandlerBuilder_RxJava(client, request).join();

        client.close();
    }
 

public static void main(String[] args) {

        KinesisAsyncClient client = KinesisAsyncClient.create();

        SubscribeToShardRequest request = SubscribeToShardRequest.builder()
                .consumerARN(CONSUMER_ARN)
                .shardId("shardId-000000000000")
                .startingPosition(StartingPosition.builder().type(ShardIteratorType.LATEST).build())
                .build();

        responseHandlerBuilder_Reactor(client, request).join();

        client.close();
    }
 

/**
 * Generate request per consumer/shard combination
 * @return a lit of completablefutures
 */
private List<CompletableFuture<?>> generateSubscribeToShardFutures() {
    List<CompletableFuture<?>> completableFutures = new ArrayList<>();
    for (int i = 0; i < CONSUMER_COUNT; i++) {
        final int consumerIndex = i;
        for (int j = 0; j < SHARD_COUNT; j++) {
            final int shardIndex = j;
            TestSubscribeToShardResponseHandler responseHandler =
                new TestSubscribeToShardResponseHandler(consumerIndex, shardIndex);
            CompletableFuture<Void> completableFuture =
                asyncClient.subscribeToShard(b -> b.shardId(shardIds.get(shardIndex))
                                                   .consumerARN(consumerArns.get(consumerIndex))
                                                   .startingPosition(s -> s.type(ShardIteratorType.TRIM_HORIZON)),
                                             responseHandler)
                           .thenAccept(b -> {
                               // Only verify data if all events have been received and the received data is not empty.
                               // It is possible the received data is empty because there is no record at the position
                               // event with TRIM_HORIZON.
                               if (responseHandler.allEventsReceived && !responseHandler.receivedData.isEmpty()) {
                                   assertThat(producedData).as(responseHandler.id + " has not received all events"
                                                               + ".").containsSequence(responseHandler.receivedData);
                               }
                           });
            completableFutures.add(completableFuture);
        }
    }
    return completableFutures;
}
 

@Test
public void testGetIteratorForBogusStream() {
    try {
        client.getShardIterator(GetShardIteratorRequest.builder()
                                                       .streamName("bogus-stream-name")
                                                       .shardId("bogus-shard-id")
                                                       .shardIteratorType(ShardIteratorType.LATEST)
                                                       .build());
        Assert.fail("Expected ResourceNotFoundException");
    } catch (ResourceNotFoundException exception) {
        // Ignored or expected.
    }
}
 

@Test
public void subscribeToShard_ReceivesAllData() {
    List<SdkBytes> producedData = new ArrayList<>();
    ScheduledExecutorService producer = Executors.newScheduledThreadPool(1);
    // Delay it a bit to allow us to subscribe first
    producer.scheduleAtFixedRate(() -> putRecord().ifPresent(producedData::add), 10, 1, TimeUnit.SECONDS);

    List<SdkBytes> receivedData = new ArrayList<>();
    // Add every event's data to the receivedData list
    Consumer<SubscribeToShardEvent> eventConsumer = s -> receivedData.addAll(
        s.records().stream()
         .map(Record::data)
         .collect(Collectors.toList()));
    asyncClient.subscribeToShard(r -> r.consumerARN(consumerArn)
                                       .shardId(shardId)
                                       .startingPosition(s -> s.type(ShardIteratorType.LATEST)),
                                 SubscribeToShardResponseHandler.builder()
                                                                .onEventStream(p -> p.filter(SubscribeToShardEvent.class)
                                                                                     .subscribe(eventConsumer))
                                                                .onResponse(this::verifyHttpMetadata)
                                                                .build())
               .join();
    producer.shutdown();
    // Make sure we all the data we received was data we published, we may have published more
    // if the producer isn't shutdown immediately after we finish subscribing.
    assertThat(producedData).containsSequence(receivedData);
}
 

@Test
public void subscribeToShard_smallWindow_doesNotTimeOutReads() {
    // We want sufficiently large records (relative to the initial window
    // size we're choosing) so the client has to send multiple
    // WINDOW_UPDATEs to receive them
    for (int i = 0; i < 16; ++i) {
        putRecord(64 * 1024);
    }

    KinesisAsyncClient smallWindowAsyncClient = KinesisAsyncClient.builder()
            .credentialsProvider(CREDENTIALS_PROVIDER_CHAIN)
            .httpClientBuilder(NettyNioAsyncHttpClient.builder()
                .http2Configuration(Http2Configuration.builder()
                        .initialWindowSize(16384)
                        .build()))
            .build();

    try {
        smallWindowAsyncClient.subscribeToShard(r -> r.consumerARN(consumerArn)
                        .shardId(shardId)
                        .startingPosition(s -> s.type(ShardIteratorType.TRIM_HORIZON)),
                SubscribeToShardResponseHandler.builder()
                        .onEventStream(es -> Flowable.fromPublisher(es).forEach(e -> {}))
                        .onResponse(this::verifyHttpMetadata)
                        .build())
                .join();

    } finally {
        smallWindowAsyncClient.close();
    }
}
 

private <T, R> void sequenceNumber(Supplier<T> supplier, Consumer<R> baseVerifier, IteratorApply<T> iteratorRequest,
        Function<T, WrappedRequest<R>> toRequest) {
    sequenceNumber(supplier, baseVerifier, iteratorRequest, toRequest, ShardIteratorType.AT_SEQUENCE_NUMBER);
}
 

@Test
public void subscribeReconnectTest() {
    sequenceNumber(this::stsBase, this::verifyStsBase, IteratorBuilder::reconnectRequest, WrappedRequest::wrapped,
            ShardIteratorType.AFTER_SEQUENCE_NUMBER);
}
 

/**
 * Test initialize() with the TIME_ZERO iterator instruction
 */
@Test
public final void testInitializeTimeZero() throws Exception {
    testInitializeAndFetch(ShardIteratorType.TRIM_HORIZON.toString(), ShardIteratorType.TRIM_HORIZON.toString(),
            INITIAL_POSITION_TRIM_HORIZON);
}
 

/**
 * Test initialize() with the AT_TIMESTAMP iterator instruction
 */
@Test
public final void testInitializeAtTimestamp() throws Exception {
    testInitializeAndFetch(ShardIteratorType.AT_TIMESTAMP.toString(), ShardIteratorType.AT_TIMESTAMP.toString(),
            INITIAL_POSITION_AT_TIMESTAMP);
}
 

public ShardIteratorType getIteratorType() {
    return iteratorType;
}
 

/**
 * Test initialize() with the LATEST iterator instruction
 */
@Test
public final void testInitializeLatest() throws Exception {
    testInitializeAndFetch(ShardIteratorType.LATEST.toString(), ShardIteratorType.LATEST.toString(),
            INITIAL_POSITION_LATEST);
}
 

@Test
public void testContinuesAfterSequence() {
    FanOutRecordsPublisher source = new FanOutRecordsPublisher(kinesisClient, SHARD_ID, CONSUMER_ARN);

    ArgumentCaptor<FanOutRecordsPublisher.RecordSubscription> captor = ArgumentCaptor
            .forClass(FanOutRecordsPublisher.RecordSubscription.class);
    ArgumentCaptor<FanOutRecordsPublisher.RecordFlow> flowCaptor = ArgumentCaptor
            .forClass(FanOutRecordsPublisher.RecordFlow.class);

    doNothing().when(publisher).subscribe(captor.capture());

    source.start(new ExtendedSequenceNumber("0"),
            InitialPositionInStreamExtended.newInitialPosition(InitialPositionInStream.LATEST));

    NonFailingSubscriber nonFailingSubscriber = new NonFailingSubscriber();

    source.subscribe(new ShardConsumerNotifyingSubscriber(nonFailingSubscriber, source));

    SubscribeToShardRequest expected = SubscribeToShardRequest.builder().consumerARN(CONSUMER_ARN).shardId(SHARD_ID)
            .startingPosition(StartingPosition.builder().sequenceNumber("0")
                    .type(ShardIteratorType.AT_SEQUENCE_NUMBER).build())
            .build();

    verify(kinesisClient).subscribeToShard(argThat(new SubscribeToShardRequestMatcher(expected)), flowCaptor.capture());

    flowCaptor.getValue().onEventStream(publisher);
    captor.getValue().onSubscribe(subscription);

    List<Record> records = Stream.of(1, 2, 3).map(this::makeRecord).collect(Collectors.toList());
    List<KinesisClientRecordMatcher> matchers = records.stream().map(KinesisClientRecordMatcher::new)
            .collect(Collectors.toList());

    batchEvent = SubscribeToShardEvent.builder().millisBehindLatest(100L).records(records)
            .continuationSequenceNumber("3").build();

    captor.getValue().onNext(batchEvent);
    captor.getValue().onComplete();
    flowCaptor.getValue().complete();

    ArgumentCaptor<FanOutRecordsPublisher.RecordSubscription> nextSubscribeCaptor = ArgumentCaptor
            .forClass(FanOutRecordsPublisher.RecordSubscription.class);
    ArgumentCaptor<FanOutRecordsPublisher.RecordFlow> nextFlowCaptor = ArgumentCaptor
            .forClass(FanOutRecordsPublisher.RecordFlow.class);

    SubscribeToShardRequest nextExpected = SubscribeToShardRequest.builder().consumerARN(CONSUMER_ARN)
            .shardId(SHARD_ID).startingPosition(StartingPosition.builder().sequenceNumber("3")
                    .type(ShardIteratorType.AFTER_SEQUENCE_NUMBER).build())
            .build();

    verify(kinesisClient).subscribeToShard(argThat(new SubscribeToShardRequestMatcher(nextExpected)), nextFlowCaptor.capture());
    reset(publisher);
    doNothing().when(publisher).subscribe(nextSubscribeCaptor.capture());

    nextFlowCaptor.getValue().onEventStream(publisher);
    nextSubscribeCaptor.getValue().onSubscribe(subscription);

    List<Record> nextRecords = Stream.of(4, 5, 6).map(this::makeRecord).collect(Collectors.toList());
    List<KinesisClientRecordMatcher> nextMatchers = nextRecords.stream().map(KinesisClientRecordMatcher::new)
            .collect(Collectors.toList());

    batchEvent = SubscribeToShardEvent.builder().millisBehindLatest(100L).records(nextRecords)
            .continuationSequenceNumber("6").build();
    nextSubscribeCaptor.getValue().onNext(batchEvent);

    verify(subscription, times(4)).request(1);

    assertThat(nonFailingSubscriber.received.size(), equalTo(2));

    verifyRecords(nonFailingSubscriber.received.get(0).records(), matchers);
    verifyRecords(nonFailingSubscriber.received.get(1).records(), nextMatchers);

}
 

public static GetShardIteratorRequest.Builder request(GetShardIteratorRequest.Builder builder,
        String sequenceNumber, InitialPositionInStreamExtended initialPosition) {
    return apply(builder, GetShardIteratorRequest.Builder::shardIteratorType, GetShardIteratorRequest.Builder::timestamp,
            GetShardIteratorRequest.Builder::startingSequenceNumber, initialPosition, sequenceNumber,
            ShardIteratorType.AT_SEQUENCE_NUMBER);
}
 

public static StartingPosition.Builder reconnectRequest(StartingPosition.Builder builder, String sequenceNumber,
        InitialPositionInStreamExtended initialPosition) {
    return apply(builder, StartingPosition.Builder::type, StartingPosition.Builder::timestamp,
            StartingPosition.Builder::sequenceNumber, initialPosition, sequenceNumber,
            ShardIteratorType.AFTER_SEQUENCE_NUMBER);
}