Refactor VirtualDisplay creation and cleanup to its own utils

Bug: 365512241
Test: atest WmTests:TaskStackChangedListenerTest
Flag: EXEMPT test cleanup
Change-Id: Ie247dc30a2ce5b789b6c695334214917137362a7

2 files changed, 98 insertions, 27 deletions

diff --git a/services/tests/wmtests/src/com/android/server/wm/TaskStackChangedListenerTest.java b/services/tests/wmtests/src/com/android/server/wm/TaskStackChangedListenerTest.java
index c876663dd749..8a068cc7837a 100644
--- a/services/tests/wmtests/src/com/android/server/wm/TaskStackChangedListenerTest.java
+++ b/services/tests/wmtests/src/com/android/server/wm/TaskStackChangedListenerTest.java
@@ -16,9 +16,6 @@
 
 package com.android.server.wm;
 
-import static android.hardware.display.DisplayManager.VIRTUAL_DISPLAY_FLAG_OWN_CONTENT_ONLY;
-import static android.hardware.display.DisplayManager.VIRTUAL_DISPLAY_FLAG_PRESENTATION;
-import static android.hardware.display.DisplayManager.VIRTUAL_DISPLAY_FLAG_PUBLIC;
 import static android.os.Build.HW_TIMEOUT_MULTIPLIER;
 
 import static androidx.test.platform.app.InstrumentationRegistry.getInstrumentation;
@@ -41,10 +38,7 @@ import android.content.Context;
 import android.content.Intent;
 import android.content.pm.ActivityInfo;
 import android.graphics.Color;
-import android.graphics.PixelFormat;
-import android.hardware.display.DisplayManager;
 import android.hardware.display.VirtualDisplay;
-import android.media.ImageReader;
 import android.os.Bundle;
 import android.os.RemoteException;
 import android.os.SystemClock;
@@ -57,13 +51,14 @@ import android.widget.LinearLayout;
 import androidx.test.filters.MediumTest;
 
 import com.android.server.wm.utils.CommonUtils;
+import com.android.server.wm.utils.VirtualDisplayTestRule;
 
 import org.junit.After;
 import org.junit.Before;
+import org.junit.Rule;
 import org.junit.Test;
 
 import java.util.ArrayList;
-import java.util.Arrays;
 import java.util.concurrent.ArrayBlockingQueue;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
