From d19c0cd612704b771044ae28ac9a15e2ceef99fa Mon Sep 17 00:00:00 2001
From: Rishab Nahata <rnnahata@amazon.com>
Date: Wed, 4 Sep 2024 13:29:25 +0530
Subject: [PATCH] Minor refactor

Signed-off-by: Rishab Nahata <rnnahata@amazon.com>
---
 .../gateway/ShardsBatchGatewayAllocator.java    | 17 +++--------------
 .../gateway/GatewayAllocatorTests.java          |  2 +-
 .../gateway/TestShardBatchGatewayAllocator.java |  6 +++---
 3 files changed, 7 insertions(+), 18 deletions(-)

diff --git a/server/src/main/java/org/opensearch/gateway/ShardsBatchGatewayAllocator.java b/server/src/main/java/org/opensearch/gateway/ShardsBatchGatewayAllocator.java
index 116a7f94a512e..5e2dcbcd70b40 100644
--- 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
index 3222d19ecd066..61b53d688d330 100644
--- 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
index c134c9ca72d5d..c2ff228a6bf3a 100644
--- 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<>();