Use Awaitility in two of the mongodb unit tests instead of depending …

…on Thread.sleep(). (#4288) Signed-off-by: David Venable <dlv@amazon.com>
opensearch-project · Mar 15, 2024 · f05703e · f05703e
1 parent e121375
commit f05703e
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 3 deletions.
diff --git a/...db/src/test/java/org/opensearch/dataprepper/plugins/mongo/export/ExportSchedulerTest.java b/...db/src/test/java/org/opensearch/dataprepper/plugins/mongo/export/ExportSchedulerTest.java
@@ -15,6 +15,7 @@
 import org.opensearch.dataprepper.plugins.mongo.coordination.partition.ExportPartition;
 import org.opensearch.dataprepper.plugins.mongo.coordination.partition.GlobalState;
 
+import java.time.Duration;
 import java.time.Instant;
 import java.util.List;
 import java.util.Map;
@@ -25,6 +26,7 @@
 import java.util.concurrent.Executors;
 import java.util.stream.Collectors;
 
+import static org.awaitility.Awaitility.await;
 import static org.hamcrest.MatcherAssert.assertThat;
 import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.is;
@@ -108,7 +110,9 @@ void test_export_run() throws InterruptedException {
 
         ExecutorService executorService = Executors.newSingleThreadExecutor();
         executorService.submit(() -> exportScheduler.run());
-        Thread.sleep(100);
+        await()
+                .atMost(Duration.ofSeconds(2))
+                .untilAsserted(() -> verify(coordinator, times(2)).createPartition(any()));
         executorService.shutdownNow();
 
 
@@ -161,7 +165,9 @@ void test_export_run_multiple_partitions() throws InterruptedException {
 
         ExecutorService executorService = Executors.newSingleThreadExecutor();
         executorService.submit(() -> exportScheduler.run());
-        Thread.sleep(100);
+        await()
+                .atMost(Duration.ofSeconds(2))
+                .untilAsserted(() -> verify(coordinator, times(4)).createPartition(any()));
         executorService.shutdownNow();
 
 

diff --git a/...ngodb/src/test/java/org/opensearch/dataprepper/plugins/mongo/export/ExportWorkerTest.java b/...ngodb/src/test/java/org/opensearch/dataprepper/plugins/mongo/export/ExportWorkerTest.java
@@ -28,12 +28,14 @@
 import org.opensearch.dataprepper.plugins.mongo.configuration.MongoDBSourceConfig;
 import org.opensearch.dataprepper.plugins.mongo.coordination.partition.DataQueryPartition;
 
+import java.time.Duration;
 import java.util.List;
 import java.util.Optional;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.Future;
 
+import static org.awaitility.Awaitility.await;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyInt;
 import static org.mockito.ArgumentMatchers.anyString;
@@ -160,7 +162,9 @@ public void test_shouldProcessPartitionSuccess(final String partitionKey) throws
             }
         });
 
-        Thread.sleep(100);
+        await()
+                .atMost(Duration.ofSeconds(2))
+                .untilAsserted(() -> verify(mongoClient).getDatabase(eq("test")));
         executorService.shutdownNow();
         // Then dependencies are called
         verify(mongoClient).getDatabase(eq("test"));