From ba7d90c250406d9476d9716f8729bba66d523b64 Mon Sep 17 00:00:00 2001
From: kkewwei <kkewwei@163.com>
Date: Wed, 25 Sep 2024 12:12:25 +0800
Subject: [PATCH] Fix Flaky Test
 org.opensearch.gateway.RecoveryFromGatewayIT.testShardStoreFetchMultiNodeMultiIndexesUsingBatchAction

Signed-off-by: kkewwei <kkewwei@163.com>
---
 .../gateway/RecoveryFromGatewayIT.java          | 17 +++++++++++++----
 1 file changed, 13 insertions(+), 4 deletions(-)

diff --git a/server/src/internalClusterTest/java/org/opensearch/gateway/RecoveryFromGatewayIT.java b/server/src/internalClusterTest/java/org/opensearch/gateway/RecoveryFromGatewayIT.java
index bcf23a37c0010..84f4dee80519a 100644
--- a/server/src/internalClusterTest/java/org/opensearch/gateway/RecoveryFromGatewayIT.java
+++ b/server/src/internalClusterTest/java/org/opensearch/gateway/RecoveryFromGatewayIT.java
@@ -1366,7 +1366,8 @@ public void testSingleShardStoreFetchUsingBatchAction() throws ExecutionExceptio
         DiscoveryNode[] nodes = getDiscoveryNodes();
         TransportNodesListShardStoreMetadataBatch.NodesStoreFilesMetadataBatch response = prepareAndSendRequest(
             new String[] { indexName },
-            nodes
+            nodes,
+            false
         );
         Index index = resolveIndex(indexName);
         ShardId shardId = new ShardId(index, 0);
@@ -1379,12 +1380,14 @@ public void testSingleShardStoreFetchUsingBatchAction() throws ExecutionExceptio
 
     public void testShardStoreFetchMultiNodeMultiIndexesUsingBatchAction() throws Exception {
         internalCluster().startNodes(2);
+        ensureStableCluster(2);
         String indexName1 = "test1";
         String indexName2 = "test2";
         DiscoveryNode[] nodes = getDiscoveryNodes();
         TransportNodesListShardStoreMetadataBatch.NodesStoreFilesMetadataBatch response = prepareAndSendRequest(
             new String[] { indexName1, indexName2 },
-            nodes
+            nodes,
+            true
         );
         ClusterSearchShardsResponse searchShardsResponse = client().admin().cluster().prepareSearchShards(indexName1, indexName2).get();
         for (ClusterSearchShardsGroup clusterSearchShardsGroup : searchShardsResponse.getGroups()) {
@@ -1406,7 +1409,8 @@ public void testShardStoreFetchNodeNotConnectedUsingBatchAction() {
         String indexName = "test";
         TransportNodesListShardStoreMetadataBatch.NodesStoreFilesMetadataBatch response = prepareAndSendRequest(
             new String[] { indexName },
-            new DiscoveryNode[] { nonExistingNode }
+            new DiscoveryNode[] { nonExistingNode },
+            false
         );
         assertTrue(response.hasFailures());
         assertEquals(1, response.failures().size());
@@ -1418,6 +1422,7 @@ public void testShardStoreFetchCorruptedIndexUsingBatchAction() throws Exception
         String index1Name = "test1";
         String index2Name = "test2";
         prepareIndices(new String[] { index1Name, index2Name }, 1, 1);
+
         Map<ShardId, ShardAttributes> shardAttributesMap = prepareRequestMap(new String[] { index1Name, index2Name }, 1);
         Index index1 = resolveIndex(index1Name);
         ShardId shardId1 = new ShardId(index1, 0);
@@ -1513,10 +1518,14 @@ private void prepareIndices(String[] indices, int numberOfPrimaryShards, int num
 
     private TransportNodesListShardStoreMetadataBatch.NodesStoreFilesMetadataBatch prepareAndSendRequest(
         String[] indices,
-        DiscoveryNode[] nodes
+        DiscoveryNode[] nodes,
+        boolean ensureGreen
     ) {
         Map<ShardId, ShardAttributes> shardAttributesMap = null;
         prepareIndices(indices, 1, 1);
+        if (ensureGreen) {
+            ensureGreen(indices);
+        }
         shardAttributesMap = prepareRequestMap(indices, 1);
         TransportNodesListShardStoreMetadataBatch.NodesStoreFilesMetadataBatch response;
         return ActionTestUtils.executeBlocking(