From d847a00473f341872d2428a05a95294120aa38e2 Mon Sep 17 00:00:00 2001
From: Anshu Agarwal <anshukag@amazon.com>
Date: Thu, 26 Oct 2023 10:07:28 +0530
Subject: [PATCH] Fix weighted routing flaky test

Signed-off-by: Anshu Agarwal <anshukag@amazon.com>
---
 .../opensearch/search/SearchWeightedRoutingIT.java  | 13 ++++++++-----
 1 file changed, 8 insertions(+), 5 deletions(-)

diff --git a/server/src/internalClusterTest/java/org/opensearch/search/SearchWeightedRoutingIT.java b/server/src/internalClusterTest/java/org/opensearch/search/SearchWeightedRoutingIT.java
index 5207dab83f1d9..f508e1dea1818 100644
--- a/server/src/internalClusterTest/java/org/opensearch/search/SearchWeightedRoutingIT.java
+++ b/server/src/internalClusterTest/java/org/opensearch/search/SearchWeightedRoutingIT.java
@@ -501,8 +501,11 @@ public void testShardRoutingWithNetworkDisruption_FailOpenEnabled() throws Excep
 
         logger.info("--> creating network partition disruption");
         final String clusterManagerNode1 = internalCluster().getClusterManagerName();
-        Set<String> nodesInOneSide = Stream.of(clusterManagerNode1, nodeMap.get("b").get(0)).collect(Collectors.toCollection(HashSet::new));
-        Set<String> nodesInOtherSide = Stream.of(nodeMap.get("a").get(0)).collect(Collectors.toCollection(HashSet::new));
+        Set<String> nodesInOneSide =
+            Stream.of(nodeMap.get("a").get(0)).collect(Collectors.toCollection(HashSet::new));
+        Set<String> nodesInOtherSide =
+            Stream.of(clusterManagerNode1, nodeMap.get("b").get(0), nodeMap.get("c").get(0))
+                .collect(Collectors.toCollection(HashSet::new));
 
         NetworkDisruption networkDisruption = new NetworkDisruption(
             new NetworkDisruption.TwoPartitions(nodesInOneSide, nodesInOtherSide),
@@ -870,8 +873,7 @@ private void assertSearchInAZ(String az) {
             SearchStats.Stats searchStats = stat.getIndices().getSearch().getTotal();
             if (stat.getNode().isDataNode()) {
                 if (stat.getNode().getId().equals(dataNodeId)) {
-                    Assert.assertTrue(searchStats.getFetchCount() > 0L);
-                    Assert.assertTrue(searchStats.getQueryCount() > 0L);
+                    Assert.assertTrue(searchStats.getFetchCount() > 0L || searchStats.getQueryCount() > 0L);
                 }
             }
         }
@@ -945,7 +947,6 @@ public void testSearchAggregationWithNetworkDisruption_FailOpenEnabled() throws
         }
 
         logger.info("--> network disruption is stopped");
-        networkDisruption.stopDisrupting();
 
         for (int i = 0; i < 50; i++) {
             try {
@@ -962,6 +963,8 @@ public void testSearchAggregationWithNetworkDisruption_FailOpenEnabled() throws
                 fail("search should not fail");
             }
         }
+        networkDisruption.stopDisrupting();
+
         assertSearchInAZ("b");
         assertSearchInAZ("c");
         assertNoSearchInAZ("a");