Using awaitility to improve the consistency of the RSSSourceTest.

Signed-off-by: David Venable <[email protected]>
opensearch-project · Jul 12, 2023 · 6d152bd · 6d152bd
1 parent d122289
commit 6d152bd
Showing 1 changed file with 9 additions and 2 deletions.
diff --git a/...rss-source/src/test/java/org/opensearch/dataprepper/plugins/source/rss/RSSSourceTest.java b/...rss-source/src/test/java/org/opensearch/dataprepper/plugins/source/rss/RSSSourceTest.java
@@ -18,6 +18,7 @@
 
 import java.time.Duration;
 
+import static org.awaitility.Awaitility.await;
 import static org.mockito.ArgumentMatchers.anyCollection;
 import static org.mockito.ArgumentMatchers.anyInt;
 import static org.mockito.Mockito.atLeast;
@@ -62,9 +63,15 @@ public void tearDown() {
     @Test
     void test_ExecutorService_keep_writing_Events_to_Buffer() throws Exception {
         rssSource.start(buffer);
-        Thread.sleep(pollingFrequency.toMillis());
+        await().atMost(pollingFrequency.multipliedBy(2))
+                        .untilAsserted(() -> {
+                            verify(buffer, atLeastOnce()).writeAll(anyCollection(), anyInt());
+                        });
         verify(buffer, atLeastOnce()).writeAll(anyCollection(), anyInt());
-        Thread.sleep(pollingFrequency.toMillis());
+        await().atMost(pollingFrequency.multipliedBy(2))
+                .untilAsserted(() -> {
+                    verify(buffer, atLeast(2)).writeAll(anyCollection(), anyInt());
+                });
         verify(buffer, atLeast(2)).writeAll(anyCollection(), anyInt());
     }
 }