Address comments

Signed-off-by: Arpit Bandejiya <[email protected]>
opensearch-project · Mar 24, 2024 · cc3f2c5 · cc3f2c5
1 parent 60ccd7b
commit cc3f2c5
Showing 1 changed file with 11 additions and 0 deletions.
diff --git a/...er/src/test/java/org/opensearch/cluster/routing/allocation/BalanceConfigurationTests.java b/...er/src/test/java/org/opensearch/cluster/routing/allocation/BalanceConfigurationTests.java
@@ -751,6 +751,17 @@ private ClusterState initCluster(
         return applyAllocationUntilNoChange(clusterState, strategy);
     }
 
+    private ClusterState addNodes(ClusterState clusterState, AllocationService strategy, int numberOfNodes) {
+        logger.info("now, start [{}] more node, check that rebalancing will happen because we set it to always", numberOfNodes);
+        DiscoveryNodes.Builder nodes = DiscoveryNodes.builder(clusterState.nodes());
+        for (int i = 0; i < numberOfNodes; i++) {
+            nodes.add(newNode("node" + (clusterState.nodes().getSize() + i)));
+        }
+        clusterState = ClusterState.builder(clusterState).nodes(nodes.build()).build();
+        clusterState = strategy.reroute(clusterState, "reroute");
+        return applyStartedShardsUntilNoChange(clusterState, strategy);
+    }
+
     private ClusterState addNode(ClusterState clusterState, AllocationService strategy) {
         logger.info("now, start 1 more node, check that rebalancing will happen because we set it to always");
         clusterState = ClusterState.builder(clusterState)