Add test

Signed-off-by: Rishab Nahata <[email protected]>
imRishN · Jul 30, 2024 · ee8e7c5 · ee8e7c5
1 parent 89cc56d
commit ee8e7c5
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 2 deletions.
diff --git a/server/src/main/java/org/opensearch/gateway/ShardsBatchGatewayAllocator.java b/server/src/main/java/org/opensearch/gateway/ShardsBatchGatewayAllocator.java
@@ -794,11 +794,20 @@ public int getNumberOfStoreShardBatches() {
         return batchIdToStoreShardBatch.size();
     }
 
-    private void setPrimaryBatchAllocatorTimeout(TimeValue primaryShardsBatchGatewayAllocatorTimeout) {
+    protected void setPrimaryBatchAllocatorTimeout(TimeValue primaryShardsBatchGatewayAllocatorTimeout) {
         this.primaryShardsBatchGatewayAllocatorTimeout = primaryShardsBatchGatewayAllocatorTimeout;
     }
 
-    private void setReplicaBatchAllocatorTimeout(TimeValue replicaShardsBatchGatewayAllocatorTimeout) {
+    protected void setReplicaBatchAllocatorTimeout(TimeValue replicaShardsBatchGatewayAllocatorTimeout) {
         this.replicaShardsBatchGatewayAllocatorTimeout = replicaShardsBatchGatewayAllocatorTimeout;
     }
+
+    // for tests
+    public Set<ShardId> getTimedOutPrimaryShardIds() {
+        return timedOutPrimaryShardIds;
+    }
+
+    public Set<ShardId> getTimedOutReplicaShardIds() {
+        return timedOutReplicaShardIds;
+    }
 }
diff --git a/server/src/test/java/org/opensearch/gateway/GatewayAllocatorTests.java b/server/src/test/java/org/opensearch/gateway/GatewayAllocatorTests.java
@@ -32,6 +32,7 @@
 import org.opensearch.cluster.routing.allocation.decider.AllocationDeciders;
 import org.opensearch.common.collect.Tuple;
 import org.opensearch.common.settings.Settings;
+import org.opensearch.common.unit.TimeValue;
 import org.opensearch.common.util.BatchRunnableExecutor;
 import org.opensearch.common.util.set.Sets;
 import org.opensearch.core.index.shard.ShardId;
@@ -368,6 +369,18 @@ public void testCreatePrimaryAndReplicaExecutorOfSizeTwo() {
         assertEquals(executor.getTimeoutAwareRunnables().size(), 2);
     }
 
+    public void testCollectTimedOutShards() {
+        createIndexAndUpdateClusterState(2, 50, 2);
+        testShardsBatchGatewayAllocator.setPrimaryBatchAllocatorTimeout(TimeValue.ZERO);
+        testShardsBatchGatewayAllocator.setReplicaBatchAllocatorTimeout(TimeValue.ZERO);
+        BatchRunnableExecutor executor = testShardsBatchGatewayAllocator.allocateAllUnassignedShards(testAllocation, true);
+        executor.run();
+        assertEquals(100, testShardsBatchGatewayAllocator.getTimedOutPrimaryShardIds().size());
+        executor = testShardsBatchGatewayAllocator.allocateAllUnassignedShards(testAllocation, false);
+        executor.run();
+        assertEquals(100, testShardsBatchGatewayAllocator.getTimedOutReplicaShardIds().size());
+    }
+
     private void createIndexAndUpdateClusterState(int count, int numberOfShards, int numberOfReplicas) {
         if (count == 0) return;
         Metadata.Builder metadata = Metadata.builder();