Fix flaky test SegmentReplicationStatsIT.testMultipleIndices (opensea…

…rch-project#12070) Signed-off-by: Poojita Raj <[email protected]>
rayshrey · Mar 18, 2024 · cae24a4 · cae24a4
1 parent 1d147d4
commit cae24a4
Showing 1 changed file with 32 additions and 29 deletions.
diff --git a/...nternalClusterTest/java/org/opensearch/indices/replication/SegmentReplicationStatsIT.java b/...nternalClusterTest/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)