Add test

Author: adarshsanjeev

server/src/main/java/org/apache/druid/server/coordinator/loading/HttpLoadQueuePeon.java

@@ -201,10 +201,9 @@ private HistoricalLoadingCapabilities fetchSegmentLoadingCapabilities()
       );
     }
     catch (Throwable th) {
-      log.error("Received error while fetching historical capabilities from Server[%s].", serverId);
       throw DruidException.forPersona(DruidException.Persona.OPERATOR)
                           .ofCategory(DruidException.Category.RUNTIME_FAILURE)
-                          .build(th, "Failed to fetch historical capabilities");
+                          .build(th, "Received error while fetching historical capabilities from Server[%s].", serverId);
     }
   }
 
@@ -216,16 +215,7 @@ private void doSegmentManagement()
     }
 
     final SegmentLoadingMode loadingMode = loadingModeSupplier.get();
-    int batchSize;
-    if (config.getBatchSize() != null) {
-      batchSize = config.getBatchSize();
-    } else if (SegmentLoadingMode.TURBO.equals(loadingMode)) {
-      batchSize = serverCapabilities.getNumTurboLoadingThreads();
-    } else if (SegmentLoadingMode.NORMAL.equals(loadingMode)) {
-      batchSize = serverCapabilities.getNumLoadingThreads();
-    } else {
-      throw DruidException.defensive().build("unsupported loading mode");
-    }
+    int batchSize = calculateBatchSize(loadingMode);
 
     final List<DataSegmentChangeRequest> newRequests = new ArrayList<>(batchSize);
 
@@ -391,6 +381,20 @@ private void logRequestFailure(Throwable t)
     }
   }
 
+  @VisibleForTesting
+  int calculateBatchSize(SegmentLoadingMode loadingMode)
+  {
+    if (config.getBatchSize() != null) {
+      return config.getBatchSize();
+    } else if (SegmentLoadingMode.TURBO.equals(loadingMode)) {
+      return serverCapabilities.getNumTurboLoadingThreads();
+    } else if (SegmentLoadingMode.NORMAL.equals(loadingMode)) {
+      return serverCapabilities.getNumLoadingThreads();
+    } else {
+      throw DruidException.defensive().build("unsupported loading mode");
+    }
+  }
+
   private void handleResponseStatus(DataSegmentChangeRequest changeRequest, SegmentChangeStatus status)
   {
     changeRequest.go(

server/src/test/java/org/apache/druid/server/coordination/SegmentLoadDropHandlerTest.java

@@ -245,7 +245,7 @@ public void testProcessBatch() throws Exception
       runnable.run();
     }
 
-    result = handler.processBatch(ImmutableList.of(new SegmentChangeRequestLoad(segment1)), SegmentLoadingMode.NORMAL).get();
+    result = handler.processBatch(ImmutableList.of(new SegmentChangeRequestLoad(segment1)), SegmentLoadingMode.TURBO).get();
     Assert.assertEquals(SegmentChangeStatus.SUCCESS, result.get(0).getStatus());
 
     Assert.assertEquals(ImmutableList.of(segment1), segmentAnnouncer.getObservedSegments());

server/src/test/java/org/apache/druid/server/coordinator/loading/HttpLoadQueuePeonTest.java

@@ -103,7 +103,7 @@ public void setUp()
         ),
         httpClient.callbackExecutor,
         () -> SegmentLoadingMode.NORMAL,
-        new HistoricalLoadingCapabilities(1, 1)
+        new HistoricalLoadingCapabilities(1, 3)
     );
     httpLoadQueuePeon.start();
   }
@@ -333,6 +333,30 @@ public void testLoadRateIsChangedWhenLoadSucceeds() throws InterruptedException
     );
   }
 
+  @Test
+  public void testBatchSize()
+  {
+    Assert.assertEquals(10, httpLoadQueuePeon.calculateBatchSize(SegmentLoadingMode.NORMAL));
+
+    httpLoadQueuePeon = new HttpLoadQueuePeon(
+        "http://dummy:4000",
+        MAPPER,
+        httpClient,
+        new HttpLoadQueuePeonConfig(null, null, null),
+        new WrappingScheduledExecutorService(
+            "HttpLoadQueuePeonTest-%s",
+            httpClient.processingExecutor,
+            true
+        ),
+        httpClient.callbackExecutor,
+        () -> SegmentLoadingMode.NORMAL,
+        new HistoricalLoadingCapabilities(1, 3)
+    );
+
+    Assert.assertEquals(1, httpLoadQueuePeon.calculateBatchSize(SegmentLoadingMode.NORMAL));
+    Assert.assertEquals(3, httpLoadQueuePeon.calculateBatchSize(SegmentLoadingMode.TURBO));
+  }
+
   private LoadPeonCallback markSegmentProcessed(DataSegment segment)
   {
     return success -> httpClient.processedSegments.add(segment);

server/src/test/java/org/apache/druid/server/http/CoordinatorDynamicConfigTest.java

@@ -545,18 +545,6 @@ public void testTurboLoadingNodes()
     Assert.assertEquals(SegmentLoadingMode.TURBO, config.getLoadingModeForServer("localhost:8083"));
   }
 
-  @Test
-  public void testTurboLoadingNodes()
-  {
-    CoordinatorDynamicConfig config = CoordinatorDynamicConfig
-        .builder()
-        .withTurboLoadingNodes(ImmutableSet.of("localhost:8083"))
-        .build();
-
-    Assert.assertEquals(SegmentLoadingMode.NORMAL, config.getLoadingModeForServer("localhost:8082"));
-    Assert.assertEquals(SegmentLoadingMode.TURBO, config.getLoadingModeForServer("localhost:8083"));
-  }
-
   @Test
   public void testEqualsAndHashCode()
   {