Add tests

Signed-off-by: Rishab Nahata <rnnahata@amazon.com>
imRishN · Aug 13, 2024 · fa0ad4e · fa0ad4e
1 parent f75154d
commit fa0ad4e
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 0 deletions.
diff --git a/server/src/test/java/org/opensearch/gateway/PrimaryShardBatchAllocatorTests.java b/server/src/test/java/org/opensearch/gateway/PrimaryShardBatchAllocatorTests.java
@@ -51,6 +51,7 @@
 import java.util.stream.Collectors;
 
 import static org.opensearch.cluster.routing.UnassignedInfo.Reason.CLUSTER_RECOVERED;
+import static org.opensearch.cluster.routing.UnassignedInfo.Reason.INDEX_CREATED;
 
 public class PrimaryShardBatchAllocatorTests extends OpenSearchAllocationTestCase {
 
@@ -322,6 +323,21 @@ public void testAllocateUnassignedBatchOnTimeoutWithNoShards() {
         assertEquals(0, ignoredShards.size());
     }
 
+    public void testAllocateUnassignedBatchOnTimeoutSkipIgnoringNewPrimaryShards() {
+        ClusterSettings clusterSettings = new ClusterSettings(Settings.EMPTY, ClusterSettings.BUILT_IN_CLUSTER_SETTINGS);
+        AllocationDeciders allocationDeciders = randomAllocationDeciders(Settings.builder().build(), clusterSettings, random());
+        setUpShards(1);
+        final RoutingAllocation routingAllocation = routingAllocationWithOnePrimary(allocationDeciders, INDEX_CREATED);
+        ShardRouting shardRouting = routingAllocation.routingTable().getIndicesRouting().get("test").shard(shardId.id()).primaryShard();
+
+        Set<ShardId> shardIds = new HashSet<>();
+        shardIds.add(shardRouting.shardId());
+        batchAllocator.allocateUnassignedBatchOnTimeout(shardIds, routingAllocation, true);
+
+        List<ShardRouting> ignoredShards = routingAllocation.routingNodes().unassigned().ignored();
+        assertEquals(0, ignoredShards.size());
+    }
+
     private RoutingAllocation routingAllocationWithOnePrimary(
         AllocationDeciders deciders,
         UnassignedInfo.Reason reason,

diff --git a/server/src/test/java/org/opensearch/gateway/ReplicaShardBatchAllocatorTests.java b/server/src/test/java/org/opensearch/gateway/ReplicaShardBatchAllocatorTests.java
@@ -744,6 +744,24 @@ public void testAllocateUnassignedBatchOnTimeoutWithAlreadyRecoveringReplicaShar
         assertThat(allocation.routingNodes().unassigned().ignored().size(), equalTo(0));
     }
 
+    public void testAllocateUnassignedBatchOnTimeoutSkipIgnoringNewReplicaShards() {
+        RoutingAllocation allocation = onePrimaryOnNode1And1Replica(
+            yesAllocationDeciders(),
+            Settings.EMPTY,
+            UnassignedInfo.Reason.INDEX_CREATED
+        );
+        final RoutingNodes.UnassignedShards.UnassignedIterator iterator = allocation.routingNodes().unassigned().iterator();
+        Set<ShardId> shards = new HashSet<>();
+        while (iterator.hasNext()) {
+            ShardRouting sr = iterator.next();
+            if (sr.primary() == false) {
+                shards.add(sr.shardId());
+            }
+        }
+        testBatchAllocator.allocateUnassignedBatchOnTimeout(shards, allocation, false);
+        assertThat(allocation.routingNodes().unassigned().ignored().size(), equalTo(0));
+    }
+
     private RoutingAllocation onePrimaryOnNode1And1Replica(AllocationDeciders deciders) {
         return onePrimaryOnNode1And1Replica(deciders, Settings.EMPTY, UnassignedInfo.Reason.CLUSTER_RECOVERED);
     }