diff --git a/datastream-server-restli/src/test/java/com/linkedin/datastream/server/TestCoordinator.java b/datastream-server-restli/src/test/java/com/linkedin/datastream/server/TestCoordinator.java
index e8b2445d8..84283fb6d 100644
--- a/datastream-server-restli/src/test/java/com/linkedin/datastream/server/TestCoordinator.java
+++ b/datastream-server-restli/src/test/java/com/linkedin/datastream/server/TestCoordinator.java
@@ -4062,6 +4062,8 @@ protected synchronized void handleEvent(CoordinatorEvent event) {
     Datastream testDatastream =
         DatastreamTestUtils.createAndStoreDatastreams(zkClient, testCluster, connectorType, streamName)[0];
 
+    // Blocking until we have performed atleast 5 retries of LeaderDoAssignmentEvent with a newly elected leader.
+    PollUtils.poll(() -> shadowListWithPreviousAndNewHeadPairsAtNewLeaderDoAssignmentEvent.size() > 5, 50, 2000);
     coordinator.stop();
     zkClient.close();
     coordinator.getDatastreamCache().getZkclient().close();