KAFKA-18844: Fix flaky QuorumControllerTest.testBalancePartitionLeade…

…rs()

Signed-off-by: PoAn Yang <[email protected]>

metadata/src/test/java/org/apache/kafka/controller/QuorumControllerTest.java

@@ -373,7 +373,7 @@ public  void testElrEnabledByDefault() throws Throwable {
                 )).
                 build()
         ) {
-            controlEnv.activeController(true);
+            controlEnv.activeController();
             assertTrue(controlEnv.activeController().configurationControl().clusterConfig().containsKey(TopicConfig.MIN_IN_SYNC_REPLICAS_CONFIG));
         }
     }

metadata/src/test/java/org/apache/kafka/controller/QuorumControllerTestEnv.java

@@ -135,10 +135,6 @@ private QuorumControllerTestEnv(
     }
 
     QuorumController activeController() throws InterruptedException {
-        return activeController(false);
-    }
-
-    QuorumController activeController(boolean waitForActivation) throws InterruptedException {
         AtomicReference<QuorumController> value = new AtomicReference<>(null);
         TestUtils.retryOnExceptionWithTimeout(20000, 3, () -> {
             LeaderAndEpoch leader = logEnv.leaderAndEpoch();
@@ -155,16 +151,14 @@ QuorumController activeController(boolean waitForActivation) throws InterruptedE
             }
         });
 
-        if (waitForActivation) {
-            try {
-                // ControllerActivation happens after curClaimEpoch is set, so we need to put something on
-                // the end of the queue and wait for it to complete before returning the active controller.
-                value.get()
-                    .appendReadEvent("wait for activation", OptionalLong.empty(), () -> null)
-                    .get(20000, TimeUnit.MILLISECONDS);
-            } catch (Throwable t) {
-                throw new RuntimeException("Failed while waiting for controller activation", t);
-            }
+        try {
+            // ControllerActivation happens after curClaimEpoch is set, so we need to put something on
+            // the end of the queue and wait for it to complete before returning the active controller.
+            value.get()
+                .appendReadEvent("wait for activation", OptionalLong.empty(), () -> null)
+                .get(20000, TimeUnit.MILLISECONDS);
+        } catch (Throwable t) {
+            throw new RuntimeException("Failed while waiting for controller activation", t);
         }
 
         return value.get();