Repository: skiptests/astraea
Pull request #182 (2022-01-06 11:46:07)

Fix partitionScore
這支PR修正了 partitionScore讀不到 internel topics的問題，以及新增argument，使其可以決定是否要在計算分數時忽略internel topics，以及是否要輸出已經balance之Βroker

qoo332001: 原本只能讀到partition的leader，更新後已可以使用broker中所有的replica來計算分數
目前輸出長這樣
```
$ java -jar app-0.0.1-SNAPSHOT-all.jar score --bootstrap.servers 192.168.103.39:19217

broker: 1001
test0-0: 113.98
test0-1: 127.33
test0-3: 135.33
test1-1: 147.78
test1-2: 126.98
test1-3: 146.55
test2-0: 126.09
test2-1: 115.84

broker: 1002
test0-0: 8.01
test0-2: 8.69
test1-0: 7.13
test1-1: 10.34
test1-3: 10.26
test2-0: 8.84
test2-2: 8.41
test2-3: 9.45

The following brokers are balanced: [1003]
The following partitions are balanced: [__consumer_offsets-0, __consumer_offsets-1, __consumer_offsets-2, __consumer_offsets-3, __consumer_offsets-4, __consumer_offsets-5, __consumer_offsets-6, __consumer_offsets-7, __consumer_offsets-8, __consumer_offsets-9, __consumer_offsets-10, __consumer_offsets-11, __consumer_offsets-12, __consumer_offsets-13, __consumer_offsets-14, __consumer_offsets-15, __consumer_offsets-16, __consumer_offsets-17, __consumer_offsets-18, __consumer_offsets-19, __consumer_offsets-20, __consumer_offsets-21, __consumer_offsets-22, __consumer_offsets-23, __consumer_offsets-24, __consumer_offsets-25, __consumer_offsets-26, __consumer_offsets-27, __consumer_offsets-28, __consumer_offsets-29, __consumer_offsets-30, __consumer_offsets-31, __consumer_offsets-32, __consumer_offsets-33, __consumer_offsets-34, __consumer_offsets-35, __consumer_offsets-36, __consumer_offsets-37, __consumer_offsets-38, __consumer_offsets-39, __consumer_offsets-40, __consumer_offsets-41, __consumer_offsets-42, __consumer_offsets-43, __consumer_offsets-44, __consumer_offsets-45, __consumer_offsets-46, __consumer_offsets-47, __consumer_offsets-48, __consumer_offsets-49, test3-0, test3-0, test3-1, test3-2, test3-3]
```
忽略internal topics:
```
$ java -jar app-0.0.1-SNAPSHOT-all.jar score --bootstrap.servers 192.168.103.39:19217 --exclude.internal.topics

broker: 1001
test0-0: 18.69
test0-1: 32.78
test0-3: 41.22
test1-1: 54.36
test1-2: 32.41
test1-3: 53.06
test2-0: 31.47
test2-1: 20.65

broker: 1002
test0-0: 1.23
test0-2: 1.6
test1-0: 0.75
test1-1: 2.51
test1-3: 2.46
test2-0: 1.69
test2-2: 1.45
test2-3: 2.02

The following brokers are balanced: [1003]
The following partitions are balanced: [test3-0, test3-0, test3-1, test3-2, test3-3]
```
chia7712: 現在的畫面還不錯，不過我在想是否有機會更延伸一點讓我們的算分程式輸出更有"linux"的味道，#240接著討論

Patch:
diff --git a/README.md b/README.md
index 68505e1a..20d5a504 100644
--- a/README.md
+++ b/README.md
@@ -353,6 +353,8 @@ This tool will score the partition on brokers, the higher score the heavier load
 ### Partition Score Configurations
 
 1. --bootstrap.servers: the server to connect to
+2. --exclude.internal.topics: True if you want to ignore internal topics like _consumer_offsets while counting score. 
+3. --hide.balanced: True if you want to hide topics and partitions thar already balanced.:q
 
 ## Kafka Replica Syncing Monitor
 
diff --git a/app/src/main/java/org/astraea/topic/cost/CalculateUtils.java b/app/src/main/java/org/astraea/topic/cost/CalculateUtils.java
index e7204bb9..b6f84ac8 100644
--- a/app/src/main/java/org/astraea/topic/cost/CalculateUtils.java
+++ b/app/src/main/java/org/astraea/topic/cost/CalculateUtils.java
@@ -1,7 +1,8 @@
 package org.astraea.topic.cost;
 
