Update test

Signed-off-by: Rishab Nahata <[email protected]>
imRishN · Jul 27, 2024 · a7f0f0b · a7f0f0b
1 parent 09fea0c
commit a7f0f0b
Showing 1 changed file with 33 additions and 4 deletions.
diff --git a/server/src/test/java/org/opensearch/common/util/BatchRunnableExecutorTests.java b/server/src/test/java/org/opensearch/common/util/BatchRunnableExecutorTests.java
@@ -15,6 +15,7 @@
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.List;
+import java.util.concurrent.CountDownLatch;
 import java.util.function.Supplier;
 
 import static org.mockito.Mockito.atMost;
@@ -42,33 +43,53 @@ public void setupRunnables() {
     public void testRunWithoutTimeout() {
         setupRunnables();
         timeoutSupplier = () -> TimeValue.timeValueSeconds(1);
-        BatchRunnableExecutor executor = new BatchRunnableExecutor(runnableList, timeoutSupplier);
+        CountDownLatch countDownLatch = new CountDownLatch(1);
+        BatchRunnableExecutor executor = new BatchRunnableExecutor(runnableList, timeoutSupplier) {
+            @Override
+            public void onComplete() {
+                countDownLatch.countDown();
+            }
+        };
         executor.run();
         verify(runnable1, times(1)).run();
         verify(runnable2, times(1)).run();
         verify(runnable3, times(1)).run();
         verify(runnable1, never()).onTimeout();
         verify(runnable2, never()).onTimeout();
         verify(runnable3, never()).onTimeout();
+        assertEquals(0, countDownLatch.getCount());
     }
 
     public void testRunWithTimeout() {
         setupRunnables();
         timeoutSupplier = () -> TimeValue.timeValueNanos(1);
-        BatchRunnableExecutor executor = new BatchRunnableExecutor(runnableList, timeoutSupplier);
+        CountDownLatch countDownLatch = new CountDownLatch(1);
+        BatchRunnableExecutor executor = new BatchRunnableExecutor(runnableList, timeoutSupplier) {
+            @Override
+            public void onComplete() {
+                countDownLatch.countDown();
+            }
+        };
         executor.run();
         verify(runnable1, times(1)).onTimeout();
         verify(runnable2, times(1)).onTimeout();
         verify(runnable3, times(1)).onTimeout();
         verify(runnable1, never()).run();
         verify(runnable2, never()).run();
         verify(runnable3, never()).run();
+        assertEquals(0, countDownLatch.getCount());
     }
 
     public void testRunWithPartialTimeout() {
         setupRunnables();
         timeoutSupplier = () -> TimeValue.timeValueMillis(50);
-        BatchRunnableExecutor executor = new BatchRunnableExecutor(runnableList, timeoutSupplier);
+        CountDownLatch countDownLatch = new CountDownLatch(1);
+        BatchRunnableExecutor executor = new BatchRunnableExecutor(runnableList, timeoutSupplier) {
+            @Override
+            public void onComplete() {
+                countDownLatch.countDown();
+            }
+        };
         doAnswer(invocation -> {
             Thread.sleep(100);
             return null;
@@ -81,17 +102,25 @@ public void testRunWithPartialTimeout() {
         verify(runnable3, atMost(1)).onTimeout();
         verify(runnable2, atMost(1)).onTimeout();
         verify(runnable3, atMost(1)).onTimeout();
+        assertEquals(0, countDownLatch.getCount());
     }
 
     public void testRunWithEmptyRunnableList() {
         setupRunnables();
-        BatchRunnableExecutor executor = new BatchRunnableExecutor(Collections.emptyList(), timeoutSupplier);
+        CountDownLatch countDownLatch = new CountDownLatch(1);
+        BatchRunnableExecutor executor = new BatchRunnableExecutor(Collections.emptyList(), timeoutSupplier) {
+            @Override
+            public void onComplete() {
+                countDownLatch.countDown();
+            }
+        };
         executor.run();
         verify(runnable1, never()).onTimeout();
         verify(runnable2, never()).onTimeout();
         verify(runnable3, never()).onTimeout();
         verify(runnable1, never()).run();
         verify(runnable2, never()).run();
         verify(runnable3, never()).run();
+        assertEquals(1, countDownLatch.getCount());
     }
 }