diff --git a/services/tests/servicestests/src/com/android/server/wm/DisplayContentTests.java b/services/tests/servicestests/src/com/android/server/wm/DisplayContentTests.java
index 0801a88..162a1a9 100644
--- a/services/tests/servicestests/src/com/android/server/wm/DisplayContentTests.java
+++ b/services/tests/servicestests/src/com/android/server/wm/DisplayContentTests.java
@@ -25,14 +25,7 @@
 
 import java.util.ArrayList;
 
-import static android.view.WindowManager.LayoutParams.TYPE_APPLICATION_ATTACHED_DIALOG;
-import static android.view.WindowManager.LayoutParams.TYPE_APPLICATION_MEDIA;
 import static android.view.WindowManager.LayoutParams.TYPE_BASE_APPLICATION;
-import static android.view.WindowManager.LayoutParams.TYPE_INPUT_METHOD;
-import static android.view.WindowManager.LayoutParams.TYPE_INPUT_METHOD_DIALOG;
-import static android.view.WindowManager.LayoutParams.TYPE_NAVIGATION_BAR;
-import static android.view.WindowManager.LayoutParams.TYPE_STATUS_BAR;
-import static android.view.WindowManager.LayoutParams.TYPE_WALLPAPER;
 import static org.junit.Assert.assertEquals;
 
 /**
@@ -48,21 +41,8 @@
 
     @Test
     public void testForAllWindows() throws Exception {
-        final DisplayContent dc = new DisplayContent(mDisplay, sWm, null, null);
-        final WindowState wallpaperWindow = createWindow(null, TYPE_WALLPAPER, dc, "wallpaper");
-        final WindowState imeWindow = createWindow(null, TYPE_INPUT_METHOD, dc, "ime");
-        final WindowState imeDialogWindow = createWindow(null, TYPE_INPUT_METHOD_DIALOG, dc,
-                "ime dialog");
-        final WindowState statusBarWindow = createWindow(null, TYPE_STATUS_BAR, dc, "status bar");
-        final WindowState navBarWindow = createWindow(null, TYPE_NAVIGATION_BAR,
-                statusBarWindow.mToken, "nav bar");
-        final WindowState appWindow = createWindow(null, TYPE_BASE_APPLICATION, dc, "app");
-        final WindowState negChildAppWindow = createWindow(appWindow, TYPE_APPLICATION_MEDIA,
-                appWindow.mToken, "negative app child");
-        final WindowState posChildAppWindow = createWindow(appWindow,
-                TYPE_APPLICATION_ATTACHED_DIALOG, appWindow.mToken, "positive app child");
-        final WindowState exitingAppWindow = createWindow(null, TYPE_BASE_APPLICATION, dc,
-                "exiting app");
+        final WindowState exitingAppWindow = createWindow(null, TYPE_BASE_APPLICATION,
+                sDisplayContent, "exiting app");
         final AppWindowToken exitingAppToken = exitingAppWindow.mAppToken;
         exitingAppToken.mIsExiting = true;
         exitingAppToken.mTask.mStack.mExitingAppTokens.add(exitingAppToken);
@@ -70,30 +50,75 @@
         final ArrayList<WindowState> windows = new ArrayList();
 
         // Test forward traversal.
-        dc.forAllWindows(w -> {windows.add(w);}, false /* traverseTopToBottom */);
+        sDisplayContent.forAllWindows(w -> {windows.add(w);}, false /* traverseTopToBottom */);
 
-        assertEquals(wallpaperWindow, windows.get(0));
+        assertEquals(sWallpaperWindow, windows.get(0));
         assertEquals(exitingAppWindow, windows.get(1));
