Refactoring + fixing more bugs

Signed-off-by: Kyle Liberti <kliberti@redhat.com>

Author: kyguy

cluster-operator/src/main/java/io/strimzi/operator/cluster/operator/assembly/KafkaRebalanceAssemblyOperator.java

@@ -343,15 +343,19 @@ private static Future<MapAndStatus<ConfigMap, KafkaRebalanceStatus>> updateProgr
 
                     if (existingConfigMap == null && desiredConfigMap == null) {
                         return Future.succeededFuture(desiredStatusAndMap);
-                    } else if (existingConfigMap != null && desiredStatusAndMap.getLoadMap() != null) {
-                        // Ensure desiredConfigMap retains broker load information if it exists.
-                        desiredConfigMap.getData().put(BROKER_LOAD_KEY, existingConfigMap.getData().get(BROKER_LOAD_KEY));
-                    } else if (desiredConfigMap == null) {
-                        desiredStatusAndMap.setLoadMap(existingConfigMap);
-                        desiredConfigMap = existingConfigMap;
                     }
 
-                    // Add progress information to KafkaRebalance resource status and ConfigMap.
+                    if (existingConfigMap != null) {
+                        if (desiredConfigMap == null) {
+                            desiredStatusAndMap.setLoadMap(existingConfigMap);
+                            desiredConfigMap = existingConfigMap;
+                        } else {
+                            // Ensure desiredConfigMap retains broker load information if it exists.
+                            desiredConfigMap.getData().put(BROKER_LOAD_KEY, existingConfigMap.getData().get(BROKER_LOAD_KEY));
+                        }
+                    }
+
+                    // Add progress information to `KafkaRebalance` resource status and ConfigMap.
                     desiredStatus.setProgress(Map.of(REBALANCE_PROGRESS_CONFIG_MAP_KEY, configMapName));
                     KafkaRebalanceConfigMapUtils.updateRebalanceConfigMap(reconciliation, desiredState, host, cruiseControlPort, apiClient, desiredConfigMap)
                             .onFailure(exception -> {

cluster-operator/src/test/java/io/strimzi/operator/cluster/operator/assembly/KafkaRebalanceAssemblyOperatorProgressTest.java

@@ -254,7 +254,7 @@ public void testNewToPendingProposalToProposalReadyRemoveDisks(VertxTestContext
     private void krNewToPendingProposalToProposalReady(VertxTestContext context, CruiseControlEndpoints endpoint, KafkaRebalance kr, String verbose) throws IOException, URISyntaxException {
         // Set up the rebalance endpoint with the number of pending calls before a response is received.
         cruiseControlServer.setupCCRebalanceResponse(1, endpoint, verbose);
-        cruiseControlServer.setupCCUserTasksResponseNoGoals(1, 0);
+        cruiseControlServer.setupCCUserTasksResponseNoGoals(1, 0, false);
 
         Crds.kafkaRebalanceOperation(client).inNamespace(namespace).resource(kr).create();
         crdCreateKafka();
@@ -416,7 +416,7 @@ public void testNewToPendingProposalToStoppedAndRefreshRemoveDisks(VertxTestCont
     private void krNewToPendingProposalToStoppedAndRefresh(VertxTestContext context, CruiseControlEndpoints endpoint, KafkaRebalance kr, String verbose) throws IOException, URISyntaxException {
         // Set up the rebalance endpoint with the number of pending calls before a response is received.
         cruiseControlServer.setupCCRebalanceResponse(2, endpoint, verbose);
-        cruiseControlServer.setupCCUserTasksResponseNoGoals(0, 0);
+        cruiseControlServer.setupCCUserTasksResponseNoGoals(0, 0, false);
 
         Crds.kafkaRebalanceOperation(client).inNamespace(namespace).resource(kr).create();
         crdCreateKafka();
@@ -560,7 +560,7 @@ public void testAutoApprovalNewToReadyRemoveDisks(VertxTestContext context) thro
     private void krNewToProposalReadyToRebalancingToReady(VertxTestContext context, CruiseControlEndpoints endpoint, KafkaRebalance kr, String verbose) throws IOException, URISyntaxException {
         // Set up the rebalance and user tasks endpoints with the number of pending calls before a response is received.
         cruiseControlServer.setupCCRebalanceResponse(0, endpoint, verbose);
-        cruiseControlServer.setupCCUserTasksResponseNoGoals(0, 0);
+        cruiseControlServer.setupCCUserTasksResponseNoGoals(0, 0, false);
 
         Crds.kafkaRebalanceOperation(client).inNamespace(namespace).resource(kr).create();
         crdCreateKafka();
@@ -643,7 +643,7 @@ public void testNewToProposalReadyToReconciliationPausedRemoveDisks(VertxTestCon
     private void krNewToProposalReadyToReconciliationPaused(VertxTestContext context, CruiseControlEndpoints endpoint, KafkaRebalance kr, String verbose) throws IOException, URISyntaxException {
         // Set up the rebalance and user tasks endpoints with the number of pending calls before a response is received.
         cruiseControlServer.setupCCRebalanceResponse(0, endpoint, verbose);
-        cruiseControlServer.setupCCUserTasksResponseNoGoals(0, 0);
+        cruiseControlServer.setupCCUserTasksResponseNoGoals(0, 0, false);
 
         Crds.kafkaRebalanceOperation(client).inNamespace(namespace).resource(kr).create();
         crdCreateKafka();
@@ -731,7 +731,7 @@ public void testNewToProposalReadyToRebalancingToReadyThenRefreshRemoveDisks(Ver
     private void krNewToProposalReadyToRebalancingToReadyThenRefresh(VertxTestContext context, CruiseControlEndpoints endpoint, KafkaRebalance kr, String verbose) throws IOException, URISyntaxException {
         // Set up the rebalance and user tasks endpoints with the number of pending calls before a response is received.
         cruiseControlServer.setupCCRebalanceResponse(0, endpoint, verbose);
-        cruiseControlServer.setupCCUserTasksResponseNoGoals(0, 0);
+        cruiseControlServer.setupCCUserTasksResponseNoGoals(0, 0, false);
 
         Crds.kafkaRebalanceOperation(client).inNamespace(namespace).resource(kr).create();
         crdCreateKafka();
@@ -802,7 +802,7 @@ public void krNewToProposalReadyToRebalancingToRefresh(VertxTestContext context)
 
         // Set up the rebalance and user tasks endpoints with the number of pending calls before a response is received.
         cruiseControlServer.setupCCRebalanceResponse(0, CruiseControlEndpoints.REMOVE_BROKER, "true");
-        cruiseControlServer.setupCCUserTasksResponseNoGoals(0, 0);
+        cruiseControlServer.setupCCUserTasksResponseNoGoals(0, 0, false);
         cruiseControlServer.setupCCStopResponse();
 
         Crds.kafkaRebalanceOperation(client).inNamespace(namespace).resource(kr).create();
@@ -1133,7 +1133,7 @@ public void testNewToPendingProposalDeleteRemoveDisks(VertxTestContext context)
     private void krNewToPendingProposalDelete(VertxTestContext context, CruiseControlEndpoints endpoint, KafkaRebalance kr, String verbose) throws IOException, URISyntaxException {
         // Set up the rebalance endpoint with the number of pending calls before a response is received.
         cruiseControlServer.setupCCRebalanceResponse(1, endpoint, verbose);
-        cruiseControlServer.setupCCUserTasksResponseNoGoals(1, 0);
+        cruiseControlServer.setupCCUserTasksResponseNoGoals(1, 0, false);
 
         Crds.kafkaRebalanceOperation(client).inNamespace(namespace).resource(kr).create();
         crdCreateKafka();
@@ -1328,7 +1328,7 @@ public void testNewToProposalReadyToRebalancingToStoppedRemoveDisks(VertxTestCon
     private void krNewToProposalReadyToRebalancingToStopped(VertxTestContext context, CruiseControlEndpoints endpoint, KafkaRebalance kr, String verbose) throws IOException, URISyntaxException {
         // Set up the rebalance and user tasks endpoints with the number of pending calls before a response is received.
         cruiseControlServer.setupCCRebalanceResponse(0, endpoint, verbose);
-        cruiseControlServer.setupCCUserTasksResponseNoGoals(0, 2);
+        cruiseControlServer.setupCCUserTasksResponseNoGoals(0, 2, false);
         cruiseControlServer.setupCCStopResponse();
 
         Crds.kafkaRebalanceOperation(client).inNamespace(namespace).resource(kr).create();

cluster-operator/src/test/java/io/strimzi/operator/cluster/operator/assembly/KafkaRebalanceAssemblyOperatorTest.java

@@ -1135,7 +1135,7 @@ public void testRebalancingCompletedRemoveDisks(Vertx vertx, VertxTestContext co
     }
 
     private void krRebalancingCompleted(Vertx vertx, VertxTestContext context, KafkaRebalance kcRebalance) throws IOException, URISyntaxException {
-        cruiseControlServer.setupCCUserTasksResponseNoGoals(0, 0);
+        cruiseControlServer.setupCCUserTasksResponseNoGoals(0, 0, false);
 
         checkTransition(vertx, context,
                 KafkaRebalanceState.Rebalancing, KafkaRebalanceState.Ready,
@@ -1188,7 +1188,7 @@ private void krRebalancingPendingThenExecution(Vertx vertx, VertxTestContext con
         // This tests that the optimization proposal is added correctly if it was not ready when the rebalance(dryrun=false) was called.
         // The first poll should see active and then the second should see in execution and add the optimization and cancel the timer
         // so that the status is updated.
-        cruiseControlServer.setupCCUserTasksResponseNoGoals(activeCalls, inExecutionCalls);
+        cruiseControlServer.setupCCUserTasksResponseNoGoals(activeCalls, inExecutionCalls, false);
 
         checkTransition(vertx, context,
                 KafkaRebalanceState.Rebalancing, KafkaRebalanceState.Rebalancing,
@@ -1240,7 +1240,7 @@ public void testRebalancingToStoppedRemoveDisks(Vertx vertx, VertxTestContext co
     }
 
     private void krRebalancingToStopped(Vertx vertx, VertxTestContext context, KafkaRebalance kcRebalance) throws IOException, URISyntaxException {
-        cruiseControlServer.setupCCUserTasksResponseNoGoals(0, 0);
+        cruiseControlServer.setupCCUserTasksResponseNoGoals(0, 0, false);
         cruiseControlServer.setupCCStopResponse();
 
         checkTransition(vertx, context,

cluster-operator/src/test/java/io/strimzi/operator/cluster/operator/assembly/KafkaRebalanceStateMachineTest.java

@@ -122,7 +122,7 @@ public void testCCRebalanceProposalNotReady() throws IOException, URISyntaxExcep
     @Test
     public void testCCGetRebalanceUserTask() throws IOException, URISyntaxException {
 
-        cruiseControlServer.setupCCUserTasksResponseNoGoals(0, 0);
+        cruiseControlServer.setupCCUserTasksResponseNoGoals(0, 0, false);
 
         CruiseControlApi client = cruiseControlClientProvider();
         String userTaskID = MockCruiseControl.REBALANCE_NO_GOALS_RESPONSE_UTID;
@@ -495,7 +495,7 @@ public void testMockCCServerPendingCallsOverride() throws IOException, URISyntax
         int pendingCalls1 = 2;
         int pendingCalls2 = 4;
 
-        cruiseControlServer.setupCCUserTasksResponseNoGoals(0, pendingCalls1);
+        cruiseControlServer.setupCCUserTasksResponseNoGoals(0, pendingCalls1, false);
 
         CompletableFuture<CruiseControlUserTasksResponse> statusFuture = client.getUserTaskStatus(Reconciliation.DUMMY_RECONCILIATION, HOST, cruiseControlPort, userTaskID);
 
@@ -518,7 +518,7 @@ public void testMockCCServerPendingCallsOverride() throws IOException, URISyntax
         statusFuture = statusFuture.thenCompose(response -> {
             try {
                 cruiseControlServer.reset();
-                cruiseControlServer.setupCCUserTasksResponseNoGoals(0, pendingCalls2);
+                cruiseControlServer.setupCCUserTasksResponseNoGoals(0, pendingCalls2, false);
             } catch (IOException | URISyntaxException e) {
                 return CompletableFuture.failedFuture(e);
             }
@@ -545,7 +545,7 @@ public void testMockCCServerPendingCallsOverride() throws IOException, URISyntax
     }
 
     private void runTest(String userTaskID, int pendingCalls) throws IOException, URISyntaxException {
-        cruiseControlServer.setupCCUserTasksResponseNoGoals(0, pendingCalls);
+        cruiseControlServer.setupCCUserTasksResponseNoGoals(0, pendingCalls, false);
 
         CruiseControlApi client = cruiseControlClientProvider();