Fix flaky test SegmentReplicationStatsIT.testMultipleIndices (opensearch-project#12070) (opensearch-project#12100)

(cherry picked from commit 2f8d267)

Signed-off-by: Poojita Raj <poojiraj@amazon.com>
Signed-off-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>

Author: opensearch-trigger-bot[bot]

server/src/internalClusterTest/java/org/opensearch/indices/replication/SegmentReplicationStatsIT.java

@@ -268,12 +268,12 @@ public void testMultipleIndices() throws Exception {
         internalCluster().startClusterManagerOnlyNode();
         final String index_2 = "tst-index-2";
         List<String> nodes = new ArrayList<>();
-        final String primaryNode = internalCluster().startNode();
+        final String primaryNode = internalCluster().startDataOnlyNode();
         nodes.add(primaryNode);
         createIndex(INDEX_NAME, index_2);
 
         ensureYellowAndNoInitializingShards(INDEX_NAME, index_2);
-        nodes.add(internalCluster().startNode());
+        nodes.add(internalCluster().startDataOnlyNode());
         ensureGreen(INDEX_NAME, index_2);
 
         final long numDocs = scaledRandomIntBetween(50, 100);
@@ -284,44 +284,47 @@ public void testMultipleIndices() throws Exception {
         refresh(INDEX_NAME, index_2);
         waitForSearchableDocs(INDEX_NAME, numDocs, nodes);
         waitForSearchableDocs(index_2, numDocs, nodes);
+        ensureSearchable(INDEX_NAME, index_2);
 
         final IndexShard index_1_primary = getIndexShard(primaryNode, INDEX_NAME);
         final IndexShard index_2_primary = getIndexShard(primaryNode, index_2);
 
         assertTrue(index_1_primary.routingEntry().primary());
         assertTrue(index_2_primary.routingEntry().primary());
 
-        // test both indices are returned in the response.
-        SegmentReplicationStatsResponse segmentReplicationStatsResponse = client().admin()
-            .indices()
-            .prepareSegmentReplicationStats()
-            .execute()
-            .actionGet();
+        assertBusy(() -> {
+            // test both indices are returned in the response.
+            SegmentReplicationStatsResponse segmentReplicationStatsResponse = dataNodeClient().admin()
+                .indices()
+                .prepareSegmentReplicationStats()
+                .execute()
+                .actionGet();
 
-        Map<String, List<SegmentReplicationPerGroupStats>> replicationStats = segmentReplicationStatsResponse.getReplicationStats();
-        assertEquals(2, replicationStats.size());
-        List<SegmentReplicationPerGroupStats> replicationPerGroupStats = replicationStats.get(INDEX_NAME);
-        assertEquals(1, replicationPerGroupStats.size());
-        SegmentReplicationPerGroupStats perGroupStats = replicationPerGroupStats.get(0);
-        assertEquals(perGroupStats.getShardId(), index_1_primary.shardId());
-        Set<SegmentReplicationShardStats> replicaStats = perGroupStats.getReplicaStats();
-        assertEquals(1, replicaStats.size());
-        for (SegmentReplicationShardStats replica : replicaStats) {
-            assertNotNull(replica.getCurrentReplicationState());
-        }
+            Map<String, List<SegmentReplicationPerGroupStats>> replicationStats = segmentReplicationStatsResponse.getReplicationStats();
+            assertEquals(2, replicationStats.size());
+            List<SegmentReplicationPerGroupStats> replicationPerGroupStats = replicationStats.get(INDEX_NAME);
+            assertEquals(1, replicationPerGroupStats.size());
+            SegmentReplicationPerGroupStats perGroupStats = replicationPerGroupStats.get(0);
+            assertEquals(perGroupStats.getShardId(), index_1_primary.shardId());
+            Set<SegmentReplicationShardStats> replicaStats = perGroupStats.getReplicaStats();
+            assertEquals(1, replicaStats.size());
+            for (SegmentReplicationShardStats replica : replicaStats) {
+                assertNotNull(replica.getCurrentReplicationState());
+            }
 
-        replicationPerGroupStats = replicationStats.get(index_2);
-        assertEquals(1, replicationPerGroupStats.size());
-        perGroupStats = replicationPerGroupStats.get(0);
-        assertEquals(perGroupStats.getShardId(), index_2_primary.shardId());
-        replicaStats = perGroupStats.getReplicaStats();
-        assertEquals(1, replicaStats.size());
-        for (SegmentReplicationShardStats replica : replicaStats) {
-            assertNotNull(replica.getCurrentReplicationState());
-        }
+            replicationPerGroupStats = replicationStats.get(index_2);
+            assertEquals(1, replicationPerGroupStats.size());
+            perGroupStats = replicationPerGroupStats.get(0);
+            assertEquals(perGroupStats.getShardId(), index_2_primary.shardId());
+            replicaStats = perGroupStats.getReplicaStats();
+            assertEquals(1, replicaStats.size());
+            for (SegmentReplicationShardStats replica : replicaStats) {
+                assertNotNull(replica.getCurrentReplicationState());
+            }
+        }, 30, TimeUnit.SECONDS);
 
         // test only single index queried.
-        segmentReplicationStatsResponse = client().admin()
+        SegmentReplicationStatsResponse segmentReplicationStatsResponse = dataNodeClient().admin()
             .indices()
             .prepareSegmentReplicationStats()
             .setIndices(index_2)