Refactor testBalanceIncrementallyStartNodes (#89733)

Currently, testBalanceIncrementallyStartNodes relies on exactly 2 reroutes to
start all initializing shards after adding node-3. This might change with a
different shard allocator.

server/src/test/java/org/elasticsearch/cluster/routing/allocation/IndexBalanceTests.java

@@ -265,41 +265,24 @@ public class IndexBalanceTests extends ESAllocationTestCase {
             assertThat(clusterState.routingTable().index("test1").shard(i).replicaShards().get(0).state(), equalTo(STARTED));
         }
 
-        logger.info("Add another node and perform rerouting, relocate shards to new node");
+        logger.info("Add another node and perform reroute to relocate shards to the new node");
         clusterState = ClusterState.builder(clusterState).nodes(DiscoveryNodes.builder(clusterState.nodes()).add(newNode("node3"))).build();
-        newState = strategy.reroute(clusterState, "reroute");
-        assertThat(newState, not(equalTo(clusterState)));
-        clusterState = newState;
-
-        logger.info("Reroute, nothing should change");
-        newState = strategy.reroute(clusterState, "reroute");
-        assertThat(newState, equalTo(clusterState));
-
-        logger.info("Start the backup shard");
-        newState = startInitializingShardsAndReroute(strategy, clusterState);
-        assertThat(newState, not(equalTo(clusterState)));
-        clusterState = newState;
+        clusterState = applyStartedShardsUntilNoChange(clusterState, strategy);
 
         assertThat(clusterState.routingTable().index("test").size(), equalTo(3));
-
-        newState = startInitializingShardsAndReroute(strategy, clusterState);
-        assertThat(newState, not(equalTo(clusterState)));
-        clusterState = newState;
-        RoutingNodes routingNodes = clusterState.getRoutingNodes();
-
         assertThat(clusterState.routingTable().index("test1").size(), equalTo(3));
 
-        assertThat(routingNodes.node("node1").numberOfShardsWithState(STARTED), equalTo(4));
-        assertThat(routingNodes.node("node2").numberOfShardsWithState(STARTED), equalTo(4));
-        assertThat(routingNodes.node("node3").numberOfShardsWithState(STARTED), equalTo(4));
+        assertThat(clusterState.getRoutingNodes().node("node1").numberOfShardsWithState(STARTED), equalTo(4));
+        assertThat(clusterState.getRoutingNodes().node("node2").numberOfShardsWithState(STARTED), equalTo(4));
+        assertThat(clusterState.getRoutingNodes().node("node3").numberOfShardsWithState(STARTED), equalTo(4));
 
-        assertThat(routingNodes.node("node1").shardsWithState("test", STARTED).size(), equalTo(2));
-        assertThat(routingNodes.node("node2").shardsWithState("test", STARTED).size(), equalTo(2));
-        assertThat(routingNodes.node("node3").shardsWithState("test", STARTED).size(), equalTo(2));
+        assertThat(clusterState.getRoutingNodes().node("node1").shardsWithState("test", STARTED).size(), equalTo(2));
+        assertThat(clusterState.getRoutingNodes().node("node2").shardsWithState("test", STARTED).size(), equalTo(2));
+        assertThat(clusterState.getRoutingNodes().node("node3").shardsWithState("test", STARTED).size(), equalTo(2));
 
-        assertThat(routingNodes.node("node1").shardsWithState("test1", STARTED).size(), equalTo(2));
-        assertThat(routingNodes.node("node2").shardsWithState("test1", STARTED).size(), equalTo(2));
-        assertThat(routingNodes.node("node3").shardsWithState("test1", STARTED).size(), equalTo(2));
+        assertThat(clusterState.getRoutingNodes().node("node1").shardsWithState("test1", STARTED).size(), equalTo(2));
+        assertThat(clusterState.getRoutingNodes().node("node2").shardsWithState("test1", STARTED).size(), equalTo(2));
+        assertThat(clusterState.getRoutingNodes().node("node3").shardsWithState("test1", STARTED).size(), equalTo(2));
     }
 
     public void testBalanceAllNodesStartedAddIndex() {

server/src/test/java/org/elasticsearch/cluster/routing/allocation/RoutingNodesIntegrityTests.java

@@ -42,7 +42,6 @@ import static org.elasticsearch.cluster.metadata.IndexMetadata.SETTING_NUMBER_OF
 import static org.elasticsearch.cluster.routing.ShardRoutingState.INITIALIZING;
 import static org.elasticsearch.cluster.routing.ShardRoutingState.STARTED;
 import static org.hamcrest.Matchers.equalTo;
-import static org.hamcrest.Matchers.not;
 import static org.hamcrest.Matchers.oneOf;
 
 public class RoutingNodesIntegrityTests extends ESAllocationTestCase {
@@ -142,59 +141,39 @@ public class RoutingNodesIntegrityTests extends ESAllocationTestCase {
             org.elasticsearch.cluster.ClusterName.CLUSTER_NAME_SETTING.getDefault(Settings.EMPTY)
         ).metadata(metadata).routingTable(initialRoutingTable).build();
 
-        logger.info("Adding one node and performing rerouting");
+        logger.info("Adding node-1 and performing reroute");
         clusterState = ClusterState.builder(clusterState).nodes(DiscoveryNodes.builder().add(newNode("node1"))).build();
-
         clusterState = strategy.reroute(clusterState, "reroute");
 
-        logger.info("Add another node and perform rerouting, nothing will happen since primary not started");
+        logger.info("Add node-2 and perform reroute, nothing will happen since primary not started");
         clusterState = ClusterState.builder(clusterState).nodes(DiscoveryNodes.builder(clusterState.nodes()).add(newNode("node2"))).build();
         clusterState = strategy.reroute(clusterState, "reroute");
 
-        logger.info("Start the primary shard");
-        clusterState = startInitializingShardsAndReroute(strategy, clusterState);
-
-        logger.info("Reroute, nothing should change");
-        clusterState = strategy.reroute(clusterState, "reroute");
-
-        logger.info("Start the backup shard");
-        clusterState = startInitializingShardsAndReroute(strategy, clusterState);
+        logger.info("Start the all shards");
+        clusterState = startInitializingShardsAndReroute(strategy, clusterState); // primaries
+        clusterState = startInitializingShardsAndReroute(strategy, clusterState); // replicas
 
-        clusterState = startInitializingShardsAndReroute(strategy, clusterState);
-
-        logger.info("Add another node and perform rerouting, relocate shards to new node");
+        logger.info("Add node-3 and perform reroute, relocate shards to new node");
         clusterState = ClusterState.builder(clusterState).nodes(DiscoveryNodes.builder(clusterState.nodes()).add(newNode("node3"))).build();
         clusterState = strategy.reroute(clusterState, "reroute");
 
-        logger.info("Reroute, nothing should change");
-        ClusterState newState = strategy.reroute(clusterState, "reroute");
-        assertThat(newState, equalTo(clusterState));
-
-        logger.info("Start the backup shard");
-        newState = startInitializingShardsAndReroute(strategy, clusterState);
-        assertThat(newState, not(equalTo(clusterState)));
-        clusterState = newState;
+        logger.info("Await all shards reallocate");
+        clusterState = applyStartedShardsUntilNoChange(clusterState, strategy);
 
         assertThat(clusterState.routingTable().index("test").size(), equalTo(3));
-
-        newState = startInitializingShardsAndReroute(strategy, clusterState);
-        assertThat(newState, not(equalTo(clusterState)));
-        clusterState = newState;
-        RoutingNodes routingNodes = clusterState.getRoutingNodes();
-
         assertThat(clusterState.routingTable().index("test1").size(), equalTo(3));
 
-        assertThat(routingNodes.node("node1").numberOfShardsWithState(STARTED), equalTo(4));
-        assertThat(routingNodes.node("node2").numberOfShardsWithState(STARTED), equalTo(4));
-        assertThat(routingNodes.node("node3").numberOfShardsWithState(STARTED), equalTo(4));
+        assertThat(clusterState.getRoutingNodes().node("node1").numberOfShardsWithState(STARTED), equalTo(4));
+        assertThat(clusterState.getRoutingNodes().node("node2").numberOfShardsWithState(STARTED), equalTo(4));
+        assertThat(clusterState.getRoutingNodes().node("node3").numberOfShardsWithState(STARTED), equalTo(4));
 
-        assertThat(routingNodes.node("node1").shardsWithState("test", STARTED).size(), equalTo(2));
-        assertThat(routingNodes.node("node2").shardsWithState("test", STARTED).size(), equalTo(2));
-        assertThat(routingNodes.node("node3").shardsWithState("test", STARTED).size(), equalTo(2));
+        assertThat(clusterState.getRoutingNodes().node("node1").shardsWithState("test", STARTED).size(), equalTo(2));
+        assertThat(clusterState.getRoutingNodes().node("node2").shardsWithState("test", STARTED).size(), equalTo(2));
+        assertThat(clusterState.getRoutingNodes().node("node3").shardsWithState("test", STARTED).size(), equalTo(2));
 
-        assertThat(routingNodes.node("node1").shardsWithState("test1", STARTED).size(), equalTo(2));
-        assertThat(routingNodes.node("node2").shardsWithState("test1", STARTED).size(), equalTo(2));
-        assertThat(routingNodes.node("node3").shardsWithState("test1", STARTED).size(), equalTo(2));
+        assertThat(clusterState.getRoutingNodes().node("node1").shardsWithState("test1", STARTED).size(), equalTo(2));
+        assertThat(clusterState.getRoutingNodes().node("node2").shardsWithState("test1", STARTED).size(), equalTo(2));
+        assertThat(clusterState.getRoutingNodes().node("node3").shardsWithState("test1", STARTED).size(), equalTo(2));
     }
 
     public void testBalanceAllNodesStartedAddIndex() {