Add retries to CrtHttpClientApiCallTimeoutTest for race condition and BatchingEfficiencyUnderLoadTest with sleep to allow batch flush completion

Author: joviegas

services/sqs/src/test/java/software/amazon/awssdk/services/sqs/batchmanager/BatchingEfficiencyUnderLoadTest.java

@@ -31,11 +31,11 @@
 import java.util.stream.Collectors;
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
-import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.RegisterExtension;
 import software.amazon.awssdk.auth.credentials.AwsBasicCredentials;
 import software.amazon.awssdk.auth.credentials.StaticCredentialsProvider;
 import software.amazon.awssdk.services.sqs.SqsAsyncClient;
+import software.amazon.awssdk.testutils.retry.RetryableTest;
 
 
 /**
@@ -84,7 +84,7 @@ void tearDown() {
     /**
      * Test runs heavy load and expects average batch sizes to be close to max.
      */
-    @Test
+    @RetryableTest(maxRetries = 3)
     void sendMessage_whenHighLoadScenario_shouldEfficientlyBatchMessages() throws Exception {
         int expectedBatchSize = 25; // more than double the actual max of 10
         int rateLimit = 1000 / SEND_FREQUENCY_MILLIS * expectedBatchSize;
@@ -126,7 +126,7 @@ void sendMessage_whenHighLoadScenario_shouldEfficientlyBatchMessages() throws Ex
     /**
      * Test runs a load that should cause an average batch size of 5.
      */
-    @Test
+    @RetryableTest(maxRetries = 3)
     void sendMessage_whenMediumLoadScenario_shouldCreateHalfSizeBatches() throws Exception {
         int expectedBatchSize = 5;
         int rateLimit = 1000 / SEND_FREQUENCY_MILLIS * expectedBatchSize;
@@ -157,7 +157,7 @@ void sendMessage_whenMediumLoadScenario_shouldCreateHalfSizeBatches() throws Exc
             .isLessThan(messageCount / 3d);
     }
 
-    @Test
+    @RetryableTest(maxRetries = 3)
     void sendMessage_whenLowLoadScenario_shouldCreateSmallBatches() throws Exception {
         int expectedBatchSize = 1;
         int rateLimit = 1000 / SEND_FREQUENCY_MILLIS * expectedBatchSize;
@@ -217,5 +217,8 @@ private void runThroughputTest(int messageCount, int rateLimit) throws Interrupt
 
         executor.shutdown();
         executor.awaitTermination(10, TimeUnit.SECONDS);
+        
+        // Wait for batch manager to flush all pending batches
+        Thread.sleep(SEND_FREQUENCY_MILLIS * 3);
     }
 }

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

@@ -26,14 +26,15 @@
 import java.net.URI;
 import java.time.Duration;
 import org.junit.jupiter.api.BeforeEach;
-import org.junit.jupiter.api.Test;
 import software.amazon.awssdk.auth.credentials.AwsBasicCredentials;
 import software.amazon.awssdk.core.exception.ApiCallAttemptTimeoutException;
 import software.amazon.awssdk.core.exception.ApiCallTimeoutException;
 import software.amazon.awssdk.http.crt.AwsCrtHttpClient;
 import software.amazon.awssdk.regions.Region;
 import software.amazon.awssdk.services.protocolrestjson.ProtocolRestJsonClient;
 import software.amazon.awssdk.services.protocolrestjson.ProtocolRestJsonClientBuilder;
+import software.amazon.awssdk.testutils.retry.RetryableTest;
+import software.amazon.awssdk.crt.http.HttpException;
 
 @WireMockTest
 public class CrtHttpClientApiCallTimeoutTest {
@@ -49,7 +50,7 @@ public void setup(WireMockRuntimeInfo wiremock) {
                                               .credentialsProvider(() -> AwsBasicCredentials.create("akid", "skid"));
     }
 
-    @Test
+    @RetryableTest(maxRetries = 3)
     void apiCallAttemptExceeded_shouldThrowApiCallAttemptTimeoutException() {
         stubFor(any(anyUrl()).willReturn(aResponse().withStatus(200).withBody("{}").withFixedDelay(2000)));
         try (ProtocolRestJsonClient client =
@@ -61,7 +62,7 @@ void apiCallAttemptExceeded_shouldThrowApiCallAttemptTimeoutException() {
         }
     }
 
-    @Test
+    @RetryableTest(maxRetries = 3)
     void apiCallExceeded_shouldThrowApiCallAttemptTimeoutException() {
         stubFor(any(anyUrl()).willReturn(aResponse().withStatus(200).withBody("{}").withFixedDelay(2000)));
         try (ProtocolRestJsonClient client = clientBuilder.overrideConfiguration(b -> b.apiCallTimeout(Duration.ofMillis(10)))