Fix flaky test RecoveryFromGatewayIT.testMultipleReplicaShardAssignme…

…ntWithDelayedAllocationAndDifferentNodeStartTimeInBatchMode Signed-off-by: Swetha Guptha <[email protected]>
SwethaGuptha · Jun 18, 2024 · 92b14c2 · 92b14c2
1 parent 112704b
commit 92b14c2
Showing 1 changed file with 5 additions and 14 deletions.
diff --git a/server/src/internalClusterTest/java/org/opensearch/gateway/RecoveryFromGatewayIT.java b/server/src/internalClusterTest/java/org/opensearch/gateway/RecoveryFromGatewayIT.java
@@ -101,6 +101,7 @@
 import java.util.Map;
 import java.util.Set;
 import java.util.concurrent.ExecutionException;
+import java.util.concurrent.TimeUnit;
 import java.util.stream.Collectors;
 import java.util.stream.IntStream;
 
@@ -903,26 +904,16 @@ public void testMultipleReplicaShardAssignmentWithDelayedAllocationAndDifferentN
         assertTrue(clusterRerouteResponse.isAcknowledged());
         ensureStableCluster(6);
         waitUntil(
-            () -> client().admin().cluster().health(Requests.clusterHealthRequest().timeout("5m")).actionGet().getInitializingShards() == 0
+            () -> client().admin().cluster().health(Requests.clusterHealthRequest().timeout("5m")).actionGet().getActiveShards() == 3, 2, TimeUnit.MINUTES
         );
-
         health = client().admin().cluster().health(Requests.clusterHealthRequest().timeout("5m")).actionGet();
         assertFalse(health.isTimedOut());
         assertEquals(YELLOW, health.getStatus());
         assertEquals(1, health.getUnassignedShards());
         assertEquals(1, health.getDelayedUnassignedShards());
-        allocationExplainResponse = client().admin()
-            .cluster()
-            .prepareAllocationExplain()
-            .setIndex("test")
-            .setShard(0)
-            .setPrimary(false)
-            .get();
-        assertEquals(
-            AllocationDecision.ALLOCATION_DELAYED,
-            allocationExplainResponse.getExplanation().getShardAllocationDecision().getAllocateDecision().getAllocationDecision()
-        );
-
+        waitUntil(() ->
+            AllocationDecision.ALLOCATION_DELAYED.equals(client().admin().cluster().prepareAllocationExplain()
+                .setIndex("test").setShard(0).setPrimary(false).get().getExplanation().getShardAllocationDecision().getAllocateDecision().getAllocationDecision()), 2, TimeUnit.MINUTES);
         logger.info("--> restarting the node 0");
         internalCluster().startDataOnlyNode(
             Settings.builder().put("node.name", nodesWithReplicaShards.get(1)).put(replicaNode1DataPathSettings).build()