Fix flaky tests fromSegmentReplicationAllocationIT

server/src/internalClusterTest/java/org/opensearch/indices/replication/SegmentReplicationAllocationIT.java

@@ -25,6 +25,7 @@
 import org.opensearch.test.junit.annotations.TestLogging;
 
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.TimeUnit;
@@ -169,14 +170,16 @@ public void testSingleIndexShardAllocation() throws Exception {
 
         // Remove a node
         internalCluster().stopRandomNode(InternalTestCluster.nameFilter(nodeNames.get(0)));
-        ensureGreen(TimeValue.timeValueSeconds(60));
+        internalCluster().validateClusterFormed();
+        ensureGreen(TimeValue.timeValueSeconds(100));
         state = client().admin().cluster().prepareState().execute().actionGet().getState();
         logger.info(ShardAllocations.printShardDistribution(state));
         verifyPerIndexPrimaryBalance();
 
         // Add a new node
         internalCluster().startDataOnlyNode();
-        ensureGreen(TimeValue.timeValueSeconds(60));
+        internalCluster().validateClusterFormed();
+        ensureGreen(TimeValue.timeValueSeconds(100));
         state = client().admin().cluster().prepareState().execute().actionGet().getState();
         logger.info(ShardAllocations.printShardDistribution(state));
         verifyPerIndexPrimaryBalance();
@@ -250,24 +253,32 @@ public void testAllocationAndRebalanceWithDisruption() throws Exception {
         internalCluster().startClusterManagerOnlyNode();
         final int maxReplicaCount = 2;
         final int maxShardCount = 2;
-        // Create higher number of nodes than number of shards to reduce chances of SameShardAllocationDecider kicking-in
+        final int numberOfIndices = randomIntBetween(1, 3);
+        final int maxPossibleShards = numberOfIndices * maxShardCount * (1 + maxReplicaCount);
+
+        List<List<Integer>> shardAndReplicaCounts = new ArrayList<>();
+        int shardCount, replicaCount, totalShards = 0;
+        for (int i = 0; i < numberOfIndices; i++) {
+            shardCount = randomIntBetween(1, maxShardCount);
+            replicaCount = randomIntBetween(1, maxReplicaCount);
+            shardAndReplicaCounts.add(Arrays.asList(shardCount, replicaCount));
+            totalShards += shardCount * (1 + replicaCount);
+        }
+        // Create a strictly higher number of nodes than the number of shards to reduce chances of SameShardAllocationDecider kicking-in
         // and preventing primary relocations
-        final int nodeCount = randomIntBetween(5, 10);
-        final int numberOfIndices = randomIntBetween(1, 10);
+        final int nodeCount = randomIntBetween(totalShards, maxPossibleShards) + 1;
         final float buffer = randomIntBetween(1, 4) * 0.10f;
-
         logger.info("--> Creating {} nodes", nodeCount);
         final List<String> nodeNames = new ArrayList<>();
         for (int i = 0; i < nodeCount; i++) {
             nodeNames.add(internalCluster().startNode());
         }
         setAllocationRelocationStrategy(true, true, buffer);
 
-        int shardCount, replicaCount;
         ClusterState state;
         for (int i = 0; i < numberOfIndices; i++) {
-            shardCount = randomIntBetween(1, maxShardCount);
-            replicaCount = randomIntBetween(1, maxReplicaCount);
+            shardCount = shardAndReplicaCounts.get(i).get(0);
+            replicaCount = shardAndReplicaCounts.get(i).get(1);
             logger.info("--> Creating index test{} with primary {} and replica {}", i, shardCount, replicaCount);
             createIndex("test" + i, shardCount, replicaCount, i % 2 == 0);
             ensureGreen(TimeValue.timeValueSeconds(60));