From fab2d416afbe43e1ec17b0d8309c3db438571920 Mon Sep 17 00:00:00 2001
From: Joshua Kim <kimjos@amazon.com>
Date: Thu, 29 Apr 2021 17:12:47 -0700
Subject: [PATCH] Adding guardrails test

---
 .../proxies/KinesisProxyTest.java             | 20 ++++++++++++-------
 1 file changed, 13 insertions(+), 7 deletions(-)

diff --git a/src/test/java/com/amazonaws/services/kinesis/clientlibrary/proxies/KinesisProxyTest.java b/src/test/java/com/amazonaws/services/kinesis/clientlibrary/proxies/KinesisProxyTest.java
index a847b89f..553d451a 100644
--- a/src/test/java/com/amazonaws/services/kinesis/clientlibrary/proxies/KinesisProxyTest.java
+++ b/src/test/java/com/amazonaws/services/kinesis/clientlibrary/proxies/KinesisProxyTest.java
@@ -435,26 +435,32 @@ public class KinesisProxyTest {
         verify(mockClient).listShards(any());
     }
 
-
     @Test
     public void testNoDuplicateShardsInPartialFailure() {
         proxy.setCachedShardMap(null);
         ListShardsResult firstPage = new ListShardsResult().withShards(shards.subList(0, 2)).withNextToken(NEXT_TOKEN);
         ListShardsResult lastPage = new ListShardsResult().withShards(shards.subList(2, shards.size())).withNextToken(null);
-
-        when(mockClient.listShards(any()))
-                .thenReturn(firstPage)
-                .thenThrow(new RuntimeException("Failed!"))
-                .thenReturn(firstPage)
-                .thenReturn(lastPage);
+        when(mockClient.listShards(any())).thenReturn(firstPage).thenThrow(new RuntimeException("Failed!")) // First call
+                .thenReturn(firstPage).thenReturn(lastPage); // second call
 
         try {
             proxy.getShardList();
         } catch (Exception e) {
+            // Do nothing
         }
         assertEquals(shards, proxy.getShardList());
     }
 
+    @Test
+    public void testDuplicateShardResponseDedupedGracefully() {
+        proxy.setCachedShardMap(null);
+        ListShardsResult pageOfShards = new ListShardsResult().withShards(shards).withNextToken(null);
+        when(mockClient.listShards(any())).thenReturn(pageOfShards);
+
+        proxy.getShardList();
+        assertEquals(shards, proxy.getShardList());
+    }
+
     private void mockListShardsForSingleResponse(List<Shard> shards) {
         when(mockClient.listShards(any())).thenReturn(listShardsResult);
         when(listShardsResult.getShards()).thenReturn(shards);