-        assertEquals(negChildAppWindow, windows.get(2));
-        assertEquals(appWindow, windows.get(3));
-        assertEquals(posChildAppWindow, windows.get(4));
-        assertEquals(statusBarWindow, windows.get(5));
-        assertEquals(navBarWindow, windows.get(6));
-        assertEquals(imeWindow, windows.get(7));
-        assertEquals(imeDialogWindow, windows.get(8));
+        assertEquals(sChildAppWindowBelow, windows.get(2));
+        assertEquals(sAppWindow, windows.get(3));
+        assertEquals(sChildAppWindowAbove, windows.get(4));
+        assertEquals(sDockedDividerWindow, windows.get(5));
+        assertEquals(sStatusBarWindow, windows.get(6));
+        assertEquals(sNavBarWindow, windows.get(7));
+        assertEquals(sImeWindow, windows.get(8));
+        assertEquals(sImeDialogWindow, windows.get(9));
 
         // Test backward traversal.
         windows.clear();
-        dc.forAllWindows(w -> {windows.add(w);}, true /* traverseTopToBottom */);
+        sDisplayContent.forAllWindows(w -> {windows.add(w);}, true /* traverseTopToBottom */);
 
-        assertEquals(wallpaperWindow, windows.get(8));
-        assertEquals(exitingAppWindow, windows.get(7));
-        assertEquals(negChildAppWindow, windows.get(6));
-        assertEquals(appWindow, windows.get(5));
-        assertEquals(posChildAppWindow, windows.get(4));
-        assertEquals(statusBarWindow, windows.get(3));
-        assertEquals(navBarWindow, windows.get(2));
-        assertEquals(imeWindow, windows.get(1));
-        assertEquals(imeDialogWindow, windows.get(0));
+        assertEquals(sWallpaperWindow, windows.get(9));
+        assertEquals(exitingAppWindow, windows.get(8));
+        assertEquals(sChildAppWindowBelow, windows.get(7));
+        assertEquals(sAppWindow, windows.get(6));
+        assertEquals(sChildAppWindowAbove, windows.get(5));
+        assertEquals(sDockedDividerWindow, windows.get(4));
+        assertEquals(sStatusBarWindow, windows.get(3));
+        assertEquals(sNavBarWindow, windows.get(2));
+        assertEquals(sImeWindow, windows.get(1));
+        assertEquals(sImeDialogWindow, windows.get(0));
+    }
+
+    @Test
+    public void testForAllWindows_WithImeTarget() throws Exception {
+        final WindowState imeAppTarget =
+                createWindow(null, TYPE_BASE_APPLICATION, sDisplayContent, "imeAppTarget");
+
+        sWm.mInputMethodTarget = imeAppTarget;
+
+        final ArrayList<WindowState> windows = new ArrayList();
+
+        // Test forward traversal.
+        sDisplayContent.forAllWindows(w -> {windows.add(w);}, false /* traverseTopToBottom */);
+
+        assertEquals(sWallpaperWindow, windows.get(0));
+        assertEquals(sChildAppWindowBelow, windows.get(1));
+        assertEquals(sAppWindow, windows.get(2));
+        assertEquals(sChildAppWindowAbove, windows.get(3));
+        assertEquals(imeAppTarget, windows.get(4));
+        assertEquals(sImeWindow, windows.get(5));
+        assertEquals(sImeDialogWindow, windows.get(6));
+        assertEquals(sDockedDividerWindow, windows.get(7));
+        assertEquals(sStatusBarWindow, windows.get(8));
+        assertEquals(sNavBarWindow, windows.get(9));
+
+        // Test backward traversal.
+        windows.clear();
+        sDisplayContent.forAllWindows(w -> {windows.add(w);}, true /* traverseTopToBottom */);
+
+        assertEquals(sWallpaperWindow, windows.get(9));
+        assertEquals(sChildAppWindowBelow, windows.get(8));
+        assertEquals(sAppWindow, windows.get(7));
+        assertEquals(sChildAppWindowAbove, windows.get(6));
+        assertEquals(imeAppTarget, windows.get(5));
+        assertEquals(sImeWindow, windows.get(4));
+        assertEquals(sImeDialogWindow, windows.get(3));
+        assertEquals(sDockedDividerWindow, windows.get(2));
+        assertEquals(sStatusBarWindow, windows.get(1));
+        assertEquals(sNavBarWindow, windows.get(0));
+
+        // Clean-up
+        sWm.mInputMethodTarget = null;
+        imeAppTarget.removeImmediately();
     }
 }
