From 250a33b344740e24082574dd224f2ea63a022fcc Mon Sep 17 00:00:00 2001
From: kkewwei <kkewwei@163.com>
Date: Thu, 26 Sep 2024 22:21:56 +0800
Subject: [PATCH] Fix Flaky Test
 org.opensearch.gateway.RecoveryFromGatewayIT.testShardStoreFetchMultiNodeMultiIndexesUsingBatchAction

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

diff --git a/server/src/internalClusterTest/java/org/opensearch/gateway/RecoveryFromGatewayIT.java b/server/src/internalClusterTest/java/org/opensearch/gateway/RecoveryFromGatewayIT.java
index 3aef3cdfd063d..8d0a3b5788a70 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,13 +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();
-        System.out.println();
         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()) {
@@ -1393,7 +1395,6 @@ public void testShardStoreFetchMultiNodeMultiIndexesUsingBatchAction() throws Ex
             ShardRouting[] shardRoutings = clusterSearchShardsGroup.getShards();
             assertEquals(2, shardRoutings.length);
             for (ShardRouting shardRouting : shardRoutings) {
-                System.out.println("----------------" + shardRouting.shortSummary());
                 TransportNodesListShardStoreMetadataBatch.NodeStoreFilesMetadata nodeStoreFilesMetadata = response.getNodesMap()
                     .get(shardRouting.currentNodeId())
                     .getNodeStoreFilesMetadataBatch()
@@ -1408,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());
@@ -1515,10 +1517,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(
@@ -1534,11 +1540,8 @@ private void assertNodeStoreFilesMetadataSuccessCase(
         assertNull(nodeStoreFilesMetadata.getStoreFileFetchException());
         TransportNodesListShardStoreMetadataHelper.StoreFilesMetadata storeFileMetadata = nodeStoreFilesMetadata.storeFilesMetadata();
         assertFalse(storeFileMetadata.isEmpty());
-        System.out.println("----------" + storeFileMetadata);
         assertEquals(shardId, storeFileMetadata.shardId());
         assertNotNull(storeFileMetadata.peerRecoveryRetentionLeases());
-        System.out.println("----------" + storeFileMetadata.peerRecoveryRetentionLeases());
-
     }
 
     private void assertNodeGatewayStartedShardsHappyCase(GatewayStartedShard gatewayStartedShard) {