diff --git a/pinot-controller/src/test/java/org/apache/pinot/controller/helix/ControllerTest.java b/pinot-controller/src/test/java/org/apache/pinot/controller/helix/ControllerTest.java
index 2cd41d987d29..13d0a4695a0a 100644
--- a/pinot-controller/src/test/java/org/apache/pinot/controller/helix/ControllerTest.java
+++ b/pinot-controller/src/test/java/org/apache/pinot/controller/helix/ControllerTest.java
@@ -608,7 +608,8 @@ public void stopFakeInstance(String instanceId) {
 
   public void stopAndDropFakeInstance(String instanceId) {
     stopFakeInstance(instanceId);
-    _helixResourceManager.dropInstance(instanceId);
+    TestUtils.waitForCondition(aVoid -> _helixResourceManager.dropInstance(instanceId).isSuccessful(), 60_000L,
+        "Failed to drop fake instance: " + instanceId);
   }
 
   public static Schema createDummySchema(String tableName) {
diff --git a/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/PinotHelixResourceManagerStatelessTest.java b/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/PinotHelixResourceManagerStatelessTest.java
index ba4f061c39ab..ecfd025d17cb 100644
--- a/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/PinotHelixResourceManagerStatelessTest.java
+++ b/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/PinotHelixResourceManagerStatelessTest.java
@@ -187,7 +187,7 @@ public void testGetDataInstanceAdminEndpoints()
     }, 60_000L, "Failed to update the admin port");
 
     // Remove the new added server
-    _helixResourceManager.dropInstance(serverName);
+    assertTrue(_helixResourceManager.dropInstance(serverName).isSuccessful());
     TestUtils.waitForCondition(aVoid -> {
       try {
         _helixResourceManager.getDataInstanceAdminEndpoints(Collections.singleton(serverName));
@@ -220,7 +220,7 @@ public void testAddRemoveInstance() {
     assertTrue(allInstances.contains(instanceName));
 
     // Remove the added instance
-    _helixResourceManager.dropInstance(instanceName);
+    assertTrue(_helixResourceManager.dropInstance(instanceName).isSuccessful());
     allInstances = _helixResourceManager.getAllInstances();
     assertFalse(allInstances.contains(instanceName));
   }
diff --git a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/PartialUpsertTableRebalanceIntegrationTest.java b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/PartialUpsertTableRebalanceIntegrationTest.java
index caf3dfbce8eb..8afcbe917713 100644
--- a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/PartialUpsertTableRebalanceIntegrationTest.java
+++ b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/PartialUpsertTableRebalanceIntegrationTest.java
@@ -155,14 +155,11 @@ public void testRebalance()
     waitForRebalanceToComplete(rebalanceResult, 600_000L);
     waitForAllDocsLoaded(600_000L);
 
-    _resourceManager.disableInstance(serverStarter1.getInstanceId());
-    _resourceManager.disableInstance(serverStarter2.getInstanceId());
-
-    _resourceManager.dropInstance(serverStarter1.getInstanceId());
-    _resourceManager.dropInstance(serverStarter2.getInstanceId());
-
     serverStarter1.stop();
     serverStarter2.stop();
+    TestUtils.waitForCondition(aVoid -> _resourceManager.dropInstance(serverStarter1.getInstanceId()).isSuccessful()
+            && _resourceManager.dropInstance(serverStarter2.getInstanceId()).isSuccessful(), 60_000L,
+        "Failed to drop servers");
   }
 
   @Test