Wait for the removal of overlay display

If a test creates virtual display and overlay (simulate) display,
before overlay display is removed completely, it may be counted
when creating new virtual displays that leads to assertion failed
by the number of displays.

Now the overlay display session ensures its display is ready when
the setting is set, and remembers the created displays accordingly.
Then it guarantees that the overlay display is removed completely
after the session is closed.

Bug: 134774470
Bug: 136635756
Test: atest DisplayTests#testNonDefaultDisplayResourcesConfiguration

Change-Id: I736eb00cd28357d0411a1bfd2c86038144f9dfe3
Merged-In: I78da48d33f34b42bb3f7b3d12b1c891d23f82b00
diff --git a/tests/framework/base/windowmanager/src/android/server/wm/DisplayTests.java b/tests/framework/base/windowmanager/src/android/server/wm/DisplayTests.java
index 2fc2f1e..f51f0b6 100644
--- a/tests/framework/base/windowmanager/src/android/server/wm/DisplayTests.java
+++ b/tests/framework/base/windowmanager/src/android/server/wm/DisplayTests.java
@@ -73,7 +73,6 @@
     }
 
     @Test
-    @FlakyTest(bugId = 129521230)
     public void testNonDefaultDisplayResourcesConfiguration() throws Exception {
         final int smallDisplaySize = 1000;
         final int longDisplaySize = 1920;
diff --git a/tests/framework/base/windowmanager/src/android/server/wm/MultiDisplayTestBase.java b/tests/framework/base/windowmanager/src/android/server/wm/MultiDisplayTestBase.java
index 32a20ae8..ef78bc7b 100644
--- a/tests/framework/base/windowmanager/src/android/server/wm/MultiDisplayTestBase.java
+++ b/tests/framework/base/windowmanager/src/android/server/wm/MultiDisplayTestBase.java
@@ -74,6 +74,7 @@
 import java.util.Collections;
 import java.util.List;
 import java.util.function.Consumer;
+import java.util.function.Predicate;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
 
@@ -244,6 +245,18 @@
         tapOnDisplay(bounds.centerX(), bounds.centerY(), displayId);
     }
 
+    private void waitForDisplayGone(Predicate<WindowManagerState.Display> displayPredicate) {
+        for (int retry = 1; retry <= 5; retry++) {
+            mAmWmState.computeState(true);
+            if (!mAmWmState.getWmState().getDisplays().stream().anyMatch(displayPredicate::test)) {
+                return;
+            }
+            logAlways("Waiting for hosted displays destruction... retry=" + retry);
+            SystemClock.sleep(500);
+        }
+        fail("Waiting for hosted displays destruction failed.");
+    }
+
     /**
      * This class should only be used when you need to test virtual display created by a
      * non-privileged app.
@@ -360,19 +373,13 @@
          * @return {@link ActivityDisplay} of newly created display.
          */
         private List<ActivityDisplay> simulateDisplay() throws Exception {
-            final List<ActivityDisplay> originalDs = getDisplaysStates();
-
             // Create virtual display with custom density dpi and specified size.
             mOverlayDisplayDeviceSession.set(mSimulationDisplaySize + "/" + mDensityDpi);
-            final List<ActivityDisplay> newDisplays = assertAndGetNewDisplays(1, originalDs);
-
             if (mShowSystemDecorations) {
-                for (ActivityDisplay display : newDisplays) {
-                    mOverlayDisplayDeviceSession.addAndConfigDisplayState(display,
-                            true /* requestShowSysDecors */, true /* requestShowIme */);
-                }
+                mOverlayDisplayDeviceSession.configureDisplays(
+                        true /* requestShowSysDecors */, true /* requestShowIme */);
             }
-            return newDisplays;
+            return mOverlayDisplayDeviceSession.getCreatedDisplays();
         }
 
         /**
@@ -448,55 +455,9 @@
                     + " -f 0x20000000"
                     + " --es " + KEY_COMMAND + " " + COMMAND_DESTROY_DISPLAY;
             executeShellCommand(destroyVirtualDisplayCommand);
-            waitForDisplaysDestroyed();
-        }
-
-        private void waitForDisplaysDestroyed() {
-            for (int retry = 1; retry <= 5; retry++) {
-                if (!isHostedVirtualDisplayPresent()) {
-                    return;
-                }
-                logAlways("Waiting for hosted displays destruction... retry=" + retry);
-                SystemClock.sleep(500);
-            }
-            fail("Waiting for hosted displays destruction failed.");
-        }
-
-        private boolean isHostedVirtualDisplayPresent() {
-            mAmWmState.computeState(true);
-            return mAmWmState.getWmState().getDisplays().stream().anyMatch(
+            waitForDisplayGone(
                     d -> d.getName() != null && d.getName().contains(VIRTUAL_DISPLAY_PREFIX));
         }
-
-        /**
-         * Wait for desired number of displays to be created and get their properties.
-         * @param newDisplayCount expected display count, -1 if display should not be created.
-         * @param originalDS display states before creation of new display(s).
-         * @return list of new displays, empty list if no new display is created.
-         */
-        private List<ActivityDisplay> assertAndGetNewDisplays(int newDisplayCount,
-                List<ActivityDisplay> originalDS) {
-            final int originalDisplayCount = originalDS.size();
-
-            // Wait for the display(s) to be created and get configurations.
-            final List<ActivityDisplay> ds = getDisplayStateAfterChange(
-                    originalDisplayCount + newDisplayCount);
-            if (newDisplayCount != -1) {
-                assertEquals("New virtual display(s) must be created",
-                        originalDisplayCount + newDisplayCount, ds.size());
-            } else {
-                assertEquals("New virtual display must not be created",
-                        originalDisplayCount, ds.size());
-                return Collections.emptyList();
-            }
-
-            // Find the newly added display(s).
-            final List<ActivityDisplay> newDisplays = findNewDisplayStates(originalDS, ds);
-            assertThat("New virtual display must be created",
-                    newDisplays, hasSize(newDisplayCount));
-
-            return newDisplays;
-        }
     }
 
     // TODO(b/112837428): Merge into VirtualDisplaySession when all usages are migrated.
