From 963da1c62eab892426bc61ce6805f8b2b5262d5c Mon Sep 17 00:00:00 2001
From: "opensearch-trigger-bot[bot]"
 <98922864+opensearch-trigger-bot[bot]@users.noreply.github.com>
Date: Fri, 3 Nov 2023 14:51:22 +0530
Subject: [PATCH] Fix weighted routing flaky test (#10955) (#11028)

(cherry picked from commit cb39d0057e7277e02a99e2cf5dabfc36855991c4)

Signed-off-by: Anshu Agarwal <anshukag@amazon.com>
---
 .../opensearch/search/SearchWeightedRoutingIT.java    | 11 ++++++-----
 1 file changed, 6 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 0975585e2984d..44a9041e97a3e 100644
--- a/server/src/internalClusterTest/java/org/opensearch/search/SearchWeightedRoutingIT.java
+++ b/server/src/internalClusterTest/java/org/opensearch/search/SearchWeightedRoutingIT.java
@@ -501,8 +501,9 @@ 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 +871,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 +945,6 @@ public void testSearchAggregationWithNetworkDisruption_FailOpenEnabled() throws
         }
 
         logger.info("--> network disruption is stopped");
-        networkDisruption.stopDisrupting();
 
         for (int i = 0; i < 50; i++) {
             try {
@@ -962,6 +961,8 @@ public void testSearchAggregationWithNetworkDisruption_FailOpenEnabled() throws
                 fail("search should not fail");
             }
         }
+        networkDisruption.stopDisrupting();
+
         assertSearchInAZ("b");
         assertSearchInAZ("c");
         assertNoSearchInAZ("a");