Fix flaky test RecoveryFromGatewayIT.testMultipleReplicaShardAssignme…

…ntWithDelayedAllocationAndDifferentNodeStartTimeInBatchMode (#14424) Signed-off-by: Swetha Guptha <[email protected]> Co-authored-by: Swetha Guptha <[email protected]>
opensearch-project · Jun 18, 2024 · 802f2e6 · 802f2e6
1 parent 120678d
commit 802f2e6
Showing 1 changed file with 19 additions and 25 deletions.
diff --git a/server/src/internalClusterTest/java/org/opensearch/gateway/RecoveryFromGatewayIT.java b/server/src/internalClusterTest/java/org/opensearch/gateway/RecoveryFromGatewayIT.java
@@ -34,7 +34,6 @@
 
 import org.apache.lucene.index.CorruptIndexException;
 import org.opensearch.Version;
-import org.opensearch.action.admin.cluster.allocation.ClusterAllocationExplainResponse;
 import org.opensearch.action.admin.cluster.configuration.AddVotingConfigExclusionsAction;
 import org.opensearch.action.admin.cluster.configuration.AddVotingConfigExclusionsRequest;
 import org.opensearch.action.admin.cluster.configuration.ClearVotingConfigExclusionsAction;
@@ -101,6 +100,8 @@
 import java.util.Map;
 import java.util.Set;
 import java.util.concurrent.ExecutionException;
+import java.util.concurrent.TimeUnit;
+import java.util.function.BooleanSupplier;
 import java.util.stream.Collectors;
 import java.util.stream.IntStream;
 
@@ -883,17 +884,20 @@ public void testMultipleReplicaShardAssignmentWithDelayedAllocationAndDifferentN
         assertEquals(YELLOW, health.getStatus());
         assertEquals(2, health.getUnassignedShards());
         // shard should be unassigned because of Allocation_Delayed
-        ClusterAllocationExplainResponse allocationExplainResponse = client().admin()
-            .cluster()
-            .prepareAllocationExplain()
-            .setIndex("test")
-            .setShard(0)
-            .setPrimary(false)
-            .get();
-        assertEquals(
-            AllocationDecision.ALLOCATION_DELAYED,
-            allocationExplainResponse.getExplanation().getShardAllocationDecision().getAllocateDecision().getAllocationDecision()
+        BooleanSupplier delayedShardAllocationStatusVerificationSupplier = () -> AllocationDecision.ALLOCATION_DELAYED.equals(
+            client().admin()
+                .cluster()
+                .prepareAllocationExplain()
+                .setIndex("test")
+                .setShard(0)
+                .setPrimary(false)
+                .get()
+                .getExplanation()
+                .getShardAllocationDecision()
+                .getAllocateDecision()
+                .getAllocationDecision()
         );
+        waitUntil(delayedShardAllocationStatusVerificationSupplier, 2, TimeUnit.MINUTES);
 
         logger.info("--> restarting the node 1");
         internalCluster().startDataOnlyNode(
@@ -903,26 +907,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(delayedShardAllocationStatusVerificationSupplier, 2, TimeUnit.MINUTES);
         logger.info("--> restarting the node 0");
         internalCluster().startDataOnlyNode(
             Settings.builder().put("node.name", nodesWithReplicaShards.get(1)).put(replicaNode1DataPathSettings).build()