From 2f8d267adbdfef175004912e877bdbd0814f27f1 Mon Sep 17 00:00:00 2001
From: Poojita Raj <poojiraj@amazon.com>
Date: Tue, 30 Jan 2024 16:15:41 -0800
Subject: [PATCH] Fix flaky test SegmentReplicationStatsIT.testMultipleIndices
 (#12070)

Signed-off-by: Poojita Raj <poojiraj@amazon.com>
---
 .../SegmentReplicationStatsIT.java            | 61 ++++++++++---------
 1 file changed, 32 insertions(+), 29 deletions(-)

diff --git a/server/src/internalClusterTest/java/org/opensearch/indices/replication/SegmentReplicationStatsIT.java b/server/src/internalClusterTest/java/org/opensearch/indices/replication/SegmentReplicationStatsIT.java
index 766471fdc0756..89aef6f0be1a6 100644
--- a/server/src/internalClusterTest/java/org/opensearch/indices/replication/SegmentReplicationStatsIT.java
+++ b/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,6 +284,7 @@ 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);
@@ -291,37 +292,39 @@ public void testMultipleIndices() throws Exception {
         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)