Add code coverage tests for New Index creation in PrimaryShardBatchAl…

…locator Signed-off-by: Pranshu Shukla <[email protected]>
opensearch-project · May 17, 2024 · 8139e64 · 8139e64
1 parent 4700be3
commit 8139e64
Showing 1 changed file with 43 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
@@ -111,6 +111,49 @@ public void testMakeAllocationDecisionForReplicaShard() {
         assertFalse(allDecisions.get(replicaShards.get(0)).isDecisionTaken());
     }
 
+    public void testDecisionTakenForNewIndexShards() {
+
+        // Simulating for unassigned info = EMPTY_STORE (In the case of new index creations)
+        final ShardId shard1 = new ShardId("test", "_na_", 0);
+        final ShardId shard2 = new ShardId("test", "_na_", 1);
+
+        // Creating metadata for such a scenario.
+        Metadata metadata = Metadata.builder()
+            .put(
+                IndexMetadata.builder(shard1.getIndexName())
+                    .settings(settings(Version.CURRENT))
+                    .numberOfShards(2)
+                    .numberOfReplicas(1)
+                    .putInSyncAllocationIds(shard1.id(), Sets.newHashSet())
+                    .putInSyncAllocationIds(shard2.id(), Sets.newHashSet())
+            )
+            .build();
+
+        RoutingTable.Builder routingTableBuilder = RoutingTable.builder();
+        routingTableBuilder.addAsNew(metadata.index(shard1.getIndex()));
+
+        ClusterState state = ClusterState.builder(org.opensearch.cluster.ClusterName.CLUSTER_NAME_SETTING.getDefault(Settings.EMPTY))
+            .metadata(metadata)
+            .routingTable(routingTableBuilder.build())
+            .nodes(DiscoveryNodes.builder().add(node1).add(node2).add(node3))
+            .build();
+        RoutingAllocation newIndexRoutingAllocation = new RoutingAllocation(noAllocationDeciders(), new RoutingNodes(state, false), state, null, null, System.nanoTime());
+
+        ShardRouting primaryShard1 = newIndexRoutingAllocation.routingTable().getIndicesRouting().get("test").shard(shard1.id()).primaryShard();
+        ShardRouting primaryShard2 = newIndexRoutingAllocation.routingTable().getIndicesRouting().get("test").shard(shard2.id()).primaryShard();
+
+        List<ShardRouting> shards = new ArrayList<>();
+        shards.add(primaryShard1);
+        shards.add(primaryShard2);
+
+        HashMap<ShardRouting, AllocateUnassignedDecision> allDecisions = batchAllocator.makeAllocationDecision(shards, newIndexRoutingAllocation, logger);
+        // verify we get decisions for all the shards
+        assertEquals(shards.size(), allDecisions.size());
+        assertEquals(shards, new ArrayList<>(allDecisions.keySet()));
+        assertFalse(allDecisions.get(shards.get(0)).isDecisionTaken());
+        assertFalse(allDecisions.get(shards.get(1)).isDecisionTaken());
+    }
+
     public void testMakeAllocationDecisionDataFetched() {
         final RoutingAllocation allocation = routingAllocationWithOnePrimary(noAllocationDeciders(), CLUSTER_RECOVERED, "allocId1");