Wait for windows that are in WM but not in AM.

It is possible that between request of states from AM and
WM activity will show up, which will result in window reported
in WM state and no corresponding activity in AM state.

Bug: 28786255
Change-Id: I196e03b69f548a8862299a78d53cb70f247b2e76
diff --git a/hostsidetests/services/activitymanager/src/android/server/cts/ActivityAndWindowManagersState.java b/hostsidetests/services/activitymanager/src/android/server/cts/ActivityAndWindowManagersState.java
index ee4afbb..a1583d9d 100644
--- a/hostsidetests/services/activitymanager/src/android/server/cts/ActivityAndWindowManagersState.java
+++ b/hostsidetests/services/activitymanager/src/android/server/cts/ActivityAndWindowManagersState.java
@@ -326,6 +326,14 @@
                     return false;
                 }
             }
+
+            for (WindowTask wTask : wStack.mTasks) {
+                if (aStack.getTask(wTask.mTaskId) == null) {
+                    log("Task is in WM but not in AM, waiting for it to settle, taskId="
+                            + wTask.mTaskId);
+                    return false;
+                }
+            }
         }
         return true;
     }
diff --git a/hostsidetests/services/activitymanager/src/android/server/cts/ActivityManagerFreeformStackTests.java b/hostsidetests/services/activitymanager/src/android/server/cts/ActivityManagerFreeformStackTests.java
index f86d9ec..abe067f 100644
--- a/hostsidetests/services/activitymanager/src/android/server/cts/ActivityManagerFreeformStackTests.java
+++ b/hostsidetests/services/activitymanager/src/android/server/cts/ActivityManagerFreeformStackTests.java
@@ -36,7 +36,7 @@
 
         launchActivityInStack(FREEFORM_ACTIVITY, FREEFORM_WORKSPACE_STACK_ID);
 
-        mAmWmState.computeState(mDevice, new String[] {FREEFORM_ACTIVITY});
+        mAmWmState.computeState(mDevice, new String[] {FREEFORM_ACTIVITY, TEST_ACTIVITY});
 
         if (!supportsFreeform()) {
             mAmWmState.assertDoesNotContainStack(
diff --git a/hostsidetests/services/activitymanager/src/android/server/cts/ActivityManagerState.java b/hostsidetests/services/activitymanager/src/android/server/cts/ActivityManagerState.java
index 7e8ee04..d392e00 100644
--- a/hostsidetests/services/activitymanager/src/android/server/cts/ActivityManagerState.java
+++ b/hostsidetests/services/activitymanager/src/android/server/cts/ActivityManagerState.java
@@ -339,6 +339,15 @@
         List<ActivityTask> getTasks() {
             return new ArrayList(mTasks);
         }
+
+        ActivityTask getTask(int taskId) {
+            for (ActivityTask task : mTasks) {
+                if (taskId == task.mTaskId) {
+                    return task;
+                }
+            }
+            return null;
+        }
     }
 
     static class ActivityTask extends ActivityContainer {