-import java.util.HashMap;
-import java.util.Map;
+import java.util.*;
+import java.util.function.Function;
+import java.util.stream.Collectors;
 import org.apache.kafka.common.TopicPartition;
 
 public class CalculateUtils {
@@ -9,20 +10,29 @@ public class CalculateUtils {
       Map<Integer, Map<TopicPartition, Integer>> brokerPartitionSize,
       Map<String, Integer> retentionMillis) {
     Map<Integer, Map<TopicPartition, Double>> brokerPartitionLoad = new HashMap<>();
-    double load;
-    for (var broker : brokerPartitionSize.keySet()) {
-      Map<TopicPartition, Double> partitionLoad = new HashMap<>();
-      for (var partition : brokerPartitionSize.get(broker).keySet()) {
-        load =
-            (double) brokerPartitionSize.get(broker).get(partition)
-                / retentionMillis.get(partition.topic());
-        partitionLoad.put(partition, load);
-      }
-      brokerPartitionLoad.put(broker, partitionLoad);
-    }
+
+    brokerPartitionSize
+        .keySet()
+        .forEach(
+            (broker) -> {
+              Map<TopicPartition, Double> partitionLoad =
+                  brokerPartitionSize.get(broker).keySet().stream()
+                      .filter(partition -> retentionMillis.containsKey(partition.topic()))
+                      .collect(
+                          Collectors.toMap(
+                              Function.identity(),
+                              partition ->
+                                  (double) brokerPartitionSize.get(broker).get(partition)
+                                      / retentionMillis.get(partition.topic())));
+              brokerPartitionLoad.put(broker, partitionLoad);
+            });
     return brokerPartitionLoad;
   }
 
+  private static Double countSum(Set<Double> in) {
+    return in.stream().mapToDouble(i -> i).sum();
+  }
+
   public static Map<Integer, Map<TopicPartition, Double>> getScore(
       Map<Integer, Map<TopicPartition, Double>> load) {
     Map<Integer, Double> brokerLoad = new HashMap<>();
@@ -30,50 +40,75 @@ public class CalculateUtils {
     Map<Integer, Double> partitionSD = new HashMap<>();
     Map<Integer, Double> partitionMean = new HashMap<>();
     Map<Integer, Map<TopicPartition, Double>> brokerPartitionScore = new HashMap<>();
+    double brokerSD;
+    Set<Double> loadSet = new HashSet<>();
+    Set<Double> LoadSQR = new HashSet<>();
 
-    double mean = 0f, LoadSQR = 0f, SD, brokerSD;
-    int partitionNum;
-    for (var broker : load.keySet()) {
-      for (var topicPartition : load.get(broker).keySet()) {
-        mean += load.get(broker).get(topicPartition);
-        partitionLoad.put(topicPartition, load.get(broker).get(topicPartition));
-        LoadSQR += Math.pow(load.get(broker).get(topicPartition), 2);
-      }
-      partitionNum = load.get(broker).keySet().size();
-      brokerLoad.put(broker, mean);
-      mean /= partitionNum;
-      partitionMean.put(broker, mean);
-      SD = Math.pow((LoadSQR - mean * mean * partitionNum) / partitionNum, 0.5);
-      partitionSD.put(broker, SD);
-      mean = 0f;
-      LoadSQR = 0f;
-    }
+    load.keySet()
+        .forEach(
+            broker -> {
+              load.get(broker)
+                  .keySet()
+                  .forEach(
+                      tp -> {
+                        loadSet.add(load.get(broker).get(tp));
+                        partitionLoad.put(tp, load.get(broker).get(tp));
+                        LoadSQR.add(Math.pow(load.get(broker).get(tp), 2));
+                      });
+              var loadSum = countSum(loadSet);
+              var partitionNum = load.get(broker).keySet().size();
+              brokerLoad.put(broker, loadSum);
+              var mean = loadSum / load.get(broker).size();
+              partitionMean.put(broker, loadSum / load.get(broker).size());
+              var SD =
+                  Math.pow((countSum(LoadSQR) - mean * mean * partitionNum) / partitionNum, 0.5);
+              partitionSD.put(broker, SD);
+              loadSet.clear();
+              LoadSQR.clear();
+            });
 
-    for (var i : brokerLoad.keySet()) {
-      mean += brokerLoad.get(i);
-      LoadSQR += Math.pow(brokerLoad.get(i), 2);
-    }
-    mean /= brokerLoad.keySet().size();
+    brokerLoad
+        .keySet()
+        .forEach(
+            broker -> {
+              loadSet.add(brokerLoad.get(broker));
+              LoadSQR.add(Math.pow(brokerLoad.get(broker), 2));
+            });
+    var brokerLoadMean = countSum(loadSet) / brokerLoad.keySet().size();
+    var brokerLoadSQR = countSum(LoadSQR);
     brokerSD =
         Math.pow(
-            (LoadSQR - mean * mean * brokerLoad.keySet().size()) / brokerLoad.keySet().size(), 0.5);
-    for (var broker : load.keySet()) {
-      Map<TopicPartition, Double> partitionScore = new HashMap<>();
-      for (var topicPartition : load.get(broker).keySet()) {
-        if (brokerLoad.get(broker) - mean > 0) {
-          partitionScore.put(
-              topicPartition,
-              ((brokerLoad.get(broker) - mean) / brokerSD)
-                  * ((partitionLoad.get(topicPartition) - partitionMean.get(broker))
-                      / partitionSD.get(broker))
-                  * 60.0);
-
-        } else {
-          partitionScore.put(topicPartition, 0.0);
-        }
-        brokerPartitionScore.put(broker, partitionScore);
-      }
-    }
+            (brokerLoadSQR - brokerLoadMean * brokerLoadMean * brokerLoad.keySet().size())
+                / brokerLoad.keySet().size(),
+            0.5);
+    load.keySet()
+        .forEach(
+            broker -> {
+              Map<TopicPartition, Double> partitionScore =
+                  new TreeMap<>(
+                      Comparator.comparing(TopicPartition::topic)
+                          .thenComparing(TopicPartition::partition));
+              load.get(broker)
+                  .keySet()
+                  .forEach(
+                      topicPartition -> {
+                        if (brokerLoad.get(broker) - brokerLoadMean > 0) {
+                          partitionScore.put(
+                              topicPartition,
+                              Math.round(
+                                      (((brokerLoad.get(broker) - brokerLoadMean) / brokerSD)
+                                              * ((partitionLoad.get(topicPartition)
+                                                      - partitionMean.get(broker))
+                                                  / partitionSD.get(broker))
+                                              * 60.0)
+                                          * 100.0)
+                                  / 100.0);
+                        } else {
+                          partitionScore.put(topicPartition, 0.0);
+                        }
+                        brokerPartitionScore.put(broker, partitionScore);
+                      });
+            });
     return brokerPartitionScore;
   }
 }
