DO NOT MERGE: Allow any suitable stack for TestActivity

With change ag/3595036, the stack for the TestActivity is no longer
guaranteed to be FULLSCREEN_WORKSPACE_STACK_ID. Instead, check that it
goes on a suitable fullscreen-type stack, and that the same display
contains the Assistant stack.

Bug: 69859187
Test: Passes on VR Device and Phone

Change-Id: I7b721ca1e35a435181fef5f26e3d8654691ea35f
diff --git a/hostsidetests/services/activityandwindowmanager/activitymanager/src/android/server/cts/ActivityManagerAssistantStackTests.java b/hostsidetests/services/activityandwindowmanager/activitymanager/src/android/server/cts/ActivityManagerAssistantStackTests.java
index e1e108d..80a8e5f 100644
--- a/hostsidetests/services/activityandwindowmanager/activitymanager/src/android/server/cts/ActivityManagerAssistantStackTests.java
+++ b/hostsidetests/services/activityandwindowmanager/activitymanager/src/android/server/cts/ActivityManagerAssistantStackTests.java
@@ -85,8 +85,8 @@
         // Dock a task
         launchActivity(TEST_ACTIVITY);
         launchActivityInDockStack(DOCKED_ACTIVITY);
-        mAmWmState.assertContainsStack("Must contain fullscreen stack.",
-                FULLSCREEN_WORKSPACE_STACK_ID);
+        mAmWmState.assertHasFullscreenTypeStackOnDisplay("Must contain fullscreen stack.",
+                mAssistantDisplayId);
         mAmWmState.assertContainsStack("Must contain docked stack.", DOCKED_STACK_ID);
 
         // Enable the assistant and launch an assistant activity, ensure it is on top
@@ -112,8 +112,8 @@
         // Dock a task
         launchActivity(TEST_ACTIVITY);
         launchActivityInDockStack(DOCKED_ACTIVITY);
-        mAmWmState.assertContainsStack("Must contain fullscreen stack.",
-                FULLSCREEN_WORKSPACE_STACK_ID);
+        mAmWmState.assertHasFullscreenTypeStackOnDisplay("Must contain fullscreen stack.",
+                mAssistantDisplayId);
         mAmWmState.assertContainsStack("Must contain docked stack.", DOCKED_STACK_ID);
 
         enableAssistant();
@@ -132,10 +132,9 @@
         // Ensure that the fullscreen stack is on top and the test activity is now visible
         mAmWmState.waitForValidState(mDevice, TEST_ACTIVITY, FULLSCREEN_WORKSPACE_STACK_ID);
         mAmWmState.assertFocusedActivity("TestActivity should be resumed", TEST_ACTIVITY);
-        mAmWmState.assertFrontStack("Fullscreen stack should be on top.",
-                FULLSCREEN_WORKSPACE_STACK_ID);
-        mAmWmState.assertFocusedStack("Fullscreen stack should be focused.",
-                FULLSCREEN_WORKSPACE_STACK_ID);
+        mAmWmState.assertActivityInFullscreenTypeStack(
+                "TestActivity should go on a suitable fullscreen-type stack.",
+                TEST_ACTIVITY);
 
         // Now, tell it to finish itself and ensure that the assistant stack is brought back forward
         executeShellCommand("am broadcast -a " + TEST_ACTIVITY_ACTION_FINISH_SELF);
@@ -155,10 +154,9 @@
         mAmWmState.waitForValidState(mDevice, TEST_ACTIVITY, FULLSCREEN_WORKSPACE_STACK_ID);
         mAmWmState.waitForActivityState(mDevice, TEST_ACTIVITY, STATE_RESUMED);
         mAmWmState.assertFocusedActivity("TestActivity should be resumed", TEST_ACTIVITY);
-        mAmWmState.assertFrontStack("Fullscreen stack should be on top.",
-                FULLSCREEN_WORKSPACE_STACK_ID);
-        mAmWmState.assertFocusedStack("Fullscreen stack should be focused.",
-                FULLSCREEN_WORKSPACE_STACK_ID);
+        mAmWmState.assertActivityInFullscreenTypeStack(
+                "TestActivity should go on a suitable fullscreen-type stack.",
+                TEST_ACTIVITY);
     }
 
     public void testDisallowEnterPiPFromAssistantStack() throws Exception {
diff --git a/hostsidetests/services/activityandwindowmanager/util/src/android/server/cts/ActivityAndWindowManagersState.java b/hostsidetests/services/activityandwindowmanager/util/src/android/server/cts/ActivityAndWindowManagersState.java
index df16346..0f232e7 100644
--- a/hostsidetests/services/activityandwindowmanager/util/src/android/server/cts/ActivityAndWindowManagersState.java
+++ b/hostsidetests/services/activityandwindowmanager/util/src/android/server/cts/ActivityAndWindowManagersState.java
@@ -459,6 +459,18 @@
         assertTrue(msg, mWmState.containsStack(stackId));
     }
 
+    void assertHasFullscreenTypeStackOnDisplay(String msg, int displayId) {
+        assertTrue(msg, mAmState.containsSuitableStackOnDisplay(displayId));
+    }
+
+    void assertActivityInFullscreenTypeStack(String msg, String activityName)
+            throws Exception {
+        final int stackId = mAmState.getTaskByActivityName(activityName).mStackId;
+        assertTrue(mAmState.isSuitableStack(stackId));
+        assertFrontStack(msg, stackId);
+        assertFocusedStack(msg, stackId);
+    }
+
     void assertDoesNotContainStack(String msg, int stackId) throws Exception {
         assertFalse(msg, mAmState.containsStack(stackId));
         assertFalse(msg, mWmState.containsStack(stackId));
diff --git a/hostsidetests/services/activityandwindowmanager/util/src/android/server/cts/ActivityManagerState.java b/hostsidetests/services/activityandwindowmanager/util/src/android/server/cts/ActivityManagerState.java
index b8f6811..2e53997 100644
--- a/hostsidetests/services/activityandwindowmanager/util/src/android/server/cts/ActivityManagerState.java
+++ b/hostsidetests/services/activityandwindowmanager/util/src/android/server/cts/ActivityManagerState.java
@@ -33,8 +33,7 @@
 import java.util.regex.Pattern;
 import java.util.regex.Matcher;
 
-import static android.server.cts.ActivityManagerTestBase.HOME_STACK_ID;
-import static android.server.cts.ActivityManagerTestBase.RECENTS_STACK_ID;
+import static android.server.cts.ActivityManagerTestBase.*;
 import static android.server.cts.StateLogger.log;
 import static android.server.cts.StateLogger.logE;
 
@@ -235,6 +234,21 @@
         return null;
     }
 
+    boolean containsSuitableStackOnDisplay(int displayId) {
+        for (ActivityStack stack : mDisplayStacks.get(displayId)) {
+            if (isSuitableStack(stack.mStackId)) {
+                return true;
+            }
+        }
+        return false;
+    }
+
+    boolean isSuitableStack(int stackId) {
+        return stackId != HOME_STACK_ID
+                && stackId != DOCKED_STACK_ID
+                && stackId != PINNED_STACK_ID;
+    }
+
     int getStackPosition(int stackId) {
         for (Integer displayId : mDisplayStacks.keySet()) {
             List<ActivityStack> stacks = mDisplayStacks.get(displayId);