diff --git a/hadoop-ozone/integration-test/src/test/java/org/apache/hadoop/ozone/TestContainerBalancerOperations.java b/hadoop-ozone/integration-test/src/test/java/org/apache/hadoop/ozone/TestContainerBalancerOperations.java
index a2a534f839d7..b759e7e8cbeb 100644
--- a/hadoop-ozone/integration-test/src/test/java/org/apache/hadoop/ozone/TestContainerBalancerOperations.java
+++ b/hadoop-ozone/integration-test/src/test/java/org/apache/hadoop/ozone/TestContainerBalancerOperations.java
@@ -33,6 +33,7 @@
 import org.apache.hadoop.hdds.scm.client.ScmClient;
 import org.apache.hadoop.hdds.scm.container.balancer.ContainerBalancerConfiguration;
 import org.apache.hadoop.hdds.scm.container.placement.algorithms.SCMContainerPlacementCapacity;
+import org.apache.ozone.test.GenericTestUtils;
 import org.junit.jupiter.api.AfterAll;
 import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.Test;
@@ -97,15 +98,17 @@ public void testContainerBalancerCLIOperations() throws Exception {
     running = containerBalancerClient.getContainerBalancerStatus();
     assertTrue(running);
 
-    // waiting for balance completed.
-    // TODO: this is a temporary implementation for now
-    // modify this after balancer is fully completed
-    try {
-      Thread.sleep(20000);
-    } catch (InterruptedException e) { }
-
-    running = containerBalancerClient.getContainerBalancerStatus();
-    assertFalse(running);
+    GenericTestUtils.waitFor(
+        () -> {
+          try {
+            return !containerBalancerClient.getContainerBalancerStatus();
+          } catch (IOException e) {
+            return false;
+          }
+        },
+        100,
+        30000
+    );
 
     // test normally start , and stop it before balance is completed
     containerBalancerClient.startContainerBalancer(threshold, iterations,