diff --git a/services/tests/servicestests/src/com/android/server/wm/WindowLayersControllerTests.java b/services/tests/servicestests/src/com/android/server/wm/WindowLayersControllerTests.java
index b9fc3ff..c8650bf 100644
--- a/services/tests/servicestests/src/com/android/server/wm/WindowLayersControllerTests.java
+++ b/services/tests/servicestests/src/com/android/server/wm/WindowLayersControllerTests.java
@@ -16,7 +16,6 @@
 
 package com.android.server.wm;
 
-import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
@@ -27,13 +26,7 @@
 import static android.view.WindowManager.LayoutParams.TYPE_APPLICATION_ATTACHED_DIALOG;
 import static android.view.WindowManager.LayoutParams.TYPE_APPLICATION_MEDIA_OVERLAY;
 import static android.view.WindowManager.LayoutParams.TYPE_BASE_APPLICATION;
-import static android.view.WindowManager.LayoutParams.TYPE_DOCK_DIVIDER;
-import static android.view.WindowManager.LayoutParams.TYPE_INPUT_METHOD;
-import static android.view.WindowManager.LayoutParams.TYPE_INPUT_METHOD_DIALOG;
-import static android.view.WindowManager.LayoutParams.TYPE_NAVIGATION_BAR;
-import static android.view.WindowManager.LayoutParams.TYPE_STATUS_BAR;
 import static android.view.WindowManager.LayoutParams.TYPE_SYSTEM_OVERLAY;
