From 26c737cc2a3a36e27913274e3139dc8243cd747b Mon Sep 17 00:00:00 2001
From: Joshua Kim <kimjos@amazon.com>
Date: Tue, 7 Apr 2020 03:22:56 -0400
Subject: [PATCH] Adding in memory implememtation for pending checkpoint state

---
 .../software/amazon/kinesis/checkpoint/InMemoryCheckpointer.java | 1 +
 1 file changed, 1 insertion(+)

diff --git a/amazon-kinesis-client/src/test/java/software/amazon/kinesis/checkpoint/InMemoryCheckpointer.java b/amazon-kinesis-client/src/test/java/software/amazon/kinesis/checkpoint/InMemoryCheckpointer.java
index 83327931..a6190dfc 100644
--- a/amazon-kinesis-client/src/test/java/software/amazon/kinesis/checkpoint/InMemoryCheckpointer.java
+++ b/amazon-kinesis-client/src/test/java/software/amazon/kinesis/checkpoint/InMemoryCheckpointer.java
@@ -71,6 +71,7 @@ public class InMemoryCheckpointer implements Checkpointer {
     @Override
     public void prepareCheckpoint(String leaseKey, ExtendedSequenceNumber pendingCheckpoint, byte[] pendingCheckpointState, String concurrencyToken) throws KinesisClientLibException {
         pendingCheckpoints.put(leaseKey, pendingCheckpoint);
+        pendingCheckpointStates.put(leaseKey, pendingCheckpointState);
     }
 
     @Override