diff --git a/app/src/main/java/org/astraea/topic/cost/GetPartitionInf.java b/app/src/main/java/org/astraea/topic/cost/GetPartitionInf.java
index 4216ea8b..613694b1 100644
--- a/app/src/main/java/org/astraea/topic/cost/GetPartitionInf.java
+++ b/app/src/main/java/org/astraea/topic/cost/GetPartitionInf.java
@@ -8,21 +8,31 @@ import org.astraea.topic.TopicAdmin;
 public class GetPartitionInf {
   static Map<Integer, Map<TopicPartition, Integer>> getSize(TopicAdmin client) {
     Map<Integer, Map<TopicPartition, Integer>> brokerPartitionSize = new HashMap<>();
-    for (var broker : client.brokerIds()) {
-      Map<TopicPartition, Integer> partitionSize = new HashMap<>();
-      for (var partition : client.replicas(client.publicTopics().keySet()).keySet())
-        if (client.replicas(client.publicTopics().keySet()).get(partition).get(0).broker()
-            == broker)
-          partitionSize.put(
-              partition,
-              (int) client.replicas(client.publicTopics().keySet()).get(partition).get(0).size());
-      brokerPartitionSize.put(broker, partitionSize);
-    }
+    client
+        .brokerIds()
+        .forEach(
+            (broker) -> {
+              TreeMap<TopicPartition, Integer> partitionSize =
+                  new TreeMap<>(
+                      Comparator.comparing(TopicPartition::topic)
+                          .thenComparing(TopicPartition::partition));
+              client
+                  .replicas(client.topicNames())
+                  .forEach(
+                      (tp, assignment) -> {
+                        assignment.forEach(
+                            partition -> {
+                              if (partition.broker() == broker)
+                                partitionSize.put(tp, (int) partition.size());
+                            });
+                        brokerPartitionSize.put(broker, partitionSize);
+                      });
+            });
     return brokerPartitionSize;
   }
 
   static Map<String, Integer> getRetentionMillis(TopicAdmin client) {
-    return client.publicTopics().entrySet().stream()
+    return client.topics().entrySet().stream()
         .collect(
             Collectors.toMap(
                 Map.Entry::getKey,
diff --git a/app/src/main/java/org/astraea/topic/cost/PartitionScore.java b/app/src/main/java/org/astraea/topic/cost/PartitionScore.java
index fcda459c..40def4f6 100644
--- a/app/src/main/java/org/astraea/topic/cost/PartitionScore.java
+++ b/app/src/main/java/org/astraea/topic/cost/PartitionScore.java
@@ -1,61 +1,104 @@
 package org.astraea.topic.cost;
 
+import static org.astraea.argument.ArgumentUtil.parseArgument;
+
+import com.beust.jcommander.Parameter;
+import com.beust.jcommander.converters.BooleanConverter;
 import java.util.*;
 import java.util.stream.Collectors;
 import org.apache.kafka.common.TopicPartition;
-import org.astraea.argument.ArgumentUtil;
 import org.astraea.argument.BasicArgument;
+import org.astraea.argument.validator.NotEmptyString;
 import org.astraea.topic.TopicAdmin;
 
 public class PartitionScore {
-  TopicAdmin admin;
-  Map<Integer, Map<TopicPartition, Integer>> brokerPartitionSize;
-  Map<Integer, Map<TopicPartition, Double>> score;
-  Map<Integer, Map<TopicPartition, Double>> load;
-  Map<String, Integer> retentionMillis;
-
-  public PartitionScore(String address) {
-    admin = TopicAdmin.of(address);
-    brokerPartitionSize = GetPartitionInf.getSize(admin);
-    retentionMillis = GetPartitionInf.getRetentionMillis(admin);
-    load = CalculateUtils.getLoad(brokerPartitionSize, retentionMillis);
-    score = CalculateUtils.getScore(load);
-  }
+  public static void printScore(
+      Map<Integer, Map<TopicPartition, Double>> score, Argument argument) {
+    List<TopicPartition> partitionGood = new ArrayList<>();
+    Map<Integer, Boolean> brokerGood = new HashMap<>();
+    score
+        .keySet()
+        .forEach(
+            broker -> {
+              brokerGood.put(broker, true);
+              score
+                  .get(broker)
+                  .keySet()
+                  .forEach(
+                      tp -> {
+                        if (score.get(broker).get(tp) > 0) brokerGood.put(broker, false);
+                      });
+
+              if (!brokerGood.get(broker)) {
+                System.out.println("\nbroker: " + broker);
+                score
+                    .get(broker)
+                    .keySet()
+                    .forEach(
+                        tp -> {
+                          if (score.get(broker).get(tp) > 0) {
+                            System.out.println(tp + ": " + score.get(broker).get(tp));
+                          } else {
+                            partitionGood.add(tp);
+                          }
+                        });
+              }
+            });
+    if (!argument.hideBalanced) {
+      System.out.println(
+          "\nThe following brokers are balanced: "
+              + brokerGood.entrySet().stream()
+                  .filter(Map.Entry::getValue)
+                  .map(Map.Entry::getKey)
+                  .collect(Collectors.toSet()));
 
-  public void printScore(Map<Integer, Map<TopicPartition, Double>> score) {
-    Set<TopicPartition> partitionGood = new HashSet<>();
-    Map<Integer, Boolean> BrokerGood = new HashMap<>();
-    for (var broker : score.keySet()) {
-      BrokerGood.put(broker, true);
-      for (var tp : score.get(broker).keySet())
-        if (score.get(broker).get(tp) > 0) BrokerGood.put(broker, false);
-      System.out.println();
-      if (BrokerGood.get(broker)) {
-        System.out.println("broker: " + broker + " is balanced.");
-      } else {
-        System.out.println("broker: " + broker);
-        for (var tp : score.get(broker).keySet()) {
-          if (score.get(broker).get(tp) > 0) {
-            System.out.println(tp + ": " + score.get(broker).get(tp));
-          } else {
-            partitionGood.add(tp);
-          }
-        }
-      }
+      partitionGood.sort(
+          Comparator.comparing(TopicPartition::topic).thenComparing(TopicPartition::partition));
+      System.out.println(
+          "The following partitions are balanced: "
+              + partitionGood.stream()
+                  .map(String::valueOf)
+                  .collect(Collectors.joining(", ", "[", "]")));
     }
+  }
 
-    System.out.print(
-        "\nThe following partitions are balanced: "
-            + partitionGood.stream()
-                .map(String::valueOf)
-                .collect(Collectors.joining(", ", "[", "]")));
+  public static Map<Integer, Map<TopicPartition, Double>> execute(
+      Argument argument, TopicAdmin admin) {
+    var internalTopic =
+        Set.of(
+            "__consumer_offsets",
+            "_confluent-command",
+            "_confluent-metrics",
+            "_confluent-telemetry-metrics",
+            "__transaction_state");
+    var brokerPartitionSize = GetPartitionInf.getSize(admin);
+    var retentionMillis = GetPartitionInf.getRetentionMillis(admin);
+    if (argument.excludeInternalTopic) internalTopic.forEach(retentionMillis::remove);
+    var load = CalculateUtils.getLoad(brokerPartitionSize, retentionMillis);
+    return CalculateUtils.getScore(load);
   }
 
   public static void main(String[] args) {
-    var argument = ArgumentUtil.parseArgument(new Argument(), args);
-    var partitionScore = new PartitionScore(argument.brokers);
-    partitionScore.printScore(partitionScore.score);
+    var argument = parseArgument(new Argument(), args);
+    var admin = TopicAdmin.of(argument.brokers);
+    var score = execute(argument, admin);
+    printScore(score, argument);
   }
 
-  static class Argument extends BasicArgument {}
+  static class Argument extends BasicArgument {
+    @Parameter(
+        names = {"--exclude.internal.topics"},
+        description =
+            "True if you want to ignore internal topics like _consumer_offsets while counting score.",
+        validateWith = NotEmptyString.class,
+        converter = BooleanConverter.class)
+    boolean excludeInternalTopic = false;
+
+    @Parameter(
+        names = {"--hide.balanced"},
+        description = "True if you want to hide topics and partitions thar already balanced.",
+        validateWith = NotEmptyString.class,
+        converter = BooleanConverter.class)
+    boolean hideBalanced = false;
+  }
 }


Test patch:
diff --git a/app/src/test/java/org/astraea/topic/cost/CalculateUtilsTest.java b/app/src/test/java/org/astraea/topic/cost/CalculateUtilsTest.java
index 80dd0748..0a3f6478 100644
--- a/app/src/test/java/org/astraea/topic/cost/CalculateUtilsTest.java
+++ b/app/src/test/java/org/astraea/topic/cost/CalculateUtilsTest.java
@@ -54,14 +54,14 @@ public class CalculateUtilsTest extends RequireBrokerCluster {
     assertEquals(2, Load.size());
     assertEquals(4, Load.get(0).size());
     assertEquals(4, Load.get(1).size());
-    assertEquals(0.165, round(Load.get(0).get(new TopicPartition("test0", 0))));
-    assertEquals(0.331, round(Load.get(0).get(new TopicPartition("test0", 1))));
-    assertEquals(0.496, round(Load.get(0).get(new TopicPartition("test0", 2))));
-    assertEquals(0.661, round(Load.get(0).get(new TopicPartition("test0", 3))));
-    assertEquals(0.827, round(Load.get(1).get(new TopicPartition("test1", 0))));
-    assertEquals(0.992, round(Load.get(1).get(new TopicPartition("test1", 1))));
-    assertEquals(1.157, round(Load.get(1).get(new TopicPartition("test1", 2))));
-    assertEquals(1.323, round(Load.get(1).get(new TopicPartition("test1", 3))));
+    assertEquals(0.17, round(Load.get(0).get(new TopicPartition("test0", 0))));
+    assertEquals(0.33, round(Load.get(0).get(new TopicPartition("test0", 1))));
+    assertEquals(0.50, round(Load.get(0).get(new TopicPartition("test0", 2))));
+    assertEquals(0.66, round(Load.get(0).get(new TopicPartition("test0", 3))));
+    assertEquals(0.83, round(Load.get(1).get(new TopicPartition("test1", 0))));
+    assertEquals(0.99, round(Load.get(1).get(new TopicPartition("test1", 1))));
+    assertEquals(1.16, round(Load.get(1).get(new TopicPartition("test1", 2))));
+    assertEquals(1.32, round(Load.get(1).get(new TopicPartition("test1", 3))));
   }
 
   @Test
@@ -74,13 +74,13 @@ public class CalculateUtilsTest extends RequireBrokerCluster {
     assertEquals(0.0, round(Score.get(0).get(new TopicPartition("test0", 1))));
     assertEquals(0.0, round(Score.get(0).get(new TopicPartition("test0", 2))));
     assertEquals(0.0, round(Score.get(0).get(new TopicPartition("test0", 3))));
-    assertEquals(-80.498, round(Score.get(1).get(new TopicPartition("test1", 0))));
-    assertEquals(-26.833, round(Score.get(1).get(new TopicPartition("test1", 1))));
-    assertEquals(26.833, round(Score.get(1).get(new TopicPartition("test1", 2))));
-    assertEquals(80.498, round(Score.get(1).get(new TopicPartition("test1", 3))));
+    assertEquals(-80.50, round(Score.get(1).get(new TopicPartition("test1", 0))));
+    assertEquals(-26.83, round(Score.get(1).get(new TopicPartition("test1", 1))));
+    assertEquals(26.83, round(Score.get(1).get(new TopicPartition("test1", 2))));
+    assertEquals(80.50, round(Score.get(1).get(new TopicPartition("test1", 3))));
   }
 
   double round(double score) {
-    return Math.round(1000 * score) / 1000.0;
+    return Math.round(100 * score) / 100.0;
   }
 }
diff --git a/app/src/test/java/org/astraea/topic/cost/GetPartitionInfTest.java b/app/src/test/java/org/astraea/topic/cost/GetPartitionInfTest.java
index f06c348c..51953779 100644
--- a/app/src/test/java/org/astraea/topic/cost/GetPartitionInfTest.java
+++ b/app/src/test/java/org/astraea/topic/cost/GetPartitionInfTest.java
@@ -67,8 +67,7 @@ public class GetPartitionInfTest extends RequireBrokerCluster {
   }
 
   @Test
-  void getSize() {
-
+  void testGetSize() {
     var brokerPartitionSize = GetPartitionInf.getSize(admin);
     assertEquals(3, brokerPartitionSize.size());
     assertEquals(
@@ -79,7 +78,7 @@ public class GetPartitionInfTest extends RequireBrokerCluster {
   }
 
   @Test
-  void getRetentionMillis() {
+  void testGetRetentionMillis() {
     var brokerPartitionRetentionMillis = GetPartitionInf.getRetentionMillis(admin);
     assertEquals(3, brokerPartitionRetentionMillis.size());
   }
diff --git a/app/src/test/java/org/astraea/topic/cost/PartitionScoreTest.java b/app/src/test/java/org/astraea/topic/cost/PartitionScoreTest.java
index 414018c7..bca1a744 100644
--- a/app/src/test/java/org/astraea/topic/cost/PartitionScoreTest.java
+++ b/app/src/test/java/org/astraea/topic/cost/PartitionScoreTest.java
@@ -22,7 +22,7 @@ public class PartitionScoreTest extends RequireBrokerCluster {
     Map<Integer, String> topicName = new HashMap<>();
     topicName.put(0, "testPartitionScore0");
     topicName.put(1, "testPartitionScore1");
-    topicName.put(2, "testPartitionScore2");
+    topicName.put(2, "__consumer_offsets");
     try (var admin = TopicAdmin.of(bootstrapServers())) {
       admin
           .creator()
@@ -50,7 +50,7 @@ public class PartitionScoreTest extends RequireBrokerCluster {
     var producer =
         Producer.builder().brokers(bootstrapServers()).keySerializer(Serializer.STRING).build();
     int size = 10000;
-    for (int t = 0; t <= 2; t++) {
+    for (int t = 0; t <= 1; t++) {
       for (int p = 0; p <= 3; p++) {
         producer
             .sender()
@@ -67,13 +67,15 @@ public class PartitionScoreTest extends RequireBrokerCluster {
   }
 
   @Test
-  void testGetScore() throws ExecutionException, InterruptedException {
-    PartitionScore partitionScore = new PartitionScore(bootstrapServers());
-    assertEquals(3, partitionScore.score.size());
-    assertEquals(
-        3 * 4,
-        partitionScore.score.get(0).size()
-            + partitionScore.score.get(1).size()
-            + partitionScore.score.get(2).size());
+  void testGetScore() {
+    PartitionScore.Argument argument = new PartitionScore.Argument();
+    argument.excludeInternalTopic = false;
+    var score = PartitionScore.execute(argument, admin);
+    assertEquals(3, score.size());
+    assertEquals(3 * 4, score.get(0).size() + score.get(1).size() + score.get(2).size());
+    argument.excludeInternalTopic = true;
+    score = PartitionScore.execute(argument, admin);
+    assertEquals(3, score.size());
+    assertEquals(2 * 4, score.get(0).size() + score.get(1).size() + score.get(2).size());
   }
 }