Fix intermittent test failures in AsyncApiCallTimeoutTest

test/protocol-tests/src/test/java/software/amazon/awssdk/protocol/tests/timeout/async/AsyncApiCallTimeoutTest.java

@@ -15,10 +15,6 @@
 
 package software.amazon.awssdk.protocol.tests.timeout.async;
 
-import static com.github.tomakehurst.wiremock.client.WireMock.aResponse;
-import static com.github.tomakehurst.wiremock.client.WireMock.anyUrl;
-import static com.github.tomakehurst.wiremock.client.WireMock.post;
-import static com.github.tomakehurst.wiremock.client.WireMock.stubFor;
 import static org.assertj.core.api.AssertionsForClassTypes.assertThat;
 import static org.assertj.core.api.AssertionsForClassTypes.assertThatThrownBy;
 
@@ -36,7 +32,10 @@
 import software.amazon.awssdk.core.async.AsyncResponseTransformer;
 import software.amazon.awssdk.core.exception.ApiCallTimeoutException;
 import software.amazon.awssdk.core.retry.RetryPolicy;
+import software.amazon.awssdk.http.HttpExecuteResponse;
+import software.amazon.awssdk.http.SdkHttpResponse;
 import software.amazon.awssdk.protocol.tests.timeout.BaseApiCallTimeoutTest;
+import software.amazon.awssdk.protocol.tests.util.MockAsyncHttpClient;
 import software.amazon.awssdk.regions.Region;
 import software.amazon.awssdk.services.protocolrestjson.ProtocolRestJsonAsyncClient;
 import software.amazon.awssdk.services.protocolrestjson.model.AllTypesResponse;
@@ -104,13 +103,34 @@ protected WireMockRule wireMock() {
 
     @Test
     public void increaseTimeoutInRequestOverrideConfig_shouldTakePrecedence() {
-        stubFor(post(anyUrl())
-                    .willReturn(aResponse().withStatus(200).withBody("{}").withFixedDelay(DELAY_AFTER_TIMEOUT)));
+        MockAsyncHttpClient mockClient = new MockAsyncHttpClient();
+        ProtocolRestJsonAsyncClient asyncClient = createClientWithMockClient(mockClient);
+        mockClient.stubNextResponse(mockResponse(200));
+        mockClient.withFixedDelay(DELAY_AFTER_TIMEOUT);
+
         CompletableFuture<AllTypesResponse> allTypesResponseCompletableFuture =
-            client.allTypes(b -> b.overrideConfiguration(c -> c.apiCallTimeout(Duration.ofMillis(DELAY_AFTER_TIMEOUT + 1000))));
+            asyncClient.allTypes(b -> b.overrideConfiguration(c -> c.apiCallTimeout(Duration.ofMillis(DELAY_AFTER_TIMEOUT + 1000))));
 
         AllTypesResponse response = allTypesResponseCompletableFuture.join();
         assertThat(response).isNotNull();
     }
 
+    public ProtocolRestJsonAsyncClient createClientWithMockClient(MockAsyncHttpClient mockClient) {
+        return ProtocolRestJsonAsyncClient.builder()
+                                          .region(Region.US_WEST_1)
+                                          .httpClient(mockClient)
+                                          .credentialsProvider(() -> AwsBasicCredentials.create("akid", "skid"))
+                                          .overrideConfiguration(b -> b.apiCallTimeout(Duration.ofMillis(TIMEOUT))
+                                                                       .retryPolicy(RetryPolicy.none()))
+                                          .build();
+    }
+
+    private HttpExecuteResponse mockResponse(int statusCode) {
+        return HttpExecuteResponse.builder()
+                                  .response(SdkHttpResponse.builder()
+                                                           .statusCode(statusCode)
+                                                           .build())
+                                  .build();
+    }
+
 }

test/protocol-tests/src/test/java/software/amazon/awssdk/protocol/tests/util/MockAsyncHttpClient.java

@@ -0,0 +1,89 @@
+package software.amazon.awssdk.protocol.tests.util;
+
+import static software.amazon.awssdk.utils.FunctionalUtils.invokeSafely;
+
+import java.nio.ByteBuffer;
+import java.util.Optional;
+import java.util.concurrent.CompletableFuture;
+import org.reactivestreams.Subscriber;
+import org.reactivestreams.Subscription;
+import software.amazon.awssdk.http.HttpExecuteResponse;
+import software.amazon.awssdk.http.async.AsyncExecuteRequest;
+import software.amazon.awssdk.http.async.SdkAsyncHttpClient;
+import software.amazon.awssdk.http.async.SdkHttpContentPublisher;
+import software.amazon.awssdk.utils.IoUtils;
+
+/**
+ * Mock implementation of {@link SdkAsyncHttpClient}.
+ */
+public class MockAsyncHttpClient implements SdkAsyncHttpClient {
+    private HttpExecuteResponse nextResponse;
+    private long fixedDelay;
+
+    @Override
+    public CompletableFuture<Void> execute(AsyncExecuteRequest request) {
+
+        byte[] content = nextResponse.responseBody().map(p -> invokeSafely(() -> IoUtils.toByteArray(p)))
+                                     .orElseGet(() -> new byte[0]);
+
+        request.responseHandler().onHeaders(nextResponse.httpResponse());
+        request.responseHandler().onStream(new ResponsePublisher(content));
+
+        try {
+            Thread.sleep(fixedDelay);
+        } catch (InterruptedException e) {
+        }
+
+        return CompletableFuture.completedFuture(null);
+    }
+
+    @Override
+    public void close() {
+    }
+
+    public void stubNextResponse(HttpExecuteResponse nextResponse) {
+        this.nextResponse = nextResponse;
+    }
+
+    public void withFixedDelay(long fixedDelay) {
+        this.fixedDelay = fixedDelay;
+    }
+
+    private static class ResponsePublisher implements SdkHttpContentPublisher {
+        private final byte[] content;
+
+        private ResponsePublisher(byte[] content) {
+            this.content = content;
+        }
+
+        @Override
+        public Optional<Long> contentLength() {
+            return Optional.of((long) content.length);
+        }
+
+        @Override
+        public void subscribe(Subscriber<? super ByteBuffer> s) {
+            s.onSubscribe(new Subscription() {
+                private boolean running = true;
+
+                @Override
+                public void request(long n) {
+                    if (n <= 0) {
+                        running = false;
+                        s.onError(new IllegalArgumentException("Demand must be positive"));
+                    } else if (running) {
+                        running = false;
+                        s.onNext(ByteBuffer.wrap(content));
+                        s.onComplete();
+                    }
+                }
+
+                @Override
+                public void cancel() {
+                    running = false;
+                }
+            });
+        }
+    }
+
+}