Cover new method in junit test

linkedin · Nov 26, 2024 · 7eb541d · 7eb541d
1 parent eb899fa
commit 7eb541d
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 0 deletions.
diff --git a/...control/src/main/java/com/linkedin/kafka/cruisecontrol/executor/ExecutionTaskPlanner.java b/...control/src/main/java/com/linkedin/kafka/cruisecontrol/executor/ExecutionTaskPlanner.java
@@ -568,4 +568,11 @@ Map<Integer, Integer> getSortedBrokerIdToInterBrokerMoveTaskCountMap() {
         ));
     return resultMap;
   }
+
+  /*
+   * Package private for testing.
+   */
+  Map<Integer, SortedSet<ExecutionTask>> getInterPartMoveTasksByBrokerId() {
+    return _interPartMoveTasksByBrokerId;
+  }
 }
diff --git a/...rol/src/test/java/com/linkedin/kafka/cruisecontrol/executor/ExecutionTaskPlannerTest.java b/...rol/src/test/java/com/linkedin/kafka/cruisecontrol/executor/ExecutionTaskPlannerTest.java
@@ -24,6 +24,7 @@
 import java.util.HashSet;
 import java.util.Properties;
 import java.util.Set;
+import java.util.SortedSet;
 import org.apache.kafka.clients.admin.AdminClient;
 import org.apache.kafka.clients.admin.DescribeReplicaLogDirsResult;
 import org.apache.kafka.common.Cluster;
@@ -353,6 +354,13 @@ public void testGetInterBrokerPartitionMovementWithMinIsrTasks() {
     readyBrokers.put(4, 5);
     readyBrokers.put(5, 6);
     prioritizeOneAboveMinIsrMovementPlanner.addExecutionProposals(proposals, strategyOptions, null);
+    Map<Integer, Integer> countMap = prioritizeOneAboveMinIsrMovementPlanner.getSortedBrokerIdToInterBrokerMoveTaskCountMap();
+    Map<Integer, SortedSet<ExecutionTask>> taskMap = prioritizeOneAboveMinIsrMovementPlanner.getInterPartMoveTasksByBrokerId();
+    for (Map.Entry<Integer, Integer> entry : countMap.entrySet()) {
+      int brokerId = entry.getKey();
+      int count = entry.getValue();
+      assertEquals(taskMap.get(brokerId).size(), count);
+    }
     List<ExecutionTask> partitionMovementTasks
         = prioritizeOneAboveMinIsrMovementPlanner.getInterBrokerReplicaMovementTasks(readyBrokers, Collections.emptySet(), _defaultPartitionsMaxCap);
     assertEquals("First task", _rf4PartitionMovement2, partitionMovementTasks.get(0).proposal());