Fix flaky test testStatsOnShardUnassigned in RemoteStoreStatsIT

Signed-off-by: Ashish Singh <[email protected]>
opensearch-project · Aug 2, 2023 · 8d8df2c · 8d8df2c
1 parent 11c8426
commit 8d8df2c
Showing 1 changed file with 5 additions and 3 deletions.
diff --git a/server/src/internalClusterTest/java/org/opensearch/remotestore/RemoteStoreStatsIT.java b/server/src/internalClusterTest/java/org/opensearch/remotestore/RemoteStoreStatsIT.java
@@ -432,7 +432,7 @@ public void testStatsOnShardRelocation() {
         assertEquals(destinationNode, replicaShardStat.getShardRouting().currentNodeId());
     }
 
-    public void testStatsOnShardUnassigned() throws IOException {
+    public void testStatsOnShardUnassigned() throws Exception {
         // Scenario:
         // - Create index with single primary and two replica shard
         // - Index documents
@@ -445,10 +445,12 @@ public void testStatsOnShardUnassigned() throws IOException {
         indexDocs();
         int dataNodeCountBeforeStop = client().admin().cluster().prepareHealth().get().getNumberOfDataNodes();
         internalCluster().stopRandomDataNode();
+        assertBusy(
+            () -> assertEquals(dataNodeCountBeforeStop - 1, client().admin().cluster().prepareHealth().get().getNumberOfDataNodes())
+        );
         RemoteStoreStatsResponse response = client().admin().cluster().prepareRemoteStoreStats(INDEX_NAME, "0").get();
-        int dataNodeCountAfterStop = client().admin().cluster().prepareHealth().get().getNumberOfDataNodes();
         assertEquals(dataNodeCountBeforeStop, response.getTotalShards());
-        assertEquals(dataNodeCountAfterStop, response.getSuccessfulShards());
+        assertEquals(dataNodeCountBeforeStop - 1, response.getSuccessfulShards());
     }
 
     public void testStatsOnRemoteStoreRestore() throws IOException {