fix: await until tasks are acked

Author: bevzzz

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

@@ -159,7 +159,7 @@ private void closeContext() throws Exception {
 
     try {
       context.close();
-      eof.get();
+      eof.get(5, TimeUnit.SECONDS);
     } finally {
       contextClosed = true;
     }
@@ -177,15 +177,19 @@ public void test_sendOneBatch() throws Exception {
     List<String> received = recvDataAndAck();
     Assertions.assertThat(tasks)
         .extracting(TaskHandle::id).containsExactlyInAnyOrderElementsOf(received);
-    Assertions.assertThat(tasks)
-        .extracting(TaskHandle::isAcked).allMatch(CompletableFuture::isDone);
+
+    CompletableFuture<?>[] tasksAcked = tasks.stream()
+        .map(TaskHandle::isAcked).toArray(CompletableFuture[]::new);
+    Assertions.assertThat(CompletableFuture.allOf(tasksAcked))
+        .succeedsWithin(5, TimeUnit.SECONDS);
 
     out.beforeEof(new Event.Results(received, Collections.emptyMap()));
 
     // Since MockServer runs in the same thread as this test,
     // the context will be updated before the last emitEvent returns.
     closeContext();
 
+    // By the time context.close() returns all tasks MUST have results set.
     Assertions.assertThat(tasks).extracting(TaskHandle::result)
         .allMatch(CompletableFuture::isDone)
         .extracting(CompletableFuture::get).extracting(TaskHandle.Result::error)
@@ -207,8 +211,11 @@ public void test_drainOnClose() throws Exception {
         List<String> received = recvDataAndAck();
         Assertions.assertThat(tasks).extracting(TaskHandle::id)
             .containsExactlyInAnyOrderElementsOf(received);
-        Assertions.assertThat(tasks).extracting(TaskHandle::isAcked)
-            .allMatch(CompletableFuture::isDone);
+
+        CompletableFuture<?>[] tasksAcked = tasks.stream()
+            .map(TaskHandle::isAcked).toArray(CompletableFuture[]::new);
+        Assertions.assertThat(CompletableFuture.allOf(tasksAcked))
+            .succeedsWithin(5, TimeUnit.SECONDS);
       } catch (Exception e) {
         throw new RuntimeException(e);
       }