Minor refactor

Signed-off-by: Rishab Nahata <rnnahata@amazon.com>
imRishN · Sep 4, 2024 · d19c0cd · d19c0cd
1 parent e2da98f
commit d19c0cd
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 18 deletions.
diff --git a/server/src/main/java/org/opensearch/gateway/ShardsBatchGatewayAllocator.java b/server/src/main/java/org/opensearch/gateway/ShardsBatchGatewayAllocator.java
@@ -184,27 +184,16 @@ public void cleanCaches() {
 
     // for tests
     protected ShardsBatchGatewayAllocator() {
-        this(DEFAULT_SHARD_BATCH_SIZE);
+        this(DEFAULT_SHARD_BATCH_SIZE, null);
     }
 
-    protected ShardsBatchGatewayAllocator(long batchSize) {
-        this.rerouteService = null;
-        this.batchStartedAction = null;
-        this.primaryShardBatchAllocator = null;
-        this.batchStoreAction = null;
-        this.replicaShardBatchAllocator = null;
-        this.maxBatchSize = batchSize;
-        this.primaryShardsBatchGatewayAllocatorTimeout = null;
-        this.replicaShardsBatchGatewayAllocatorTimeout = null;
-    }
-
-    protected ShardsBatchGatewayAllocator(RerouteService rerouteService) {
+    protected ShardsBatchGatewayAllocator(long batchSize, RerouteService rerouteService) {
         this.rerouteService = rerouteService;
         this.batchStartedAction = null;
         this.primaryShardBatchAllocator = null;
         this.batchStoreAction = null;
         this.replicaShardBatchAllocator = null;
-        this.maxBatchSize = DEFAULT_SHARD_BATCH_SIZE;
+        this.maxBatchSize = batchSize;
         this.primaryShardsBatchGatewayAllocatorTimeout = null;
         this.replicaShardsBatchGatewayAllocatorTimeout = null;
     }

diff --git a/server/src/test/java/org/opensearch/gateway/GatewayAllocatorTests.java b/server/src/test/java/org/opensearch/gateway/GatewayAllocatorTests.java
@@ -444,7 +444,7 @@ public void testCollectTimedOutShardsAndScheduleReroute() throws InterruptedExce
             rerouteLatch.countDown();
         };
         CountDownLatch timedOutShardsLatch = new CountDownLatch(20);
-        testShardsBatchGatewayAllocator = new TestShardBatchGatewayAllocator(timedOutShardsLatch, rerouteService);
+        testShardsBatchGatewayAllocator = new TestShardBatchGatewayAllocator(timedOutShardsLatch, 1000, rerouteService);
         testShardsBatchGatewayAllocator.setPrimaryBatchAllocatorTimeout(TimeValue.ZERO);
         testShardsBatchGatewayAllocator.setReplicaBatchAllocatorTimeout(TimeValue.ZERO);
         BatchRunnableExecutor executor = testShardsBatchGatewayAllocator.allocateAllUnassignedShards(testAllocation, true);

diff --git a/test/framework/src/main/java/org/opensearch/test/gateway/TestShardBatchGatewayAllocator.java b/test/framework/src/main/java/org/opensearch/test/gateway/TestShardBatchGatewayAllocator.java
@@ -40,13 +40,13 @@ public TestShardBatchGatewayAllocator() {
 
     }
 
-    public TestShardBatchGatewayAllocator(CountDownLatch latch, RerouteService rerouteService) {
-        super(rerouteService);
+    public TestShardBatchGatewayAllocator(CountDownLatch latch, long maxBatchSize, RerouteService rerouteService) {
+        super(maxBatchSize, rerouteService);
         this.latch = latch;
     }
 
     public TestShardBatchGatewayAllocator(long maxBatchSize) {
-        super(maxBatchSize);
+        super(maxBatchSize, null);
     }
 
     Map<String /* node id */, Map<ShardId, ShardRouting>> knownAllocations = new HashMap<>();