test(batch): add 100ms wait to reduce flakiness in BatchContextTest

Author: bevzzz

src/it/java/io/weaviate/integration/OIDCSupportITest.java

@@ -81,7 +81,7 @@ public void test_resourceOwnerPassword() throws Exception {
     Assume.assumeTrue("WCS_DUMMY_CI_PW is not set", WCS_DUMMY_CI_PW != null && !WCS_DUMMY_CI_PW.isBlank());
     Assume.assumeTrue("no internet connection", hasInternetConnection());
 
-    // Check norwal resource owner password flow works.
+    // Check normal resource owner password flow works.
     var password = Authentication.resourceOwnerPassword(WCS_DUMMY_CI_USERNAME, WCS_DUMMY_CI_PW, List.of());
     var auth = SpyTokenProvider.spyOn(password);
     pingWeaviate(wcsContainer, auth);

src/test/java/io/weaviate/client6/v1/api/collections/batch/BatchContextTest.java

@@ -2,6 +2,7 @@
 
 import java.io.IOException;
 import java.util.ArrayList;
+import java.util.Collection;
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
@@ -171,25 +172,21 @@ public void test_drainOnClose() throws Exception {
     // Contrary the test above, we expect the objects to be sent
     // only after context.close(), as the half-empty batch will
     // be drained. Similarly, we want to ack everything as it arrives.
-    Future<?> backgroundAcks = BACKGROUND.submit(() -> {
+    BACKGROUND.submit(() -> {
       try {
         List<String> received = recvDataAndAck();
         Assertions.assertThat(tasks)
             .extracting(TaskHandle::id).containsExactlyInAnyOrderElementsOf(received);
         Assertions.assertThat(tasks)
             .extracting(TaskHandle::isAcked).allMatch(CompletableFuture::isDone);
-
-        // Wait until the Results event's been processed to guarantee
-        // that the tasks' futures are completed before asserting.
-        Future<?> results = out.emitEvent(new Event.Results(received, Collections.emptyMap()));
-        results.get();
+        out.emitEvent(new Event.Results(received, Collections.emptyMap()));
       } catch (Exception e) {
         throw new RuntimeException(e);
       }
     });
 
     context.close();
-    backgroundAcks.get(); // Wait for the "mock server" to process the data message.
+    awaitResults(tasks);
 
     Assertions.assertThat(tasks).extracting(TaskHandle::result)
         .allMatch(CompletableFuture::isDone)
@@ -268,6 +265,13 @@ public void test_backoffBacklog() throws Exception {
 
     context.close();
 
+    // Wait until the Results event's been processed to guarantee
+    // that the tasks' futures are completed before asserting.
+    CompletableFuture.allOf(
+        tasks.stream().map(TaskHandle::result)
+            .toArray(CompletableFuture[]::new))
+        .get(100, TimeUnit.MILLISECONDS);
+
     Assertions.assertThat(tasks).extracting(TaskHandle::result)
         .allMatch(CompletableFuture::isDone)
         .extracting(CompletableFuture::get).extracting(TaskHandle.Result::error)
@@ -464,6 +468,15 @@ private List<String> recvDataAndAck() throws InterruptedException {
     return ids;
   }
 
+  private void awaitResults(Collection<TaskHandle> tasks) throws Exception {
+    // Wait until the Results event's been processed to guarantee
+    // that the tasks' futures are completed before asserting.
+    CompletableFuture.allOf(
+        tasks.stream().map(TaskHandle::result)
+            .toArray(CompletableFuture[]::new))
+        .get(100, TimeUnit.MILLISECONDS);
+  }
+
   static String getBeacon(WeaviateProtoBatch.BatchReference reference) {
     return "weaviate://localhost/" + reference.getToCollection() + "/" + reference.getToUuid();
   }