@@ -76,9 +71,9 @@ import java.util.function.Predicate;
 @MediumTest
 public class TaskStackChangedListenerTest {
 
+    @Rule
+    public VirtualDisplayTestRule mVirtualDisplayTestRule = new VirtualDisplayTestRule();
     private ITaskStackListener mTaskStackListener;
-    private VirtualDisplay mVirtualDisplay;
-    private ImageReader mImageReader;
     private final ArrayList<Activity> mStartedActivities = new ArrayList<>();
 
     private static final int WAIT_TIMEOUT_MS = 5000 * HW_TIMEOUT_MULTIPLIER;
@@ -94,10 +89,6 @@ public class TaskStackChangedListenerTest {
         if (mTaskStackListener != null) {
             ActivityTaskManager.getService().unregisterTaskStackListener(mTaskStackListener);
         }
-        if (mVirtualDisplay != null) {
-            mVirtualDisplay.release();
-            mImageReader.close();
-        }
         // Finish from bottom to top.
         final int size = mStartedActivities.size();
         for (int i = 0; i < size; i++) {
@@ -116,21 +107,9 @@ public class TaskStackChangedListenerTest {
     private VirtualDisplay createVirtualDisplay() {
         final int width = 800;
         final int height = 600;
-        final int density = 160;
-        final DisplayManager displayManager = getInstrumentation().getContext().getSystemService(
-                DisplayManager.class);
-        mImageReader = ImageReader.newInstance(width, height, PixelFormat.RGBA_8888,
-                2 /* maxImages */);
-        final int flags = VIRTUAL_DISPLAY_FLAG_PRESENTATION | VIRTUAL_DISPLAY_FLAG_OWN_CONTENT_ONLY
-                | VIRTUAL_DISPLAY_FLAG_PUBLIC;
         final String name = getClass().getSimpleName() + "_VirtualDisplay";
-        mVirtualDisplay = displayManager.createVirtualDisplay(name, width, height, density,
-                mImageReader.getSurface(), flags);
-        mVirtualDisplay.setSurface(mImageReader.getSurface());
-        assertNotNull("display must be registered",
-                Arrays.stream(displayManager.getDisplays()).filter(
-                        d -> d.getName().equals(name)).findAny());
-        return mVirtualDisplay;
+        return mVirtualDisplayTestRule.createDisplayManagerAttachedVirtualDisplay(name, width,
+                height);
     }
 
     @Test
diff --git a/services/tests/wmtests/src/com/android/server/wm/utils/VirtualDisplayTestRule.java b/services/tests/wmtests/src/com/android/server/wm/utils/VirtualDisplayTestRule.java
new file mode 100644
index 000000000000..e92e6846e161
--- /dev/null
+++ b/services/tests/wmtests/src/com/android/server/wm/utils/VirtualDisplayTestRule.java
@@ -0,0 +1,92 @@
+/*
+ * Copyright 2025 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.server.wm.utils;
+
+import static android.hardware.display.DisplayManager.VIRTUAL_DISPLAY_FLAG_OWN_CONTENT_ONLY;
+import static android.hardware.display.DisplayManager.VIRTUAL_DISPLAY_FLAG_PRESENTATION;
+import static android.hardware.display.DisplayManager.VIRTUAL_DISPLAY_FLAG_PUBLIC;
+
+import static androidx.test.platform.app.InstrumentationRegistry.getInstrumentation;
+
+import static org.junit.Assert.assertNotNull;
+
+import android.graphics.PixelFormat;
+import android.hardware.display.DisplayManager;
+import android.hardware.display.VirtualDisplay;
+import android.media.ImageReader;
+
+import org.junit.rules.TestRule;
+import org.junit.runner.Description;
+import org.junit.runners.model.Statement;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.List;
+
+/** Provides wrapper test rule for creating and managing the cleanup for VirtualDisplay */
+public class VirtualDisplayTestRule implements TestRule {
+    private static final int DISPLAY_DENSITY = 160;
+
+    private final List<VirtualDisplay> mVirtualDisplays = new ArrayList<>();
+    private final List<ImageReader> mImageReaders = new ArrayList<>();
+    private final DisplayManager mDisplayManager;
+
+    public VirtualDisplayTestRule() {
+        mDisplayManager = getInstrumentation().getTargetContext().getSystemService(
+                DisplayManager.class);
+    }
+
+    @Override
+    public Statement apply(Statement base, Description description) {
+        return new Statement() {
+            @Override
+            public void evaluate() throws Throwable {
+                try {
+                    base.evaluate();
+                } finally {
+                    tearDown();
+                }
+            }
+        };
+    }
+
+    private void tearDown() {
+        mVirtualDisplays.forEach(VirtualDisplay::release);
+        mImageReaders.forEach(ImageReader::close);
+    }
+
+    /**
+     * The virtual display in WindowTestsBase#createMockSimulatedDisplay is only attached to WM
+     * DisplayWindowSettingsProvider. DisplayManager is not aware of mock simulated display and
+     * therefore couldn't be used for actual Display-related testing (e.g. display listeners).
+     * This method creates real VirtualDisplay through DisplayManager.
+     */
+    public VirtualDisplay createDisplayManagerAttachedVirtualDisplay(String name, int width,
+            int height) {
+        final ImageReader imageReader = ImageReader.newInstance(width, height,
+                PixelFormat.RGBA_8888, 2 /* maxImages */);
+        mImageReaders.add(imageReader);
+        final int flags = VIRTUAL_DISPLAY_FLAG_PRESENTATION | VIRTUAL_DISPLAY_FLAG_OWN_CONTENT_ONLY
+                | VIRTUAL_DISPLAY_FLAG_PUBLIC;
+        final VirtualDisplay virtualDisplay = mDisplayManager.createVirtualDisplay(name, width,
+                height, DISPLAY_DENSITY, imageReader.getSurface(), flags);
+        mVirtualDisplays.add(virtualDisplay);
+        assertNotNull("display must be registered", Arrays.stream(
+                mDisplayManager.getDisplays()).filter(d -> d.getName().equals(name)).findAny());
+        return virtualDisplay;
+    }
+}