Wait until the transition completes before checking lifecycle state

- The system currently sends the configuration changed callback after the
  transition down completes. Update the CTS test to ensure that we wait
  for the stack bounds to settle before trying to check the number of
  configuration change callbacks.

Bug: 64201795
Test: android.server.cts.ActivityManagerPinnedStackTests
Test: #testConfigurationChangeOrderDuringTransition
Change-Id: Ia6ace82020cad555da0b77e0a0f8913cffb2fcbe
diff --git a/hostsidetests/services/activityandwindowmanager/activitymanager/src/android/server/cts/ActivityManagerPinnedStackTests.java b/hostsidetests/services/activityandwindowmanager/activitymanager/src/android/server/cts/ActivityManagerPinnedStackTests.java
index 55b41c8..8e960d0 100644
--- a/hostsidetests/services/activityandwindowmanager/activitymanager/src/android/server/cts/ActivityManagerPinnedStackTests.java
+++ b/hostsidetests/services/activityandwindowmanager/activitymanager/src/android/server/cts/ActivityManagerPinnedStackTests.java
@@ -768,12 +768,14 @@
         executeShellCommand("am broadcast -a " + PIP_ACTIVITY_ACTION_ENTER_PIP);
         mAmWmState.waitForValidState(mDevice, PIP_ACTIVITY, PINNED_STACK_ID);
         assertPinnedStackExists();
+        waitForValidPictureInPictureCallbacks(PIP_ACTIVITY, logSeparator);
         assertValidPictureInPictureCallbackOrder(PIP_ACTIVITY, logSeparator);
 
         // Trigger it to go back to fullscreen and ensure that only triggered one configuration
         // change as well
         logSeparator = clearLogcat();
         launchActivity(PIP_ACTIVITY);
+        waitForValidPictureInPictureCallbacks(PIP_ACTIVITY, logSeparator);
         assertValidPictureInPictureCallbackOrder(PIP_ACTIVITY, logSeparator);
     }
 
@@ -1185,6 +1187,24 @@
     }
 
     /**
+     * Waits until the expected picture-in-picture callbacks have been made.
+     */
+    private void waitForValidPictureInPictureCallbacks(String activityName, String logSeparator)
+            throws Exception {
+        mAmWmState.waitFor(mDevice, (amState, wmState) -> {
+            try {
+                final ActivityLifecycleCounts lifecycleCounts = new ActivityLifecycleCounts(
+                        activityName, logSeparator);
+                return lifecycleCounts.mConfigurationChangedCount == 1 &&
+                        lifecycleCounts.mPictureInPictureModeChangedCount == 1 &&
+                        lifecycleCounts.mMultiWindowModeChangedCount == 1;
+            } catch (Exception e) {
+                return false;
+            }
+        }, "Waiting for picture-in-picture activity callbacks...");
+    }
+
+    /**
      * @return the window state for the given {@param activity}'s window.
      */
     private WindowManagerState.WindowState getWindowState(String activity) throws Exception {