From 8beb1050784f4ba675ebc7b74d5a16c4c2afed45 Mon Sep 17 00:00:00 2001
From: stair <stairdav@amazon.com>
Date: Wed, 19 Apr 2023 17:53:08 -0400
Subject: [PATCH] Bugfix: moved instantation of `Future` inside `Supplier`.

---
 .../kinesis/retrieval/KinesisClientFacade.java  |  5 +----
 .../retrieval/KinesisClientFacadeTest.java      | 17 ++++++++---------
 2 files changed, 9 insertions(+), 13 deletions(-)

diff --git a/amazon-kinesis-client/src/main/java/software/amazon/kinesis/retrieval/KinesisClientFacade.java b/amazon-kinesis-client/src/main/java/software/amazon/kinesis/retrieval/KinesisClientFacade.java
index 33a02fbc..5c9e732b 100644
--- a/amazon-kinesis-client/src/main/java/software/amazon/kinesis/retrieval/KinesisClientFacade.java
+++ b/amazon-kinesis-client/src/main/java/software/amazon/kinesis/retrieval/KinesisClientFacade.java
@@ -1,8 +1,6 @@
 package software.amazon.kinesis.retrieval;
 
-import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ExecutionException;
-import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
 
 import lombok.NonNull;
@@ -48,9 +46,8 @@ public final class KinesisClientFacade {
     public static DescribeStreamSummaryResponse describeStreamSummary(final String streamArn) {
         final DescribeStreamSummaryRequest request = KinesisRequestsBuilder
                 .describeStreamSummaryRequestBuilder().streamARN(streamArn).build();
-        final CompletableFuture<DescribeStreamSummaryResponse> future = kinesisClient.describeStreamSummary(request);
         final ServiceCallerSupplier<DescribeStreamSummaryResponse> dss =
-                () -> future.get(10, TimeUnit.SECONDS);
+                () -> kinesisClient.describeStreamSummary(request).get();
         return retryWhenThrottled(dss, 3, streamArn, "DescribeStreamSummary");
     }
 
diff --git a/amazon-kinesis-client/src/test/java/software/amazon/kinesis/retrieval/KinesisClientFacadeTest.java b/amazon-kinesis-client/src/test/java/software/amazon/kinesis/retrieval/KinesisClientFacadeTest.java
index f81616a8..7f978bdc 100644
--- a/amazon-kinesis-client/src/test/java/software/amazon/kinesis/retrieval/KinesisClientFacadeTest.java
+++ b/amazon-kinesis-client/src/test/java/software/amazon/kinesis/retrieval/KinesisClientFacadeTest.java
@@ -1,11 +1,9 @@
 package software.amazon.kinesis.retrieval;
 
 import java.util.concurrent.CompletableFuture;
-import java.util.concurrent.TimeUnit;
-import java.util.concurrent.TimeoutException;
+import java.util.concurrent.ExecutionException;
 
 import static org.junit.Assert.assertEquals;
-import static org.mockito.Matchers.anyInt;
 import static org.mockito.Mockito.any;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.times;
@@ -21,6 +19,7 @@ import org.mockito.runners.MockitoJUnitRunner;
 import software.amazon.awssdk.services.kinesis.KinesisAsyncClient;
 import software.amazon.awssdk.services.kinesis.model.DescribeStreamSummaryRequest;
 import software.amazon.awssdk.services.kinesis.model.DescribeStreamSummaryResponse;
+import software.amazon.awssdk.services.kinesis.model.LimitExceededException;
 
 @RunWith(MockitoJUnitRunner.class)
 public class KinesisClientFacadeTest {
@@ -49,19 +48,19 @@ public class KinesisClientFacadeTest {
     public void testDescribeStreamSummaryRetries() throws Exception {
         final DescribeStreamSummaryResponse expectedResponse = DescribeStreamSummaryResponse.builder().build();
         final CompletableFuture<DescribeStreamSummaryResponse> mockFuture = mock(CompletableFuture.class);
-        final TimeoutException timeoutException = new TimeoutException();
+        final ExecutionException executionException = new ExecutionException(LimitExceededException.builder().build());
 
         when(mockKinesisClient.describeStreamSummary(any(DescribeStreamSummaryRequest.class)))
                 .thenReturn(mockFuture);
-        when(mockFuture.get(anyInt(), any(TimeUnit.class)))
-                .thenThrow(timeoutException)
-                .thenThrow(timeoutException)
+        when(mockFuture.get())
+                .thenThrow(executionException)
+                .thenThrow(executionException)
                 .thenReturn(expectedResponse);
 
         final DescribeStreamSummaryResponse actualResponse = describeStreamSummary("retry me plz");
         assertEquals(expectedResponse, actualResponse);
 
-        verify(mockKinesisClient).describeStreamSummary(any(DescribeStreamSummaryRequest.class));
-        verify(mockFuture, times(3)).get(anyInt(), any(TimeUnit.class));
+        verify(mockKinesisClient, times(3)).describeStreamSummary(any(DescribeStreamSummaryRequest.class));
+        verify(mockFuture, times(3)).get();
     }
 }
\ No newline at end of file