Have forAllWindows return IME windows in order with IME target

If there is an IME target make sure to return the IME windows next
to the IME target when processing WindowContainer.forAllWindows().
Several users of forAllWindows are expecting this to be the cases so
it is close to the visual order.

Test: bit FrameworksServicesTests:com.android.server.wm.DisplayContentTests
Test: bit FrameworksServicesTests:com.android.server.wm.WindowLayersControllerTests
Change-Id: I737f88ca607ab2694391419d0d38060c03b53840
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. */