Merge "Wait for config callback to reduce flakiness" into main
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() {}
+    }
 }