-import static android.view.WindowManager.LayoutParams.TYPE_WALLPAPER;
 
 /**
  * Tests for the {@link WindowLayersController} class.
@@ -46,43 +39,6 @@
 @RunWith(AndroidJUnit4.class)
 public class WindowLayersControllerTests extends WindowTestsBase {
 
-    private static boolean sOneTimeSetupDone = false;
-    private static WindowLayersController sLayersController;
-    private static DisplayContent sDisplayContent;
-    private static WindowState sImeWindow;
-    private static WindowState sImeDialogWindow;
-    private static WindowState sStatusBarWindow;
-    private static WindowState sDockedDividerWindow;
-    private static WindowState sNavBarWindow;
-    private static WindowState sAppWindow;
-    private static WindowState sChildAppWindow;
-
-    @Before
-    public void setUp() throws Exception {
-        super.setUp();
-
-        if (sOneTimeSetupDone) {
-            return;
-        }
-        sOneTimeSetupDone = true;
-        sLayersController = new WindowLayersController(sWm);
-        sDisplayContent =
-                new DisplayContent(mDisplay, sWm, sLayersController, new WallpaperController(sWm));
-        final WindowState wallpaperWindow =
-                createWindow(null, TYPE_WALLPAPER, sDisplayContent, "wallpaperWindow");
-        sImeWindow = createWindow(null, TYPE_INPUT_METHOD, sDisplayContent, "sImeWindow");
-        sImeDialogWindow =
-                createWindow(null, TYPE_INPUT_METHOD_DIALOG, sDisplayContent, "sImeDialogWindow");
-        sStatusBarWindow = createWindow(null, TYPE_STATUS_BAR, sDisplayContent, "sStatusBarWindow");
-        sNavBarWindow =
-                createWindow(null, TYPE_NAVIGATION_BAR, sStatusBarWindow.mToken, "sNavBarWindow");
-        sDockedDividerWindow =
-                createWindow(null, TYPE_DOCK_DIVIDER, sDisplayContent, "sDockedDividerWindow");
-        sAppWindow = createWindow(null, TYPE_BASE_APPLICATION, sDisplayContent, "sAppWindow");
-        sChildAppWindow = createWindow(sAppWindow,
-                TYPE_APPLICATION_ATTACHED_DIALOG, sAppWindow.mToken, "sChildAppWindow");
-    }
-
     @Test
     public void testAssignWindowLayers_ForImeWithNoTarget() throws Exception {
         sWm.mInputMethodTarget = null;
@@ -91,7 +47,7 @@
         // The Ime has an higher base layer than app windows and lower base layer than system
         // windows, so it should be above app windows and below system windows if there isn't an IME
         // target.
-        assertWindowLayerGreaterThan(sImeWindow, sChildAppWindow);
+        assertWindowLayerGreaterThan(sImeWindow, sChildAppWindowAbove);
         assertWindowLayerGreaterThan(sImeWindow, sAppWindow);
         assertWindowLayerGreaterThan(sImeWindow, sDockedDividerWindow);
         assertWindowLayerGreaterThan(sNavBarWindow, sImeWindow);
@@ -111,7 +67,7 @@
         // Ime should be above all app windows and below system windows if it is targeting an app
         // window.
         assertWindowLayerGreaterThan(sImeWindow, imeAppTarget);
-        assertWindowLayerGreaterThan(sImeWindow, sChildAppWindow);
+        assertWindowLayerGreaterThan(sImeWindow, sChildAppWindowAbove);
         assertWindowLayerGreaterThan(sImeWindow, sAppWindow);
         assertWindowLayerGreaterThan(sImeWindow, sDockedDividerWindow);
         assertWindowLayerGreaterThan(sNavBarWindow, sImeWindow);
@@ -140,7 +96,7 @@
         assertWindowLayerGreaterThan(sImeWindow, imeAppTarget);
         assertWindowLayerGreaterThan(imeAppTargetChildAboveWindow, sImeWindow);
         assertWindowLayerGreaterThan(sImeWindow, imeAppTargetChildBelowWindow);
-        assertWindowLayerGreaterThan(sImeWindow, sChildAppWindow);
+        assertWindowLayerGreaterThan(sImeWindow, sChildAppWindowAbove);
         assertWindowLayerGreaterThan(sImeWindow, sAppWindow);
         assertWindowLayerGreaterThan(sImeWindow, sDockedDividerWindow);
         assertWindowLayerGreaterThan(sNavBarWindow, sImeWindow);
@@ -167,7 +123,7 @@
         assertWindowLayerGreaterThan(sImeWindow, imeAppTarget);
         assertWindowLayerGreaterThan(sImeWindow, appBelowImeTarget);
         assertWindowLayerGreaterThan(appAboveImeTarget, sImeWindow);
-        assertWindowLayerGreaterThan(sImeWindow, sChildAppWindow);
+        assertWindowLayerGreaterThan(sImeWindow, sChildAppWindowAbove);
         assertWindowLayerGreaterThan(sImeWindow, sAppWindow);
         assertWindowLayerGreaterThan(sImeWindow, sDockedDividerWindow);
         assertWindowLayerGreaterThan(sNavBarWindow, sImeWindow);
@@ -188,7 +144,7 @@
         // The IME target base layer is higher than all window except for the nav bar window, so the
         // IME should be above all windows except for the nav bar.
         assertWindowLayerGreaterThan(sImeWindow, imeSystemOverlayTarget);
-        assertWindowLayerGreaterThan(sImeWindow, sChildAppWindow);
+        assertWindowLayerGreaterThan(sImeWindow, sChildAppWindowAbove);
         assertWindowLayerGreaterThan(sImeWindow, sAppWindow);
         assertWindowLayerGreaterThan(sImeWindow, sDockedDividerWindow);
         assertWindowLayerGreaterThan(sImeWindow, sStatusBarWindow);
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 9681bd2..e301b19 100644
--- a/services/tests/servicestests/src/com/android/server/wm/WindowTestsBase.java
+++ b/services/tests/servicestests/src/com/android/server/wm/WindowTestsBase.java
@@ -22,7 +22,6 @@
 import android.content.Context;
 import android.os.IBinder;
 import android.support.test.InstrumentationRegistry;
-import android.view.Display;
 import android.view.IWindow;
 import android.view.WindowManager;
 
@@ -31,6 +30,15 @@
 import static android.content.res.Configuration.EMPTY;
 import static android.view.WindowManager.LayoutParams.FIRST_APPLICATION_WINDOW;
 import static android.view.WindowManager.LayoutParams.LAST_APPLICATION_WINDOW;
+import static android.view.WindowManager.LayoutParams.TYPE_APPLICATION_ATTACHED_DIALOG;
+import static android.view.WindowManager.LayoutParams.TYPE_APPLICATION_MEDIA_OVERLAY;
+import static android.view.WindowManager.LayoutParams.TYPE_BASE_APPLICATION;
+import static android.view.WindowManager.LayoutParams.TYPE_DOCK_DIVIDER;
+import static android.view.WindowManager.LayoutParams.TYPE_INPUT_METHOD;
+import static android.view.WindowManager.LayoutParams.TYPE_INPUT_METHOD_DIALOG;
+import static android.view.WindowManager.LayoutParams.TYPE_NAVIGATION_BAR;
+import static android.view.WindowManager.LayoutParams.TYPE_STATUS_BAR;
+import static android.view.WindowManager.LayoutParams.TYPE_WALLPAPER;
 import static org.mockito.Mockito.mock;
 
 /**
@@ -40,15 +48,49 @@
     static WindowManagerService sWm = null;
     private final IWindow mIWindow = new TestIWindow();
     private final Session mMockSession = mock(Session.class);
-    Display mDisplay;
     private static int sNextStackId = FIRST_DYNAMIC_STACK_ID;
     private static int sNextTaskId = 0;
 
+    private static boolean sOneTimeSetupDone = false;
+    protected static DisplayContent sDisplayContent;
+    protected static WindowLayersController sLayersController;
+    protected static WindowState sWallpaperWindow;
+    protected static WindowState sImeWindow;
+    protected static WindowState sImeDialogWindow;
+    protected static WindowState sStatusBarWindow;
+    protected static WindowState sDockedDividerWindow;
+    protected static WindowState sNavBarWindow;
+    protected static WindowState sAppWindow;
+    protected static WindowState sChildAppWindowAbove;
+    protected static WindowState sChildAppWindowBelow;
+
     @Before
     public void setUp() throws Exception {
+        if (sOneTimeSetupDone) {
+            return;
+        }
+        sOneTimeSetupDone = true;
         final Context context = InstrumentationRegistry.getTargetContext();
         sWm = TestWindowManagerPolicy.getWindowManagerService(context);
-        mDisplay = context.getDisplay();
+        sLayersController = new WindowLayersController(sWm);
+        sDisplayContent = new DisplayContent(context.getDisplay(), sWm, sLayersController,
+                new WallpaperController(sWm));
+
+        // Set-up some common windows.
+        sWallpaperWindow = createWindow(null, TYPE_WALLPAPER, sDisplayContent, "wallpaperWindow");
+        sImeWindow = createWindow(null, TYPE_INPUT_METHOD, sDisplayContent, "sImeWindow");
+        sImeDialogWindow =
+                createWindow(null, TYPE_INPUT_METHOD_DIALOG, sDisplayContent, "sImeDialogWindow");
+        sStatusBarWindow = createWindow(null, TYPE_STATUS_BAR, sDisplayContent, "sStatusBarWindow");
+        sNavBarWindow =
+                createWindow(null, TYPE_NAVIGATION_BAR, sStatusBarWindow.mToken, "sNavBarWindow");
+        sDockedDividerWindow =
+                createWindow(null, TYPE_DOCK_DIVIDER, sDisplayContent, "sDockedDividerWindow");
+        sAppWindow = createWindow(null, TYPE_BASE_APPLICATION, sDisplayContent, "sAppWindow");
+        sChildAppWindowAbove = createWindow(sAppWindow,
+                TYPE_APPLICATION_ATTACHED_DIALOG, sAppWindow.mToken, "sChildAppWindowAbove");
+        sChildAppWindowBelow = createWindow(sAppWindow,
+                TYPE_APPLICATION_MEDIA_OVERLAY, sAppWindow.mToken, "sChildAppWindowBelow");
     }
 
     /** Asserts that the first entry is greater than the second entry. */