List<Record> retrieveRecords(String kinesisStream, KinesisClient kinesisClient) {

        LOG.debug("About to retrieve logs records from Kinesis.");
        // Create ListShard request and response and designate the Kinesis stream
        final ListShardsRequest listShardsRequest = ListShardsRequest.builder().streamName(kinesisStream).build();
        final ListShardsResponse listShardsResponse = kinesisClient.listShards(listShardsRequest);
        final List<Record> recordsList = new ArrayList<>();

        // Iterate through the shards that exist
        for (Shard shard : listShardsResponse.shards()) {
            final String shardId = shard.shardId();
            final GetShardIteratorRequest getShardIteratorRequest =
                    GetShardIteratorRequest.builder()
                            .shardId(shardId)
                            .streamName(kinesisStream)
                            .shardIteratorType(ShardIteratorType.TRIM_HORIZON)
                            .build();
            String shardIterator = kinesisClient.getShardIterator(getShardIteratorRequest).shardIterator();
            boolean stayOnCurrentShard = true;
            LOG.debug("Retrieved shard id: [{}] with shard iterator: [{}]", shardId, shardIterator);
            while (stayOnCurrentShard) {
                LOG.debug("Getting more records");
                final GetRecordsRequest getRecordsRequest = GetRecordsRequest.builder().shardIterator(shardIterator).build();
                final GetRecordsResponse getRecordsResponse = kinesisClient.getRecords(getRecordsRequest);
                shardIterator = getRecordsResponse.nextShardIterator();

                for (Record record : getRecordsResponse.records()) {
                    if (isControlMessage(record)) {
                        continue;
                    }
                    recordsList.add(record);

                    if (recordsList.size() == RECORDS_SAMPLE_SIZE) {
                        LOG.debug("Returning the list of records now that sample size [{}] has been met.", RECORDS_SAMPLE_SIZE);
                        return recordsList;
                    }
                }

                if (getRecordsResponse.millisBehindLatest() == 0) {
                    LOG.debug("Found the end of the shard. No more records returned from the shard.");
                    stayOnCurrentShard = false;
                }
            }
        }
        LOG.debug("Returning the list with [{}] records.", recordsList.size());
        return recordsList;
    }
@Test
    public void testRetrieveRecords() throws IOException {

        Shard shard = Shard.builder().shardId("shardId-1234").build();
        when(kinesisClient.listShards(isA(ListShardsRequest.class)))
                .thenReturn(ListShardsResponse.builder().shards(shard).build());

        when(kinesisClient.getShardIterator(isA(GetShardIteratorRequest.class)))
                .thenReturn(GetShardIteratorResponse.builder().shardIterator("shardIterator").build());

        final Record record = Record.builder()
                .approximateArrivalTimestamp(Instant.now())
                .data(SdkBytes.fromByteArray(AWSTestingUtils.cloudWatchRawPayload()))
                .build();
        GetRecordsResponse recordsResponse = GetRecordsResponse.builder().records(record).millisBehindLatest(10000L).build();
        when(kinesisClient.getRecords(isA(GetRecordsRequest.class)))
                .thenReturn(recordsResponse)
                .thenReturn(recordsResponse)
                .thenReturn(recordsResponse)
                .thenReturn(recordsResponse)
                .thenReturn(recordsResponse)
                .thenReturn(recordsResponse)
                .thenReturn(recordsResponse)
                .thenReturn(recordsResponse)
                .thenReturn(recordsResponse)
                .thenReturn(recordsResponse)
                .thenReturn(recordsResponse)
                .thenReturn(recordsResponse);

        List<Record> fakeRecordsList = kinesisService.retrieveRecords("kinesisStream", kinesisClient);
        assertEquals(fakeRecordsList.size(), 10);
    }