diff --git a/services/tests/servicestests/src/com/android/server/wm/WindowTestsBase.java b/services/tests/servicestests/src/com/android/server/wm/WindowTestsBase.java
index 32eee84..6618a69 100644
--- a/services/tests/servicestests/src/com/android/server/wm/WindowTestsBase.java
+++ b/services/tests/servicestests/src/com/android/server/wm/WindowTestsBase.java
@@ -115,6 +115,10 @@
                 "mChildAppWindowAbove");
         mChildAppWindowBelow = createCommonWindow(mAppWindow, TYPE_APPLICATION_MEDIA_OVERLAY,
                 "mChildAppWindowBelow");
+
+        // Adding a display will cause freezing the display. Make sure to wait until it's unfrozen
+        // to not run into race conditions with the tests.
+        waitUntilHandlersIdle();
     }
 
     @After
@@ -135,6 +139,9 @@
             mDisplayContent.removeImmediately();
             sWm.mInputMethodTarget = null;
         }
+
+        // Wait until everything is really cleaned up.
+        waitUntilHandlersIdle();
     }
 
     private WindowState createCommonWindow(WindowState parent, int type, String name) {