Make tests more stable

Yelp · Aug 28, 2024 · d3963f1 · d3963f1
1 parent 18d2ad8
commit d3963f1
Showing 1 changed file with 20 additions and 0 deletions.
diff --git a/src/test/java/com/yelp/nrtsearch/server/luceneserver/nrt/NrtDataManagerTest.java b/src/test/java/com/yelp/nrtsearch/server/luceneserver/nrt/NrtDataManagerTest.java
@@ -501,6 +501,7 @@ public void testUploadManagerThread_processTask()
     nrtDataManager.enqueueUpload(copyState, List.of(refreshUploadFuture));
 
     refreshUploadFuture.get(30, java.util.concurrent.TimeUnit.SECONDS);
+    waitUntilDone(nrtDataManager);
 
     NrtFileMetaData nrtFileMetaData = new NrtFileMetaData(fileMetaData, PRIMARY_ID, "timestamp");
     NrtPointState nrtPointState =
@@ -570,6 +571,7 @@ public void testUploadManagerThread_processMultipleTasks()
 
     refreshUploadFuture.get(30, java.util.concurrent.TimeUnit.SECONDS);
     refreshUploadFuture2.get(30, java.util.concurrent.TimeUnit.SECONDS);
+    waitUntilDone(nrtDataManager);
 
     NrtFileMetaData nrtFileMetaData1 = new NrtFileMetaData(fileMetaData1, PRIMARY_ID, "timestamp");
     NrtFileMetaData nrtFileMetaData2 = new NrtFileMetaData(fileMetaData2, PRIMARY_ID, "timestamp");
@@ -629,6 +631,7 @@ public void testUploadManagerThread_sameVersion()
 
     refreshUploadFuture.get(30, java.util.concurrent.TimeUnit.SECONDS);
     refreshUploadFuture2.get(30, java.util.concurrent.TimeUnit.SECONDS);
+    waitUntilDone(nrtDataManager);
 
     NrtFileMetaData nrtFileMetaData = new NrtFileMetaData(fileMetaData, PRIMARY_ID, "timestamp");
     NrtPointState nrtPointState =
@@ -696,6 +699,7 @@ public void testUploadManagerThread_lowerVersion()
 
     refreshUploadFuture.get(30, java.util.concurrent.TimeUnit.SECONDS);
     refreshUploadFuture2.get(30, java.util.concurrent.TimeUnit.SECONDS);
+    waitUntilDone(nrtDataManager);
 
     NrtFileMetaData nrtFileMetaData = new NrtFileMetaData(fileMetaData, PRIMARY_ID, "timestamp");
     NrtPointState nrtPointState =
@@ -758,6 +762,7 @@ public void testUploadManagerThread_error()
     } catch (ExecutionException e) {
       assertEquals("java.io.IOException: error", e.getMessage());
     }
+    waitUntilDone(nrtDataManager);
     assertNull(nrtDataManager.getLastPointState());
 
     NrtFileMetaData nrtFileMetaData = new NrtFileMetaData(fileMetaData, PRIMARY_ID, "timestamp");
@@ -787,6 +792,21 @@ public void testUploadManagerThread_error()
     verifyNoMoreInteractions(mockPrimaryNode, mockRemoteBackend);
   }
 
+  private void waitUntilDone(NrtDataManager nrtDataManager) {
+    int count = 0;
+    while (nrtDataManager.getCurrentUploadTask() != null) {
+      count++;
+      try {
+        Thread.sleep(100);
+      } catch (InterruptedException e) {
+        Thread.currentThread().interrupt();
+      }
+      if (count > 100) {
+        fail("Upload tasks are not done");
+      }
+    }
+  }
+
   @Test
   public void testEnqueueNoRemoteCommit() {
     RemoteBackend mockRemoteBackend = mock(RemoteBackend.class);