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