Wait for config callback to reduce flakiness

Previously, we have already wait for config callback from
WindowContexListenerController and WindowTokenClient,
but the test is still flaky. A theory is that the client side received
the config change binder call but the config value hasn't applied to
the Resources.

This CL wait for the config callback via registering ComponentCallbacks
to ensure the Resources value is updated.

Test: atest InputMethodDialogWindowContextTest#testGetSettingsContextOnDualDisplayContent
Bug: 296622353
Change-Id: I4b491d35ae9ce73eb40c8b866ba91e9d7f6442db
diff --git a/services/tests/wmtests/src/com/android/server/wm/InputMethodDialogWindowContextTest.java b/services/tests/wmtests/src/com/android/server/wm/InputMethodDialogWindowContextTest.java
index c3db241..4165911 100644
--- a/services/tests/wmtests/src/com/android/server/wm/InputMethodDialogWindowContextTest.java
+++ b/services/tests/wmtests/src/com/android/server/wm/InputMethodDialogWindowContextTest.java
@@ -26,7 +26,6 @@
 
 import static com.google.common.truth.Truth.assertThat;
 
-import static org.junit.Assert.assertNotNull;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyInt;
 import static org.mockito.ArgumentMatchers.eq;
@@ -35,9 +34,14 @@
 import static org.mockito.Mockito.timeout;
 import static org.mockito.Mockito.verify;
 
+import static java.util.Objects.requireNonNull;
+
 import android.app.ActivityThread;
 import android.app.IApplicationThread;
+import android.app.WindowConfiguration;
+import android.content.ComponentCallbacks;
 import android.content.Context;
+import android.content.res.Configuration;
 import android.content.res.Resources;
 import android.graphics.Rect;
 import android.hardware.display.DisplayManagerGlobal;
@@ -50,6 +54,8 @@
 import android.window.WindowContextInfo;
 import android.window.WindowTokenClient;
 
+import androidx.annotation.NonNull;
+
 import com.android.server.inputmethod.InputMethodDialogWindowContext;
 
 import org.junit.After;
