Added UT to test new batches are created for failed shards

Signed-off-by: Shivansh Arora <shivansh.arora@protonmail.com>
shiv0408 · Oct 17, 2023 · 354a945 · 354a945
1 parent e35add0
commit 354a945
Show file tree

Hide file tree

Showing 3 changed files with 350 additions and 3 deletions.
diff --git a/server/src/main/java/org/opensearch/gateway/GatewayAllocator.java b/server/src/main/java/org/opensearch/gateway/GatewayAllocator.java
@@ -76,8 +76,8 @@ public class GatewayAllocator implements ExistingShardsAllocator {
 
     private final RerouteService rerouteService;
 
-    private final PrimaryShardAllocator primaryShardAllocator;
-    private final ReplicaShardAllocator replicaShardAllocator;
+    private PrimaryShardAllocator primaryShardAllocator;
+    private ReplicaShardAllocator replicaShardAllocator;
 
     private final ConcurrentMap<
         ShardId,
@@ -113,6 +113,17 @@ protected GatewayAllocator() {
         this.replicaShardAllocator = null;
     }
 
+    // for tests
+    protected void setShardAllocators(PrimaryShardAllocator primaryShardAllocator,
+                                 ReplicaShardAllocator replicaShardAllocator,
+                                 PrimaryShardBatchAllocator primaryShardBatchAllocator,
+                                 ReplicaShardBatchAllocator replicaBatchShardAllocator) {
+        this.primaryShardAllocator = primaryShardAllocator;
+        this.replicaShardAllocator = replicaShardAllocator;
+        this.primaryBatchShardAllocator = primaryShardBatchAllocator;
+        this.replicaBatchShardAllocator = replicaBatchShardAllocator;
+    }
+
     @Override
     public int getNumberOfInFlightFetches() {
         int count = 0;
@@ -190,7 +201,6 @@ public AllocateUnassignedDecision explainUnassignedShardAllocation(ShardRouting
         assert routingAllocation.debugDecision();
         boolean batchMode = routingAllocation.nodes().getMinNodeVersion().onOrAfter(Version.CURRENT);
         if (batchMode) {
-            // TODO add integ test for testing this behaviour when shard is unassigned but failed many times and is ultimately removed from batch
             if (getBatchId(unassignedShard, unassignedShard.primary()) == null) {
                 createAndUpdateBatches(routingAllocation, unassignedShard.primary());
             }