From b48ab58ec88522d47741c0444d1598d52a4b4b96 Mon Sep 17 00:00:00 2001
From: Sam Whittle <samuelw@google.com>
Date: Wed, 3 Dec 2025 13:39:06 +0100
Subject: [PATCH 1/2] Fix flaky
 AsyncBatchWriteHandlerTest.correctlyLimitConcurrency

---
 .../beam/sdk/io/aws2/common/AsyncBatchWriteHandlerTest.java    | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/sdks/java/io/amazon-web-services2/src/test/java/org/apache/beam/sdk/io/aws2/common/AsyncBatchWriteHandlerTest.java b/sdks/java/io/amazon-web-services2/src/test/java/org/apache/beam/sdk/io/aws2/common/AsyncBatchWriteHandlerTest.java
index cd7aca3c9f9b..1f0a7da35bfd 100644
--- a/sdks/java/io/amazon-web-services2/src/test/java/org/apache/beam/sdk/io/aws2/common/AsyncBatchWriteHandlerTest.java
+++ b/sdks/java/io/amazon-web-services2/src/test/java/org/apache/beam/sdk/io/aws2/common/AsyncBatchWriteHandlerTest.java
@@ -35,6 +35,7 @@
 import java.util.List;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.Future;
+import java.util.concurrent.TimeUnit;
 import java.util.function.BiFunction;
 import java.util.function.Function;
 import java.util.function.Supplier;
@@ -230,7 +231,7 @@ public void correctlyLimitConcurrency() throws Throwable {
     nextResults.complete(emptyList());
 
     handler.waitForCompletion();
-    assertThat(future).isDone();
+    eventually(5, () -> assertThat(future).isDone());
   }
 
   static class SubmitFn<T, V> implements BiFunction<String, List<T>, CompletableFuture<List<V>>> {

From 648dc7ff02b67827025f79bc91407b1d75790cf2 Mon Sep 17 00:00:00 2001
From: Sam Whittle <samuelw@google.com>
Date: Wed, 3 Dec 2025 13:53:50 +0100
Subject: [PATCH 2/2] spotless

---
 .../beam/sdk/io/aws2/common/AsyncBatchWriteHandlerTest.java      | 1 -
 1 file changed, 1 deletion(-)

diff --git a/sdks/java/io/amazon-web-services2/src/test/java/org/apache/beam/sdk/io/aws2/common/AsyncBatchWriteHandlerTest.java b/sdks/java/io/amazon-web-services2/src/test/java/org/apache/beam/sdk/io/aws2/common/AsyncBatchWriteHandlerTest.java
index 1f0a7da35bfd..056d856e442c 100644
--- a/sdks/java/io/amazon-web-services2/src/test/java/org/apache/beam/sdk/io/aws2/common/AsyncBatchWriteHandlerTest.java
+++ b/sdks/java/io/amazon-web-services2/src/test/java/org/apache/beam/sdk/io/aws2/common/AsyncBatchWriteHandlerTest.java
@@ -35,7 +35,6 @@
 import java.util.List;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.Future;
-import java.util.concurrent.TimeUnit;
 import java.util.function.BiFunction;
 import java.util.function.Function;
 import java.util.function.Supplier;