Tidyup without breaking tests

Signed-off-by: Tom Bentley <[email protected]>
kroxylicious · Sep 13, 2023 · c2b48cb · c2b48cb
1 parent a0a648f
commit c2b48cb
Showing 1 changed file with 13 additions and 3 deletions.
diff --git a/impl/src/main/java/io/kroxylicious/testing/kafka/invm/InVMKafkaCluster.java b/impl/src/main/java/io/kroxylicious/testing/kafka/invm/InVMKafkaCluster.java
@@ -270,7 +270,7 @@ public synchronized String getBootstrapServers() {
 
     private synchronized String buildBrokerList(Function<Integer, KafkaClusterConfig.KafkaEndpoints.EndpointPair> endpointFunc) {
         return servers.keySet().stream()
-                .filter(nodeId -> clusterConfig.hasBrokerRole(nodeId))
+                .filter(this::isBroker)
                 .map(endpointFunc)
                 .map(KafkaClusterConfig.KafkaEndpoints.EndpointPair::connectAddress)
                 .collect(Collectors.joining(","));
@@ -406,8 +406,18 @@ public synchronized void close() throws Exception {
      */
     private void roleOrderedShutdown(Map<Integer, Server> servers) {
         // Shutdown servers without a controller port first.
-        shutdownServers(servers, e -> clusterConfig.isPureBroker(e.getKey()));
-        shutdownServers(servers, e -> clusterConfig.hasControllerRole(e.getKey()));
+        shutdownServers(servers, e -> !isController(e.getKey()));
+        shutdownServers(servers, e -> isController(e.getKey()));
+    }
+
+    private boolean isController(Integer key) {
+        return key < clusterConfig.numNodes() // dynamically added nodes are always pure brokers
+                && clusterConfig.hasControllerRole(key);
+    }
+
+    private boolean isBroker(Integer key) {
+        return key >= clusterConfig.numNodes() // dynamically added nodes are always pure brokers
+                || clusterConfig.hasBrokerRole(key);
     }
 
     @SuppressWarnings("java:S3864") // Stream.peek is being used with caution.