@@ -58,6 +64,9 @@
 import org.junit.runner.RunWith;
 import org.mockito.Mockito;
 
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
+
 // TODO(b/157888351): Move the test to inputmethod package once we find the way to test the
 //  scenario there.
 /**
@@ -138,44 +147,96 @@
     @Test
     public void testGetSettingsContextOnDualDisplayContent() {
         final Context context = mWindowContext.get(mSecondaryDisplay.getDisplayId());
+        final MaxBoundsVerifier maxBoundsVerifier = new MaxBoundsVerifier();
+        context.registerComponentCallbacks(maxBoundsVerifier);
+
         final WindowTokenClient tokenClient = (WindowTokenClient) context.getWindowContextToken();
-        assertNotNull(tokenClient);
-        spyOn(tokenClient);
+        spyOn(requireNonNull(tokenClient));
 
         final DisplayArea.Tokens imeContainer = mSecondaryDisplay.getImeContainer();
         spyOn(imeContainer);
         assertThat(imeContainer.getRootDisplayArea()).isEqualTo(mSecondaryDisplay);
 
-        mSecondaryDisplay.mFirstRoot.placeImeContainer(imeContainer);
+        final DisplayAreaGroup firstDaGroup = mSecondaryDisplay.mFirstRoot;
+        maxBoundsVerifier.setMaxBounds(firstDaGroup.getMaxBounds());
+
+        firstDaGroup.placeImeContainer(imeContainer);
 
         verify(imeContainer, timeout(WAIT_TIMEOUT_MS)).onConfigurationChanged(
-                eq(mSecondaryDisplay.mFirstRoot.getConfiguration()));
+                eq(firstDaGroup.getConfiguration()));
         verify(tokenClient, timeout(WAIT_TIMEOUT_MS)).onConfigurationChanged(
-                eq(mSecondaryDisplay.mFirstRoot.getConfiguration()),
+                eq(firstDaGroup.getConfiguration()),
                 eq(mSecondaryDisplay.mDisplayId));
-        assertThat(imeContainer.getRootDisplayArea()).isEqualTo(mSecondaryDisplay.mFirstRoot);
+        assertThat(imeContainer.getRootDisplayArea()).isEqualTo(firstDaGroup);
+        maxBoundsVerifier.waitAndAssertMaxMetricsMatches();
         assertImeSwitchContextMetricsValidity(context, mSecondaryDisplay);
 
         // Clear the previous invocation histories in case we may count the previous
         // onConfigurationChanged invocation into the next verification.
         clearInvocations(tokenClient, imeContainer);
-        mSecondaryDisplay.mSecondRoot.placeImeContainer(imeContainer);
+        final DisplayAreaGroup secondDaGroup = mSecondaryDisplay.mSecondRoot;
+        maxBoundsVerifier.setMaxBounds(secondDaGroup.getMaxBounds());
+
+        secondDaGroup.placeImeContainer(imeContainer);
 
         verify(imeContainer, timeout(WAIT_TIMEOUT_MS)).onConfigurationChanged(
-                eq(mSecondaryDisplay.mSecondRoot.getConfiguration()));
+                eq(secondDaGroup.getConfiguration()));
         verify(tokenClient, timeout(WAIT_TIMEOUT_MS)).onConfigurationChanged(
-                eq(mSecondaryDisplay.mSecondRoot.getConfiguration()),
+                eq(secondDaGroup.getConfiguration()),
                 eq(mSecondaryDisplay.mDisplayId));
-        assertThat(imeContainer.getRootDisplayArea()).isEqualTo(mSecondaryDisplay.mSecondRoot);
+        assertThat(imeContainer.getRootDisplayArea()).isEqualTo(secondDaGroup);
+        maxBoundsVerifier.waitAndAssertMaxMetricsMatches();
         assertImeSwitchContextMetricsValidity(context, mSecondaryDisplay);
     }
 
     private void assertImeSwitchContextMetricsValidity(Context context, DisplayContent dc) {
         assertThat(context.getDisplayId()).isEqualTo(dc.getDisplayId());
 
+        final Rect imeContainerBounds = dc.getImeContainer().getBounds();
         final Rect contextBounds = context.getSystemService(WindowManager.class)
                 .getMaximumWindowMetrics().getBounds();
-        final Rect imeContainerBounds = dc.getImeContainer().getBounds();
         assertThat(contextBounds).isEqualTo(imeContainerBounds);
     }
+
+    private static final class MaxBoundsVerifier implements ComponentCallbacks {
+
+        private CountDownLatch mLatch;
+
+        private Rect mMaxBounds;
+
+        /**
+         * Sets max bounds to verify whether it matches the
+         * {@link WindowConfiguration#getMaxBounds()} reported from
+         * {@link #onConfigurationChanged(Configuration)} callback, and also resets the count down
+         * latch.
+         *
+         * @param maxBounds max bounds to verify
+         */
+        private void setMaxBounds(@NonNull Rect maxBounds) {
+            mMaxBounds = maxBounds;
+            mLatch = new CountDownLatch(1);
+        }
+
+        /**
+         * Waits for the {@link #onConfigurationChanged(Configuration)} callback whose the reported
+         * {@link WindowConfiguration#getMaxBounds()} matches {@link #mMaxBounds}.
+         */
+        private void waitAndAssertMaxMetricsMatches() {
+            try {
+                assertThat(mLatch.await(WAIT_TIMEOUT_MS, TimeUnit.MILLISECONDS)).isTrue();
+            } catch (InterruptedException e) {
+                throw new RuntimeException("Test failed because of " + e);
+            }
+        }
+
+        @Override
+        public void onConfigurationChanged(@NonNull Configuration newConfig) {
+            if (newConfig.windowConfiguration.getMaxBounds().equals(mMaxBounds)) {
+                mLatch.countDown();
+            }
+        }
+
+        @Override
+        public void onLowMemory() {}
+    }
 }