Fix flaky test opensearch-project#11256

Signed-off-by: Ashish Singh <[email protected]>
ashking94 · Jan 16, 2024 · 73fac30 · 73fac30
1 parent c7196da
commit 73fac30
Showing 1 changed file with 9 additions and 2 deletions.
diff --git a/server/src/test/java/org/opensearch/index/shard/RemoteStoreRefreshListenerTests.java b/server/src/test/java/org/opensearch/index/shard/RemoteStoreRefreshListenerTests.java
@@ -418,7 +418,7 @@ public void testTrackerData() throws Exception {
         RemoteStoreRefreshListener listener = tuple.v1();
         RemoteStoreStatsTrackerFactory trackerFactory = tuple.v2();
         RemoteSegmentTransferTracker tracker = trackerFactory.getRemoteSegmentTransferTracker(indexShard.shardId());
-        assertNoLag(tracker);
+        assertBusy(() -> assertNoLag(tracker));
         indexDocs(100, randomIntBetween(100, 200));
         indexShard.refresh("test");
         listener.afterRefresh(true);
@@ -473,6 +473,14 @@ private Tuple<RemoteStoreRefreshListener, RemoteStoreStatsTrackerFactory> mockIn
             new InternalEngineFactory()
         );
 
+        RemoteSegmentTransferTracker tracker = indexShard.getRemoteStoreStatsTrackerFactory()
+            .getRemoteSegmentTransferTracker(indexShard.shardId());
+        try {
+            assertBusy(() -> assertTrue(tracker.getTotalUploadsSucceeded() > 0));
+        } catch (Exception e) {
+            assert false;
+        }
+
         indexDocs(1, randomIntBetween(1, 100));
 
         // Mock indexShard.store().directory()
@@ -554,7 +562,6 @@ private Tuple<RemoteStoreRefreshListener, RemoteStoreStatsTrackerFactory> mockIn
         RecoverySettings recoverySettings = mock(RecoverySettings.class);
         when(recoverySettings.getMinRemoteSegmentMetadataFiles()).thenReturn(10);
         when(shard.getRecoverySettings()).thenReturn(recoverySettings);
-        RemoteSegmentTransferTracker tracker = remoteStoreStatsTrackerFactory.getRemoteSegmentTransferTracker(indexShard.shardId());
         RemoteStoreRefreshListener refreshListener = new RemoteStoreRefreshListener(shard, emptyCheckpointPublisher, tracker);
         refreshListener.afterRefresh(true);
         return Tuple.tuple(refreshListener, remoteStoreStatsTrackerFactory);