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 75feb19e..c877ed79 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 @@ -483,6 +483,26 @@ 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); + + try { + proxy.getShardList(); + } catch (Exception e) { + } + assertEquals(shards, proxy.getShardList()); + } + private void mockListShardsForSingleResponse(List shards) { when(mockClient.listShards(any())).thenReturn(listShardsResult); when(listShardsResult.getShards()).thenReturn(shards);