diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/scheduler/TestSchedulerOvercommit.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/scheduler/TestSchedulerOvercommit.java
index cc665fb9020..758fd33e8a1 100644
--- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/scheduler/TestSchedulerOvercommit.java
+++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/scheduler/TestSchedulerOvercommit.java
@@ -286,6 +286,9 @@ public void testReducePreemptAndKill() throws Exception {
     updateNodeResource(rm, nmId, 2 * GB, 2, timeout);
     waitMemory(scheduler, nm, 4 * GB, -2 * GB, INTERVAL, timeout);
 
+    // wait until MARK_CONTAINER_FOR_PREEMPTION is handled
+    rm.drainEvents();
+
     // We should receive a notification to preempt the container
     PreemptionMessage preemptMsg = am.schedule().getPreemptionMessage();
     assertPreemption(container.getId(), preemptMsg);
@@ -315,13 +318,16 @@ public void testReducePreemptAndCancel() throws Exception {
     Container container = createContainer(am, 2 * GB);
     assertMemory(scheduler, nmId, 4 * GB, 0);
 
-    // We give an overcommit time out of 2 seconds
+    // We give an overcommit time out of 1 seconds
     final int timeout = (int)TimeUnit.SECONDS.toMillis(1);
 
     // Reducing to 2GB should first preempt the container
     updateNodeResource(rm, nmId, 2 * GB, 2, timeout);
     waitMemory(scheduler, nm, 4 * GB, -2 * GB, INTERVAL, timeout);
 
+    // wait until MARK_CONTAINER_FOR_PREEMPTION is handled
+    rm.drainEvents();
+
     // We should receive a notification to preempt the container
     PreemptionMessage preemptMsg = am.schedule().getPreemptionMessage();
     assertPreemption(container.getId(), preemptMsg);
@@ -479,6 +485,9 @@ public void testEndToEnd() throws Exception {
     updateNodeResource(rm, nmId, 3 * GB, 2, 2 * 1000);
     waitMemory(scheduler, nmId, 5 * GB, -2 * GB, 200, 5 * 1000);
 
+    // wait until MARK_CONTAINER_FOR_PREEMPTION is handled
+    rm.drainEvents();
+
     PreemptionMessage preemptMsg = am.schedule().getPreemptionMessage();
     assertPreemption(c2.getId(), preemptMsg);
 
@@ -493,6 +502,9 @@ public void testEndToEnd() throws Exception {
     updateNodeResource(rm, nmId, 3 * GB, 2, 2 * 1000);
     waitMemory(scheduler, nmId, 5 * GB, -2 * GB, 200, 5 * 1000);
 
+    // wait until MARK_CONTAINER_FOR_PREEMPTION is handled
+    rm.drainEvents();
+
     preemptMsg = am.schedule().getPreemptionMessage();
     assertPreemption(c2.getId(), preemptMsg);