@@ -538,7 +499,10 @@
     }
 
     /** Helper class to save, set, and restore overlay_display_devices preference. */
-    private static class OverlayDisplayDevicesSession extends SettingsSession<String> {
+    private class OverlayDisplayDevicesSession extends SettingsSession<String> {
+        /** The displays which are created by this session. */
+        private final List<ActivityDisplay> mDisplays = new ArrayList<>();
+        /** The configured displays that need to be restored when this session is closed. */
         private final List<OverlayDisplayState> mDisplayStates = new ArrayList<>();
         private final WindowManager mWm;
 
@@ -549,23 +513,38 @@
             mWm = context.getSystemService(WindowManager.class);
         }
 
-        void addAndConfigDisplayState(ActivityDisplay display, boolean requestShowSysDecors,
-                boolean requestShowIme) {
+        List<ActivityDisplay> getCreatedDisplays() {
+            return new ArrayList<>(mDisplays);
+        }
+
+        @Override
+        public void set(String value) {
+            final List<ActivityDisplay> originalDisplays = getDisplaysStates();
+            super.set(value);
+            final int newDisplayCount = 1 + (int) value.chars().filter(ch -> ch == ';').count();
+            mDisplays.addAll(assertAndGetNewDisplays(newDisplayCount, originalDisplays));
+        }
+
+        void configureDisplays(boolean requestShowSysDecors, boolean requestShowIme) {
             SystemUtil.runWithShellPermissionIdentity(() -> {
-                final boolean showSystemDecors = mWm.shouldShowSystemDecors(display.mId);
-                final boolean showIme = mWm.shouldShowIme(display.mId);
-                mDisplayStates.add(new OverlayDisplayState(display.mId, showSystemDecors, showIme));
-                if (requestShowSysDecors != showSystemDecors) {
-                    mWm.setShouldShowSystemDecors(display.mId, requestShowSysDecors);
-                    TestUtils.waitUntil("Waiting for display show system decors",
-                            5 /* timeoutSecond */,
-                            () -> mWm.shouldShowSystemDecors(display.mId) == requestShowSysDecors);
-                }
-                if (requestShowIme != showIme) {
-                    mWm.setShouldShowIme(display.mId, requestShowIme);
-                    TestUtils.waitUntil("Waiting for display show Ime",
-                            5 /* timeoutSecond */,
-                            () -> mWm.shouldShowIme(display.mId) == requestShowIme);
+                for (ActivityDisplay display : mDisplays) {
+                    final boolean showSystemDecors = mWm.shouldShowSystemDecors(display.mId);
+                    final boolean showIme = mWm.shouldShowIme(display.mId);
+                    mDisplayStates.add(new OverlayDisplayState(
+                            display.mId, showSystemDecors, showIme));
+                    if (requestShowSysDecors != showSystemDecors) {
+                        mWm.setShouldShowSystemDecors(display.mId, requestShowSysDecors);
+                        TestUtils.waitUntil("Waiting for display show system decors",
+                                5 /* timeoutSecond */,
+                                () -> mWm.shouldShowSystemDecors(
+                                        display.mId) == requestShowSysDecors);
+                    }
+                    if (requestShowIme != showIme) {
+                        mWm.setShouldShowIme(display.mId, requestShowIme);
+                        TestUtils.waitUntil("Waiting for display show Ime",
+                                5 /* timeoutSecond */,
+                                () -> mWm.shouldShowIme(display.mId) == requestShowIme);
+                    }
                 }
             });
         }
@@ -587,6 +566,8 @@
             // Need to restore display state before display is destroyed.
             restoreDisplayStates();
             super.close();
+            waitForDisplayGone(display -> mDisplays.stream()
+                    .anyMatch(state -> state.mId == display.getDisplayId()));
         }
 
         private class OverlayDisplayState {
@@ -632,6 +613,36 @@
         return true;
     }
 
+    /**
+     * Wait for desired number of displays to be created and get their properties.
+     *
+     * @param newDisplayCount expected display count, -1 if display should not be created.
+     * @param originalDisplays display states before creation of new display(s).
+     * @return list of new displays, empty list if no new display is created.
+     */
+    private List<ActivityDisplay> assertAndGetNewDisplays(int newDisplayCount,
+            List<ActivityDisplay> originalDisplays) {
+        final int originalDisplayCount = originalDisplays.size();
+
+        // Wait for the display(s) to be created and get configurations.
+        final List<ActivityDisplay> ds = getDisplayStateAfterChange(
+                originalDisplayCount + newDisplayCount);
+        if (newDisplayCount != -1) {
+            assertEquals("New virtual display(s) must be created",
+                    originalDisplayCount + newDisplayCount, ds.size());
+        } else {
+            assertEquals("New virtual display must not be created",
+                    originalDisplayCount, ds.size());
+            return Collections.emptyList();
+        }
+
+        // Find the newly added display(s).
+        final List<ActivityDisplay> newDisplays = findNewDisplayStates(originalDisplays, ds);
+        assertThat("New virtual display must be created", newDisplays, hasSize(newDisplayCount));
+
+        return newDisplays;
+    }
+
     /** Checks if the device supports multi-display. */
     protected boolean supportsMultiDisplay() {
         return hasDeviceFeature(FEATURE_ACTIVITIES_ON_SECONDARY_DISPLAYS);
@@ -714,27 +725,10 @@
                 mExternalDisplayHelper.releaseDisplay();
                 mExternalDisplayHelper = null;
 
-                waitForHostedDisplayDestroyed();
+                waitForDisplayGone(d -> d.getDisplayId() == mDisplayId);
                 mDisplayId = INVALID_DISPLAY;
             }
         }
-
-        private void waitForHostedDisplayDestroyed() {
-            for (int retry = 1; retry <= 5; retry++) {
-                if (!isHostedVirtualDisplayPresent()) {
-                    return;
-                }
-                logAlways("Waiting for hosted displays destruction... retry=" + retry);
-                SystemClock.sleep(500);
-            }
-            fail("Waiting for hosted displays destruction failed.");
-        }
-
-        private boolean isHostedVirtualDisplayPresent() {
-            mAmWmState.computeState(true);
-            return mAmWmState.getWmState().getDisplays().stream().anyMatch(
-                    d -> d.getDisplayId() == mDisplayId);
-        }
     }
 
     public static class PrimaryDisplayStateSession implements AutoCloseable {