@Override
  public synchronized KinesisMessageBatch fetchMessages(StreamPartitionMsgOffset startMsgOffset, int timeoutMs) {
    try {
      return getKinesisMessageBatch((KinesisPartitionGroupOffset) startMsgOffset);
    } catch (ProvisionedThroughputExceededException pte) {
      LOGGER.error("Rate limit exceeded while fetching messages from Kinesis stream: {} with threshold: {}",
          pte.getMessage(), _config.getRpsLimit());
      return new KinesisMessageBatch(List.of(), (KinesisPartitionGroupOffset) startMsgOffset, false);
    }
  }
@Test
  public void testBasicConsumer() {
    KinesisClient kinesisClient = mock(KinesisClient.class);
    when(kinesisClient.getShardIterator(any(GetShardIteratorRequest.class))).thenReturn(
        GetShardIteratorResponse.builder().shardIterator(PLACEHOLDER).build());
    when(kinesisClient.getRecords(any(GetRecordsRequest.class))).thenReturn(
        GetRecordsResponse.builder().nextShardIterator(PLACEHOLDER).records(_records).build());

    KinesisConsumer kinesisConsumer = new KinesisConsumer(_kinesisConfig, kinesisClient);

    // Fetch first batch
    KinesisPartitionGroupOffset startOffset = new KinesisPartitionGroupOffset("0", "1");
    KinesisMessageBatch kinesisMessageBatch = kinesisConsumer.fetchMessages(startOffset, TIMEOUT);
    assertEquals(kinesisMessageBatch.getMessageCount(), NUM_RECORDS);
    for (int i = 0; i < NUM_RECORDS; i++) {
      assertEquals(baToString(kinesisMessageBatch.getStreamMessage(i).getValue()), DUMMY_RECORD_PREFIX + i);
    }
    assertFalse(kinesisMessageBatch.isEndOfPartitionGroup());

    // Fetch second batch
    kinesisMessageBatch = kinesisConsumer.fetchMessages(kinesisMessageBatch.getOffsetOfNextBatch(), TIMEOUT);
    assertEquals(kinesisMessageBatch.getMessageCount(), NUM_RECORDS);
    for (int i = 0; i < NUM_RECORDS; i++) {
      assertEquals(baToString(kinesisMessageBatch.getStreamMessage(i).getValue()), DUMMY_RECORD_PREFIX + i);
    }
    assertFalse(kinesisMessageBatch.isEndOfPartitionGroup());

    // Expect only 1 call to get shard iterator and 2 calls to get records
    verify(kinesisClient, times(1)).getShardIterator(any(GetShardIteratorRequest.class));
    verify(kinesisClient, times(2)).getRecords(any(GetRecordsRequest.class));
  }