DO NOT MERGE ANYWHERE Use DisplayId when verifying top task in MulitDisplay MultiDisplay does not always use a specific DisplayId when verifying top task which can lead to incorrect verification as they can differ of course. Test: atest CtsWindowManagerDeviceTestCases:android.server.wm.MultiDisplayClientTests Test: atest CtsWindowManagerDeviceTestCases:android.server.wm.MultiDisplayActivityLaunchTests Flag: TEST_ONLY Bug: 368604141 (cherry picked from https://googleplex-android-review.googlesource.com/q/commit:0aea63219254b79fa52f28bd13de930ec3ebd47e) Change-Id: I3579d00775e62484cea51768a5d7e85b763c553c

diff --git a/tests/framework/base/windowmanager/util/src/android/server/wm/ActivityManagerTestBase.java b/tests/framework/base/windowmanager/util/src/android/server/wm/ActivityManagerTestBase.java
index 40ef1ca..af629c0 100644
--- a/tests/framework/base/windowmanager/util/src/android/server/wm/ActivityManagerTestBase.java
+++ b/tests/framework/base/windowmanager/util/src/android/server/wm/ActivityManagerTestBase.java

@@ -1316,10 +1316,11 @@
     public void waitAndAssertTopResumedActivity(ComponentName activityName, int displayId,
             String message) {
         final String activityClassName = getActivityName(activityName);
-        mWmState.waitForWithAmState(state -> activityClassName.equals(state.getFocusedActivity()),
+        mWmState.waitForWithAmState(state ->
+                activityClassName.equals(state.getFocusedActivityOnDisplay(displayId)),
                 "activity to be on top");
         waitAndAssertResumedActivity(activityName, "Activity must be resumed");
-        mWmState.assertFocusedActivity(message, activityName);
+        mWmState.assertFocusedActivityOnDisplay(message, activityName, displayId);
 
         final int frontRootTaskId = mWmState.getFrontRootTaskId(displayId);
         Task frontRootTaskOnDisplay = mWmState.getRootTask(frontRootTaskId);
@@ -1328,8 +1329,8 @@
                 activityClassName,
                 frontRootTaskOnDisplay.isLeafTask() ? frontRootTaskOnDisplay.mResumedActivity
                         : frontRootTaskOnDisplay.getTopTask().mResumedActivity);
-        mWmState.assertFocusedRootTask("Top activity's rootTask must also be on top",
-                frontRootTaskId);
+        mWmState.assertFocusedRootTaskOnDisplay("Top activity's rootTask must also be on top",
+                frontRootTaskId, displayId);
     }
 
     /**

diff --git a/tests/framework/base/windowmanager/util/src/android/server/wm/WindowManagerState.java b/tests/framework/base/windowmanager/util/src/android/server/wm/WindowManagerState.java
index a3de70e..cbf4eec 100644
--- a/tests/framework/base/windowmanager/util/src/android/server/wm/WindowManagerState.java
+++ b/tests/framework/base/windowmanager/util/src/android/server/wm/WindowManagerState.java

@@ -489,6 +489,10 @@
         mWindowFramesValid = false;
     }
 
+    public String getFocusedAppOnDisplay(int displayId) {
+        return getDisplay(displayId).mFocusedApp;
+    }
+
     public String getFocusedApp() {
         return mFocusedApp;
     }
@@ -599,6 +603,10 @@
         return null;
     }
 
+    int getFocusedTaskIdOnDisplay(int displayId) {
+        return getDisplay(displayId).mFocusedRootTaskId;
+    }
+
     public int getFocusedTaskId() {
         return mTopFocusedTaskId;
     }
@@ -613,6 +621,10 @@
         return rootTask != null ? rootTask.getWindowingMode() : WINDOWING_MODE_UNDEFINED;
     }
 
+    public String getFocusedActivityOnDisplay(int displayId) {
+        return getDisplay(displayId).mResumedActivity;
+    }
+
     public String getFocusedActivity() {
         return mTopResumedActivityRecord;
     }

diff --git a/tests/framework/base/windowmanager/util/src/android/server/wm/WindowManagerStateHelper.java b/tests/framework/base/windowmanager/util/src/android/server/wm/WindowManagerStateHelper.java
index ddd9d27..8823e71 100644
--- a/tests/framework/base/windowmanager/util/src/android/server/wm/WindowManagerStateHelper.java
+++ b/tests/framework/base/windowmanager/util/src/android/server/wm/WindowManagerStateHelper.java

@@ -680,6 +680,10 @@
         assertFrontStackActivityTypeOnDisplay(msg, activityType, DEFAULT_DISPLAY);
     }
 
+    void assertFocusedRootTaskOnDisplay(String msg, int taskId, int displayId) {
+        assertEquals(msg, taskId, getFocusedTaskIdOnDisplay(displayId));
+    }
+
     public void assertFocusedRootTask(String msg, int taskId) {
         assertEquals(msg, taskId, getFocusedTaskId());
     }
@@ -693,6 +697,13 @@
         }
     }
 
+    public void assertFocusedActivityOnDisplay(final String msg, final ComponentName activityName,
+            final int displayId) {
+        final String activityComponentName = getActivityName(activityName);
+        assertEquals(msg, activityComponentName, getFocusedActivityOnDisplay(displayId));
+        assertEquals(msg, activityComponentName, getFocusedAppOnDisplay(displayId));
+    }
+
     public void assertFocusedActivity(final String msg, final ComponentName activityName) {
         final String activityComponentName = getActivityName(activityName);
         assertEquals(msg, activityComponentName, getFocusedActivity());