diff --git a/services/tests/wmtests/src/com/android/server/wm/AppTransitionControllerTest.java b/services/tests/wmtests/src/com/android/server/wm/AppTransitionControllerTest.java
index f6213bd94dddb..6b613cad4b828 100644
--- a/services/tests/wmtests/src/com/android/server/wm/AppTransitionControllerTest.java
+++ b/services/tests/wmtests/src/com/android/server/wm/AppTransitionControllerTest.java
@@ -28,6 +28,7 @@ import static android.view.WindowManager.TRANSIT_TASK_OPEN;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertTrue;
+import static org.mockito.Mockito.doCallRealMethod;
 
 import android.platform.test.annotations.Presubmit;
 import android.util.ArraySet;
@@ -56,6 +57,14 @@ public class AppTransitionControllerTest extends WindowTestsBase {
         mAppTransitionController = new AppTransitionController(mWm, mDisplayContent);
     }
 
+    @Override
+    ActivityRecord createActivityRecord(DisplayContent dc, int windowingMode, int activityType) {
+        final ActivityRecord r = super.createActivityRecord(dc, windowingMode, activityType);
+        // Ensure that ActivityRecord#setOccludesParent takes effect.
+        doCallRealMethod().when(r).fillsParent();
+        return r;
+    }
+
     @Test
     @FlakyTest(bugId = 131005232)
     public void testTranslucentOpen() {
@@ -191,6 +200,9 @@ public class AppTransitionControllerTest extends WindowTestsBase {
 
     @Test
     public void testGetAnimationTargets_exitingBeforeTransition() {
+        // Create another non-empty task so the animation target won't promote to task display area.
+        WindowTestUtils.createTestActivityRecord(
+                mDisplayContent.getDefaultTaskDisplayArea().getOrCreateRootHomeTask());
         final ActivityStack stack = createTaskStackOnDisplay(mDisplayContent);
         final ActivityRecord activity = WindowTestUtils.createTestActivityRecord(stack);
         activity.setVisible(false);
diff --git a/services/tests/wmtests/src/com/android/server/wm/AppWindowTokenTests.java b/services/tests/wmtests/src/com/android/server/wm/AppWindowTokenTests.java
index 1d13788bc5230..c8b668bac455d 100644
--- a/services/tests/wmtests/src/com/android/server/wm/AppWindowTokenTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/AppWindowTokenTests.java
@@ -158,6 +158,7 @@ public class AppWindowTokenTests extends WindowTestsBase {
         mActivity.removeImmediately();
     }
 
+    @UseTestDisplay(addWindows = W_ACTIVITY)
     @Test
     @FlakyTest(bugId = 131005232)
     public void testLandscapeSeascapeRotationByApp() {
@@ -188,6 +189,7 @@ public class AppWindowTokenTests extends WindowTestsBase {
         appWindow.removeImmediately();
     }
 
+    @UseTestDisplay(addWindows = W_ACTIVITY)
     @Test
     public void testLandscapeSeascapeRotationByPolicy() {
         // This instance has been spied in {@link TestDisplayContent}.
@@ -295,6 +297,7 @@ public class AppWindowTokenTests extends WindowTestsBase {
         mWm.mDisplayFrozen = false;
     }
 
+    @UseTestDisplay
     @Test
     public void testRespectTopFullscreenOrientation() {
         final Configuration displayConfig = mActivity.mDisplayContent.getConfiguration();
@@ -316,6 +319,7 @@ public class AppWindowTokenTests extends WindowTestsBase {
         assertEquals(Configuration.ORIENTATION_LANDSCAPE, activityConfig.orientation);
     }
 
+    @UseTestDisplay
     @Test
     public void testReportOrientationChange() {
         mActivity.setOrientation(SCREEN_ORIENTATION_LANDSCAPE);
@@ -343,7 +347,7 @@ public class AppWindowTokenTests extends WindowTestsBase {
     @Test
     public void testAddRemoveRace() {
         // There was once a race condition between adding and removing starting windows
-        final ActivityRecord appToken = mAppWindow.mActivityRecord;
+        final ActivityRecord appToken = createIsolatedTestActivityRecord();
         for (int i = 0; i < 1000; i++) {
             appToken.addStartingWindow(mPackageName,
                     android.R.style.Theme, null, "Test", 0, 0, 0, 0, null, true, true, false, true,
diff --git a/services/tests/wmtests/src/com/android/server/wm/DisplayContentTests.java b/services/tests/wmtests/src/com/android/server/wm/DisplayContentTests.java
index cd524972111d8..792b597d23d21 100644
--- a/services/tests/wmtests/src/com/android/server/wm/DisplayContentTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/DisplayContentTests.java
@@ -133,6 +133,7 @@ import java.util.List;
 @RunWith(WindowTestRunner.class)
 public class DisplayContentTests extends WindowTestsBase {
 
+    @UseTestDisplay(addAllCommonWindows = true)
     @Test
     public void testForAllWindows() {
         final WindowState exitingAppWindow = createWindow(null, TYPE_BASE_APPLICATION,
@@ -159,6 +160,7 @@ public class DisplayContentTests extends WindowTestsBase {
                 mNavBarWindow));
     }
 
+    @UseTestDisplay(addAllCommonWindows = true)
     @Test
     public void testForAllWindows_WithAppImeTarget() {
         final WindowState imeAppTarget =
@@ -180,6 +182,7 @@ public class DisplayContentTests extends WindowTestsBase {
                 mNavBarWindow));
     }
 
+    @UseTestDisplay(addAllCommonWindows = true)
     @Test
     public void testForAllWindows_WithChildWindowImeTarget() throws Exception {
         mDisplayContent.mInputMethodTarget = mChildAppWindowAbove;
@@ -197,6 +200,7 @@ public class DisplayContentTests extends WindowTestsBase {
                 mNavBarWindow));
     }
 
+    @UseTestDisplay(addAllCommonWindows = true)
     @Test
     public void testForAllWindows_WithStatusBarImeTarget() throws Exception {
         mDisplayContent.mInputMethodTarget = mStatusBarWindow;
@@ -214,6 +218,7 @@ public class DisplayContentTests extends WindowTestsBase {
                 mNavBarWindow));
     }
 
+    @UseTestDisplay(addAllCommonWindows = true)
     @Test
     public void testForAllWindows_WithNotificationShadeImeTarget() throws Exception {
         mDisplayContent.mInputMethodTarget = mNotificationShadeWindow;
@@ -231,6 +236,7 @@ public class DisplayContentTests extends WindowTestsBase {
                 mNavBarWindow));
     }
 
+    @UseTestDisplay(addAllCommonWindows = true)
     @Test
     public void testForAllWindows_WithInBetweenWindowToken() {
         // This window is set-up to be z-ordered between some windows that go in the same token like
@@ -252,6 +258,7 @@ public class DisplayContentTests extends WindowTestsBase {
                 mNavBarWindow));
     }
 
+    @UseTestDisplay(addAllCommonWindows = true)
     @Test
     public void testComputeImeTarget() {
         // Verify that an app window can be an ime target.
@@ -271,6 +278,7 @@ public class DisplayContentTests extends WindowTestsBase {
         assertEquals(childWin, imeTarget);
     }
 
+    @UseTestDisplay(addAllCommonWindows = true)
     @Test
     public void testComputeImeTarget_startingWindow() {
         ActivityRecord activity = createActivityRecord(mDisplayContent,
@@ -775,6 +783,7 @@ public class DisplayContentTests extends WindowTestsBase {
                 .setDisplayInfoOverrideFromWindowManager(dc.getDisplayId(), null);
     }
 
+    @UseTestDisplay
     @Test
     public void testClearLastFocusWhenReparentingFocusedWindow() {
         final DisplayContent defaultDisplay = mWm.getDefaultDisplayContentLocked();
@@ -808,6 +817,7 @@ public class DisplayContentTests extends WindowTestsBase {
         assertFalse(isOptionsPanelAtRight(landscapeDisplay.getDisplayId()));
     }
 
+    @UseTestDisplay(addWindows = W_INPUT_METHOD)
     @Test
     public void testInputMethodTargetUpdateWhenSwitchingOnDisplays() {
         final DisplayContent newDisplay = createNewDisplay();
@@ -921,6 +931,7 @@ public class DisplayContentTests extends WindowTestsBase {
         assertEquals(dc.getImeContainer().getParentSurfaceControl(), dc.computeImeParent());
     }
 
+    @UseTestDisplay(addWindows = W_ACTIVITY)
     @Test
     public void testComputeImeParent_app_notMatchParentBounds() {
         spyOn(mAppWindow.mActivityRecord);
@@ -978,6 +989,7 @@ public class DisplayContentTests extends WindowTestsBase {
         assertNotEquals(dc.mInputMethodInputTarget, dc.computeImeControlTarget());
     }
 
+    @UseTestDisplay(addWindows = W_ACTIVITY)
     @Test
     public void testComputeImeControlTarget_notMatchParentBounds() throws Exception {
         spyOn(mAppWindow.mActivityRecord);
@@ -1096,6 +1108,7 @@ public class DisplayContentTests extends WindowTestsBase {
         win.setHasSurface(false);
     }
 
+    @UseTestDisplay(addWindows = { W_ABOVE_ACTIVITY, W_ACTIVITY})
     @Test
     public void testRequestResizeForEmptyFrames() {
         final WindowState win = mChildAppWindowAbove;
@@ -1135,6 +1148,7 @@ public class DisplayContentTests extends WindowTestsBase {
                 is(Configuration.ORIENTATION_PORTRAIT));
     }
 
+    @UseTestDisplay(addWindows = { W_ACTIVITY, W_WALLPAPER, W_STATUS_BAR, W_NAVIGATION_BAR })
     @Test
     public void testApplyTopFixedRotationTransform() {
         final DisplayPolicy displayPolicy = mDisplayContent.getDisplayPolicy();
@@ -1234,7 +1248,9 @@ public class DisplayContentTests extends WindowTestsBase {
 
     @Test
     public void testFinishFixedRotationNoAppTransitioningTask() {
-        final ActivityRecord app = mAppWindow.mActivityRecord;
+        unblockDisplayRotation(mDisplayContent);
+        final ActivityRecord app = createActivityRecord(mDisplayContent, WINDOWING_MODE_FULLSCREEN,
+                ACTIVITY_TYPE_STANDARD);
         final Task task = app.getTask();
         final ActivityRecord app2 = new ActivityTestsBase.ActivityBuilder(mWm.mAtmService)
                 .setTask(task).build();
@@ -1255,6 +1271,7 @@ public class DisplayContentTests extends WindowTestsBase {
         assertFalse(mDisplayContent.hasTopFixedRotationLaunchingApp());
     }
 
+    @UseTestDisplay(addWindows = W_ACTIVITY)
     @Test
     public void testRotateSeamlesslyWithFixedRotation() {
         final DisplayRotation displayRotation = mDisplayContent.getDisplayRotation();
@@ -1275,13 +1292,14 @@ public class DisplayContentTests extends WindowTestsBase {
 
     @Test
     public void testNoFixedRotationWithPip() {
+        final DisplayContent displayContent = mDefaultDisplay;
+        unblockDisplayRotation(displayContent);
         // Make resume-top really update the activity state.
         setBooted(mWm.mAtmService);
         // Speed up the test by a few seconds.
         mWm.mAtmService.deferWindowLayout();
         doNothing().when(mWm).startFreezingDisplay(anyInt(), anyInt(), any(), anyInt());
 
-        final DisplayContent displayContent = mWm.mRoot.getDefaultDisplay();
         final Configuration displayConfig = displayContent.getConfiguration();
         final ActivityRecord pinnedActivity = createActivityRecord(displayContent,
                 WINDOWING_MODE_PINNED, ACTIVITY_TYPE_STANDARD);
@@ -1324,11 +1342,13 @@ public class DisplayContentTests extends WindowTestsBase {
 
     @Test
     public void testRecentsNotRotatingWithFixedRotation() {
+        unblockDisplayRotation(mDisplayContent);
         final DisplayRotation displayRotation = mDisplayContent.getDisplayRotation();
-        doCallRealMethod().when(displayRotation).updateRotationUnchecked(anyBoolean());
         // Skip freezing so the unrelated conditions in updateRotationUnchecked won't disturb.
         doNothing().when(mWm).startFreezingDisplay(anyInt(), anyInt(), any(), anyInt());
 
+        final ActivityRecord activity = createActivityRecord(mDisplayContent,
+                WINDOWING_MODE_FULLSCREEN, ACTIVITY_TYPE_STANDARD);
         final ActivityRecord recentsActivity = createActivityRecord(mDisplayContent,
                 WINDOWING_MODE_FULLSCREEN, ACTIVITY_TYPE_RECENTS);
         recentsActivity.setRequestedOrientation(SCREEN_ORIENTATION_PORTRAIT);
@@ -1345,12 +1365,12 @@ public class DisplayContentTests extends WindowTestsBase {
         // Rotation can be updated if the recents animation is animating but it is not on top, e.g.
         // switching activities in different orientations by quickstep gesture.
         mDisplayContent.mFixedRotationTransitionListener.onStartRecentsAnimation(recentsActivity);
-        mDisplayContent.setFixedRotationLaunchingAppUnchecked(mAppWindow.mActivityRecord);
+        mDisplayContent.setFixedRotationLaunchingAppUnchecked(activity);
         displayRotation.setRotation((displayRotation.getRotation() + 1) % 4);
         assertTrue(displayRotation.updateRotationUnchecked(false));
 
         // The recents activity should not apply fixed rotation if the top activity is not opaque.
-        mDisplayContent.mFocusedApp = mAppWindow.mActivityRecord;
+        mDisplayContent.mFocusedApp = activity;
         doReturn(false).when(mDisplayContent.mFocusedApp).occludesParent();
         doReturn(ROTATION_90).when(mDisplayContent).rotationForActivityInDifferentOrientation(
                 eq(recentsActivity));
diff --git a/services/tests/wmtests/src/com/android/server/wm/DisplayPolicyInsetsTests.java b/services/tests/wmtests/src/com/android/server/wm/DisplayPolicyInsetsTests.java
index 39cd76aeef9ed..402fd22416fcc 100644
--- a/services/tests/wmtests/src/com/android/server/wm/DisplayPolicyInsetsTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/DisplayPolicyInsetsTests.java
@@ -36,6 +36,8 @@ import org.junit.runner.RunWith;
 
 @SmallTest
 @Presubmit
+@WindowTestsBase.UseTestDisplay(
+        addWindows = { WindowTestsBase.W_STATUS_BAR, WindowTestsBase.W_NAVIGATION_BAR })
 @RunWith(WindowTestRunner.class)
 public class DisplayPolicyInsetsTests extends DisplayPolicyTestsBase {
 
diff --git a/services/tests/wmtests/src/com/android/server/wm/DisplayPolicyLayoutTests.java b/services/tests/wmtests/src/com/android/server/wm/DisplayPolicyLayoutTests.java
index da7c41a745286..f9de379943503 100644
--- a/services/tests/wmtests/src/com/android/server/wm/DisplayPolicyLayoutTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/DisplayPolicyLayoutTests.java
@@ -49,7 +49,6 @@ import static org.hamcrest.Matchers.is;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNotEquals;
 import static org.junit.Assert.assertThat;
-import static org.mockito.Mockito.doCallRealMethod;
 import static org.mockito.Mockito.doNothing;
 import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.spy;
@@ -65,14 +64,12 @@ import android.util.SparseArray;
 import android.view.DisplayCutout;
 import android.view.DisplayInfo;
 import android.view.InsetsState;
-import android.view.View;
 import android.view.WindowInsets.Side;
 import android.view.WindowInsets.Type;
 import android.view.WindowManager;
 
 import androidx.test.filters.SmallTest;
 
-import com.android.server.policy.WindowManagerPolicy;
 import com.android.server.wm.utils.WmDisplayCutout;
 
 import org.junit.Before;
@@ -90,6 +87,8 @@ import java.io.StringWriter;
  */
 @SmallTest
 @Presubmit
+@WindowTestsBase.UseTestDisplay(
+        addWindows = { WindowTestsBase.W_STATUS_BAR, WindowTestsBase.W_NAVIGATION_BAR })
 @RunWith(WindowTestRunner.class)
 public class DisplayPolicyLayoutTests extends DisplayPolicyTestsBase {
 
@@ -120,9 +119,6 @@ public class DisplayPolicyLayoutTests extends DisplayPolicyTestsBase {
         // Disabling this call for most tests since it can override the systemUiFlags when called.
         doReturn(0).when(mDisplayPolicy).updateSystemUiVisibilityLw();
 
-        mDisplayPolicy.mLastSystemUiFlags |= View.STATUS_BAR_TRANSPARENT;
-        mDisplayPolicy.mLastSystemUiFlags |= View.NAVIGATION_BAR_TRANSPARENT;
-
         updateDisplayFrames();
     }
 
@@ -146,10 +142,7 @@ public class DisplayPolicyLayoutTests extends DisplayPolicyTestsBase {
         mFrames = createDisplayFrames();
         mDisplayBounds.set(0, 0, mFrames.mDisplayWidth, mFrames.mDisplayHeight);
         mDisplayContent.mDisplayFrames = mFrames;
-
-        doReturn(mDisplayBounds).when(mStatusBarWindow).getBounds();
-        doReturn(mDisplayBounds).when(mNavBarWindow).getBounds();
-        doReturn(mDisplayBounds).when(mWindow).getBounds();
+        mDisplayContent.setBounds(mDisplayBounds);
     }
 
     private DisplayFrames createDisplayFrames() {
diff --git a/services/tests/wmtests/src/com/android/server/wm/DisplayPolicyTests.java b/services/tests/wmtests/src/com/android/server/wm/DisplayPolicyTests.java
index a3f9b2e71cb59..4483f8c341cf8 100644
--- a/services/tests/wmtests/src/com/android/server/wm/DisplayPolicyTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/DisplayPolicyTests.java
@@ -222,6 +222,7 @@ public class DisplayPolicyTests extends WindowTestsBase {
                         opaqueDarkNavBar, imeDrawLightNavBar, imeDrawLightNavBar));
     }
 
+    @UseTestDisplay(addWindows = W_ACTIVITY)
     @Test
     public void testComputeTopFullscreenOpaqueWindow() {
         final WindowManager.LayoutParams attrs = mAppWindow.mAttrs;
@@ -320,6 +321,8 @@ public class DisplayPolicyTests extends WindowTestsBase {
         return win;
     }
 
+    @UseTestDisplay(
+            addWindows = { W_ACTIVITY, W_STATUS_BAR, W_NAVIGATION_BAR, W_NOTIFICATION_SHADE })
     @Test
     public void testUpdateHideNavInputEventReceiver() {
         final InsetsPolicy insetsPolicy = mDisplayContent.getInsetsPolicy();
@@ -358,6 +361,7 @@ public class DisplayPolicyTests extends WindowTestsBase {
         assertNull(displayPolicy.mInputConsumer);
     }
 
+    @UseTestDisplay(addWindows = { W_NAVIGATION_BAR, W_INPUT_METHOD })
     @Test
     public void testImeMinimalSourceFrame() {
         final DisplayPolicy displayPolicy = mDisplayContent.getDisplayPolicy();
diff --git a/services/tests/wmtests/src/com/android/server/wm/DisplayPolicyTestsBase.java b/services/tests/wmtests/src/com/android/server/wm/DisplayPolicyTestsBase.java
index 1e1c3998670e4..b4e1c375993dc 100644
--- a/services/tests/wmtests/src/com/android/server/wm/DisplayPolicyTestsBase.java
+++ b/services/tests/wmtests/src/com/android/server/wm/DisplayPolicyTestsBase.java
@@ -68,6 +68,10 @@ public class DisplayPolicyTestsBase extends WindowTestsBase {
 
     @Before
     public void setUpDisplayPolicy() {
+        // Disable surface placement because it has no direct relation to layout policy and it also
+        // avoids some noises such as the display info is modified, screen frozen, config change.
+        mWm.mWindowPlacerLocked.deferLayout();
+
         mDisplayPolicy = mDisplayContent.getDisplayPolicy();
         spyOn(mDisplayPolicy);
 
@@ -100,6 +104,9 @@ public class DisplayPolicyTestsBase extends WindowTestsBase {
         mNavBarWindow.mAttrs.gravity = Gravity.BOTTOM;
         addWindow(mNavBarWindow);
         mDisplayPolicy.mLastSystemUiFlags |= View.NAVIGATION_BAR_TRANSPARENT;
+
+        // Update source frame and visibility of insets providers.
+        mDisplayContent.getInsetsStateController().onPostLayout();
     }
 
     void addWindow(WindowState win) {
diff --git a/services/tests/wmtests/src/com/android/server/wm/DisplayWindowSettingsTests.java b/services/tests/wmtests/src/com/android/server/wm/DisplayWindowSettingsTests.java
index 11c02c26ca97e..a3d3739a9015c 100644
--- a/services/tests/wmtests/src/com/android/server/wm/DisplayWindowSettingsTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/DisplayWindowSettingsTests.java
@@ -77,6 +77,7 @@ import java.nio.charset.StandardCharsets;
  */
 @SmallTest
 @Presubmit
+@WindowTestsBase.UseTestDisplay
 @RunWith(WindowTestRunner.class)
 public class DisplayWindowSettingsTests extends WindowTestsBase {
 
diff --git a/services/tests/wmtests/src/com/android/server/wm/InsetsPolicyTest.java b/services/tests/wmtests/src/com/android/server/wm/InsetsPolicyTest.java
index 87bc7f1bf7817..555906d4c9104 100644
--- a/services/tests/wmtests/src/com/android/server/wm/InsetsPolicyTest.java
+++ b/services/tests/wmtests/src/com/android/server/wm/InsetsPolicyTest.java
@@ -225,6 +225,7 @@ public class InsetsPolicyTest extends WindowTestsBase {
         assertEquals(ITYPE_STATUS_BAR, fullscreenAppControls[0].getType());
     }
 
+    @UseTestDisplay(addWindows = W_ACTIVITY)
     @Test
     public void testShowTransientBars_bothCanBeTransient_appGetsBothFakeControls() {
         addNonFocusableWindow(TYPE_STATUS_BAR, "statusBar")
@@ -259,6 +260,7 @@ public class InsetsPolicyTest extends WindowTestsBase {
                 .getSource(ITYPE_NAVIGATION_BAR).isVisible());
     }
 
+    @UseTestDisplay(addWindows = W_ACTIVITY)
     @Test
     public void testShowTransientBars_statusBarCanBeTransient_appGetsStatusBarFakeControl() {
         addNonFocusableWindow(TYPE_STATUS_BAR, "statusBar")
@@ -288,6 +290,7 @@ public class InsetsPolicyTest extends WindowTestsBase {
         }
     }
 
+    @UseTestDisplay(addWindows = W_ACTIVITY)
     @Test
     public void testAbortTransientBars_bothCanBeAborted_appGetsBothRealControls() {
         addNonFocusableWindow(TYPE_STATUS_BAR, "statusBar")
diff --git a/services/tests/wmtests/src/com/android/server/wm/InsetsStateControllerTest.java b/services/tests/wmtests/src/com/android/server/wm/InsetsStateControllerTest.java
index 0a27e1a1da68c..5e83e66536ed6 100644
--- a/services/tests/wmtests/src/com/android/server/wm/InsetsStateControllerTest.java
+++ b/services/tests/wmtests/src/com/android/server/wm/InsetsStateControllerTest.java
@@ -147,6 +147,7 @@ public class InsetsStateControllerTest extends WindowTestsBase {
         assertNull(getController().getInsetsForDispatch(app).peekSource(ITYPE_NAVIGATION_BAR));
     }
 
+    @UseTestDisplay(addWindows = W_INPUT_METHOD)
     @Test
     public void testStripForDispatch_independentSources() {
         getController().getSourceProvider(ITYPE_IME).setWindow(mImeWindow, null, null);
@@ -162,6 +163,7 @@ public class InsetsStateControllerTest extends WindowTestsBase {
         assertTrue(getController().getInsetsForDispatch(app1).getSource(ITYPE_IME).isVisible());
     }
 
+    @UseTestDisplay(addWindows = W_INPUT_METHOD)
     @Test
     public void testStripForDispatch_belowIme() {
         getController().getSourceProvider(ITYPE_IME).setWindow(mImeWindow, null, null);
@@ -173,6 +175,7 @@ public class InsetsStateControllerTest extends WindowTestsBase {
         assertTrue(getController().getInsetsForDispatch(app).getSource(ITYPE_IME).isVisible());
     }
 
+    @UseTestDisplay(addWindows = W_INPUT_METHOD)
     @Test
     public void testStripForDispatch_aboveIme() {
         getController().getSourceProvider(ITYPE_IME).setWindow(mImeWindow, null, null);
@@ -184,6 +187,7 @@ public class InsetsStateControllerTest extends WindowTestsBase {
         assertFalse(getController().getInsetsForDispatch(app).getSource(ITYPE_IME).isVisible());
     }
 
+    @UseTestDisplay(addWindows = W_INPUT_METHOD)
     @Test
     public void testStripForDispatch_imeOrderChanged() {
         // This can be the IME z-order target while app cannot be the IME z-order target.
@@ -232,6 +236,7 @@ public class InsetsStateControllerTest extends WindowTestsBase {
         assertTrue(getController().getInsetsForDispatch(app).getSource(ITYPE_IME).isVisible());
     }
 
+    @UseTestDisplay(addWindows = W_INPUT_METHOD)
     @Test
     public void testStripForDispatch_childWindow_altFocusable() {
         getController().getSourceProvider(ITYPE_IME).setWindow(mImeWindow, null, null);
@@ -249,6 +254,7 @@ public class InsetsStateControllerTest extends WindowTestsBase {
         assertFalse(getController().getInsetsForDispatch(child).getSource(ITYPE_IME).isVisible());
     }
 
+    @UseTestDisplay(addWindows = W_INPUT_METHOD)
     @Test
     public void testStripForDispatch_childWindow_splitScreen() {
         getController().getSourceProvider(ITYPE_IME).setWindow(mImeWindow, null, null);
diff --git a/services/tests/wmtests/src/com/android/server/wm/RecentsAnimationControllerTest.java b/services/tests/wmtests/src/com/android/server/wm/RecentsAnimationControllerTest.java
index 8e85e7b96d1f0..4fbdd616dc93c 100644
--- a/services/tests/wmtests/src/com/android/server/wm/RecentsAnimationControllerTest.java
+++ b/services/tests/wmtests/src/com/android/server/wm/RecentsAnimationControllerTest.java
@@ -92,7 +92,6 @@ public class RecentsAnimationControllerTest extends WindowTestsBase {
     @Mock RecentsAnimationController.RecentsAnimationCallbacks mAnimationCallbacks;
     @Mock TaskSnapshot mMockTaskSnapshot;
     private RecentsAnimationController mController;
-    private DisplayContent mDefaultDisplay;
     private ActivityStack mRootHomeTask;
 
     @Before
@@ -100,7 +99,6 @@ public class RecentsAnimationControllerTest extends WindowTestsBase {
         MockitoAnnotations.initMocks(this);
         doNothing().when(mWm.mRoot).performSurfacePlacement();
         when(mMockRunner.asBinder()).thenReturn(new Binder());
-        mDefaultDisplay = mWm.mRoot.getDefaultDisplay();
         mController = spy(new RecentsAnimationController(mWm, mMockRunner, mAnimationCallbacks,
                 DEFAULT_DISPLAY));
         mRootHomeTask = mDefaultDisplay.getDefaultTaskDisplayArea().getRootHomeTask();
@@ -321,6 +319,7 @@ public class RecentsAnimationControllerTest extends WindowTestsBase {
 
     @Test
     public void testRecentViewInFixedPortraitWhenTopAppInLandscape() {
+        unblockDisplayRotation(mDefaultDisplay);
         mWm.setRecentsAnimationController(mController);
 
         final ActivityRecord homeActivity = createHomeActivity();
@@ -365,6 +364,7 @@ public class RecentsAnimationControllerTest extends WindowTestsBase {
 
     @Test
     public void testClearFixedRotationLaunchingAppAfterCleanupAnimation() {
+        unblockDisplayRotation(mDefaultDisplay);
         final ActivityRecord homeActivity = createHomeActivity();
         homeActivity.setRequestedOrientation(ActivityInfo.SCREEN_ORIENTATION_PORTRAIT);
         final ActivityRecord activity = createActivityRecord(mDefaultDisplay,
@@ -389,6 +389,7 @@ public class RecentsAnimationControllerTest extends WindowTestsBase {
 
     @Test
     public void testWallpaperHasFixedRotationApplied() {
+        unblockDisplayRotation(mDefaultDisplay);
         mWm.setRecentsAnimationController(mController);
 
         // Create a portrait home activity, a wallpaper and a landscape activity displayed on top.
diff --git a/services/tests/wmtests/src/com/android/server/wm/SurfaceAnimationRunnerTest.java b/services/tests/wmtests/src/com/android/server/wm/SurfaceAnimationRunnerTest.java
index 7ffb2eb26721d..2233b2234c688 100644
--- a/services/tests/wmtests/src/com/android/server/wm/SurfaceAnimationRunnerTest.java
+++ b/services/tests/wmtests/src/com/android/server/wm/SurfaceAnimationRunnerTest.java
@@ -34,6 +34,7 @@ import android.animation.AnimationHandler.AnimationFrameCallbackProvider;
 import android.animation.ValueAnimator;
 import android.graphics.Matrix;
 import android.graphics.Point;
+import android.os.Handler;
 import android.os.PowerManagerInternal;
 import android.platform.test.annotations.Presubmit;
 import android.view.Choreographer;
@@ -46,11 +47,12 @@ import android.view.animation.TranslateAnimation;
 import androidx.test.filters.FlakyTest;
 import androidx.test.filters.SmallTest;
 
+import com.android.server.AnimationThread;
 import com.android.server.wm.LocalAnimationAdapter.AnimationSpec;
 
+import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
-import org.junit.runner.RunWith;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 
@@ -64,8 +66,7 @@ import java.util.concurrent.CountDownLatch;
  */
 @SmallTest
 @Presubmit
-@RunWith(WindowTestRunner.class)
-public class SurfaceAnimationRunnerTest extends WindowTestsBase {
+public class SurfaceAnimationRunnerTest {
 
     @Mock SurfaceControl mMockSurface;
     @Mock Transaction mMockTransaction;
@@ -75,6 +76,9 @@ public class SurfaceAnimationRunnerTest extends WindowTestsBase {
     private SurfaceAnimationRunner mSurfaceAnimationRunner;
     private CountDownLatch mFinishCallbackLatch;
 
+    private final Handler mAnimationThreadHandler = AnimationThread.getHandler();
+    private final Handler mSurfaceAnimationHandler = SurfaceAnimationThread.getHandler();
+
     @Before
     public void setUp() throws Exception {
         MockitoAnnotations.initMocks(this);
@@ -84,6 +88,12 @@ public class SurfaceAnimationRunnerTest extends WindowTestsBase {
                 mMockTransaction, mMockPowerManager);
     }
 
+    @After
+    public void tearDown() {
+        SurfaceAnimationThread.dispose();
+        AnimationThread.dispose();
+    }
+
     private void finishedCallback() {
         mFinishCallbackLatch.countDown();
     }
@@ -101,8 +111,7 @@ public class SurfaceAnimationRunnerTest extends WindowTestsBase {
         verify(mMockTransaction, atLeastOnce()).setMatrix(eq(mMockSurface), eq(m), any());
         verify(mMockTransaction, atLeastOnce()).setAlpha(eq(mMockSurface), eq(1.0f));
 
-        waitHandlerIdle(SurfaceAnimationThread.getHandler());
-        mFinishCallbackLatch.await(1, SECONDS);
+        waitHandlerIdle(mSurfaceAnimationHandler);
         assertFinishCallbackCalled();
 
         m.setTranslate(10, 0);
@@ -120,7 +129,7 @@ public class SurfaceAnimationRunnerTest extends WindowTestsBase {
                 .startAnimation(createTranslateAnimation(), mMockSurface, mMockTransaction,
                 this::finishedCallback);
         mSurfaceAnimationRunner.onAnimationCancelled(mMockSurface);
-        waitUntilHandlersIdle();
+        waitHandlerIdle(mAnimationThreadHandler);
         assertTrue(mSurfaceAnimationRunner.mPendingAnimations.isEmpty());
         assertFinishCallbackNotCalled();
     }
@@ -135,7 +144,7 @@ public class SurfaceAnimationRunnerTest extends WindowTestsBase {
         assertFalse(mSurfaceAnimationRunner.mRunningAnimations.isEmpty());
         mSurfaceAnimationRunner.onAnimationCancelled(mMockSurface);
         assertTrue(mSurfaceAnimationRunner.mRunningAnimations.isEmpty());
-        waitUntilHandlersIdle();
+        waitHandlerIdle(mAnimationThreadHandler);
         assertFinishCallbackNotCalled();
     }
 
@@ -180,9 +189,8 @@ public class SurfaceAnimationRunnerTest extends WindowTestsBase {
         assertTrue(mSurfaceAnimationRunner.mRunningAnimations.isEmpty());
         mSurfaceAnimationRunner.continueStartingAnimations();
         waitUntilNextFrame();
-        waitHandlerIdle(SurfaceAnimationThread.getHandler());
+        waitHandlerIdle(mSurfaceAnimationHandler);
         assertFalse(mSurfaceAnimationRunner.mRunningAnimations.isEmpty());
-        mFinishCallbackLatch.await(1, SECONDS);
         assertFinishCallbackCalled();
     }
 
@@ -204,7 +212,15 @@ public class SurfaceAnimationRunnerTest extends WindowTestsBase {
         latch.await();
     }
 
+    private static void waitHandlerIdle(Handler handler) {
+        handler.runWithScissors(() -> { }, 0 /* timeout */);
+    }
+
     private void assertFinishCallbackCalled() {
+        try {
+            assertTrue(mFinishCallbackLatch.await(5, SECONDS));
+        } catch (InterruptedException ignored) {
+        }
         assertEquals(0, mFinishCallbackLatch.getCount());
     }
 
diff --git a/services/tests/wmtests/src/com/android/server/wm/SystemServicesTestRule.java b/services/tests/wmtests/src/com/android/server/wm/SystemServicesTestRule.java
index 218261b010295..50675b03ae1fd 100644
--- a/services/tests/wmtests/src/com/android/server/wm/SystemServicesTestRule.java
+++ b/services/tests/wmtests/src/com/android/server/wm/SystemServicesTestRule.java
@@ -102,6 +102,7 @@ public class SystemServicesTestRule implements TestRule {
 
     private static final int[] TEST_USER_PROFILE_IDS = {};
 
+    private Description mDescription;
     private Context mContext;
     private StaticMockitoSession mMockitoSession;
     private ActivityManagerService mAmService;
@@ -121,6 +122,7 @@ public class SystemServicesTestRule implements TestRule {
         return new Statement() {
             @Override
             public void evaluate() throws Throwable {
+                mDescription = description;
                 Throwable throwable = null;
                 try {
                     runWithDexmakerShareClassLoader(SystemServicesTestRule.this::setUp);
@@ -373,6 +375,10 @@ public class SystemServicesTestRule implements TestRule {
         LocalServices.removeServiceForTest(StatusBarManagerInternal.class);
     }
 
+    Description getDescription() {
+        return mDescription;
+    }
+
     WindowManagerService getWindowManagerService() {
         return mWmService;
     }
diff --git a/services/tests/wmtests/src/com/android/server/wm/TaskSnapshotControllerTest.java b/services/tests/wmtests/src/com/android/server/wm/TaskSnapshotControllerTest.java
index 0c4bb9f2ac2bb..0b99e328b9bf7 100644
--- a/services/tests/wmtests/src/com/android/server/wm/TaskSnapshotControllerTest.java
+++ b/services/tests/wmtests/src/com/android/server/wm/TaskSnapshotControllerTest.java
@@ -183,6 +183,7 @@ public class TaskSnapshotControllerTest extends WindowTestsBase {
         }
     }
 
+    @UseTestDisplay(addWindows = W_ACTIVITY)
     @Test
     public void testPrepareTaskSnapshot() {
         mAppWindow.mWinAnimator.mLastAlpha = 1f;
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 b8a1c2b414544..bce1142c99bef 100644
--- a/services/tests/wmtests/src/com/android/server/wm/TaskStackChangedListenerTest.java
+++ b/services/tests/wmtests/src/com/android/server/wm/TaskStackChangedListenerTest.java
@@ -74,6 +74,7 @@ public class TaskStackChangedListenerTest {
     private IActivityManager mService;
     private ITaskStackListener mTaskStackListener;
 
+    private static final int WAIT_TIMEOUT_MS = 5000;
     private static final Object sLock = new Object();
     @GuardedBy("sLock")
     private static boolean sTaskStackChangedCalled;
@@ -490,7 +491,8 @@ public class TaskStackChangedListenerTest {
         SystemUtil.runWithShellPermissionIdentity(() -> context.startActivity(
                 new Intent(context, activityClass).addFlags(Intent.FLAG_ACTIVITY_NEW_TASK),
                 options.toBundle()));
-        final TestActivity activity = (TestActivity) monitor.waitForActivityWithTimeout(1000);
+        final TestActivity activity =
+                (TestActivity) monitor.waitForActivityWithTimeout(WAIT_TIMEOUT_MS);
         if (activity == null) {
             throw new RuntimeException("Timed out waiting for Activity");
         }
@@ -508,7 +510,7 @@ public class TaskStackChangedListenerTest {
 
     private void waitForCallback(CountDownLatch latch) {
         try {
-            final boolean result = latch.await(4, TimeUnit.SECONDS);
+            final boolean result = latch.await(WAIT_TIMEOUT_MS, TimeUnit.MILLISECONDS);
             if (!result) {
                 throw new RuntimeException("Timed out waiting for task stack change notification");
             }
@@ -560,7 +562,7 @@ public class TaskStackChangedListenerTest {
                 if (mIsResumed == isResumed) {
                     return;
                 }
-                wait(5000);
+                wait(WAIT_TIMEOUT_MS);
             }
             assertEquals("The activity resume state change timed out", isResumed, mIsResumed);
         }
diff --git a/services/tests/wmtests/src/com/android/server/wm/WallpaperControllerTests.java b/services/tests/wmtests/src/com/android/server/wm/WallpaperControllerTests.java
index 53ede60e9ac76..573e37a2d6b33 100644
--- a/services/tests/wmtests/src/com/android/server/wm/WallpaperControllerTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/WallpaperControllerTests.java
@@ -270,6 +270,7 @@ public class WallpaperControllerTests extends WindowTestsBase {
         assertEquals(WINDOWING_MODE_FULLSCREEN, token.getWindowingMode());
     }
 
+    @UseTestDisplay(addWindows = W_ACTIVITY)
     @Test
     public void testFixedRotationRecentsAnimatingTask() {
         final RecentsAnimationController recentsController = mock(RecentsAnimationController.class);
diff --git a/services/tests/wmtests/src/com/android/server/wm/WindowContainerTraversalTests.java b/services/tests/wmtests/src/com/android/server/wm/WindowContainerTraversalTests.java
index 3c0dd1e897f51..47e4559cbb130 100644
--- a/services/tests/wmtests/src/com/android/server/wm/WindowContainerTraversalTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/WindowContainerTraversalTests.java
@@ -46,6 +46,7 @@ import java.util.function.Consumer;
 @RunWith(WindowTestRunner.class)
 public class WindowContainerTraversalTests extends WindowTestsBase {
 
+    @UseTestDisplay(addWindows = { W_DOCK_DIVIDER, W_INPUT_METHOD })
     @Test
     public void testDockedDividerPosition() {
         final WindowState splitScreenWindow = createWindowOnStack(null,
diff --git a/services/tests/wmtests/src/com/android/server/wm/WindowOrganizerTests.java b/services/tests/wmtests/src/com/android/server/wm/WindowOrganizerTests.java
index 079ed973ca689..7cc19ad0ddd88 100644
--- a/services/tests/wmtests/src/com/android/server/wm/WindowOrganizerTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/WindowOrganizerTests.java
@@ -521,6 +521,7 @@ public class WindowOrganizerTests extends WindowTestsBase {
         assertEquals(ACTIVITY_TYPE_UNDEFINED, info1.topActivityType);
     }
 
+    @UseTestDisplay
     @Test
     public void testTaskInfoCallback() {
         final ArrayList<RunningTaskInfo> lastReportedTiles = new ArrayList<>();
@@ -584,6 +585,7 @@ public class WindowOrganizerTests extends WindowTestsBase {
         assertEquals(ACTIVITY_TYPE_UNDEFINED, lastReportedTiles.get(0).topActivityType);
     }
 
+    @UseTestDisplay
     @Test
     public void testHierarchyTransaction() {
         final ArrayMap<IBinder, RunningTaskInfo> lastReportedTiles = new ArrayMap<>();
diff --git a/services/tests/wmtests/src/com/android/server/wm/WindowStateTests.java b/services/tests/wmtests/src/com/android/server/wm/WindowStateTests.java
index 360d73b5bd872..ce9dd685293a7 100644
--- a/services/tests/wmtests/src/com/android/server/wm/WindowStateTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/WindowStateTests.java
@@ -36,7 +36,6 @@ import static android.view.WindowManager.LayoutParams.TYPE_APPLICATION_STARTING;
 import static android.view.WindowManager.LayoutParams.TYPE_APPLICATION_SUB_PANEL;
 import static android.view.WindowManager.LayoutParams.TYPE_BASE_APPLICATION;
 import static android.view.WindowManager.LayoutParams.TYPE_INPUT_METHOD;
-import static android.view.WindowManager.LayoutParams.TYPE_STATUS_BAR;
 
 import static com.android.dx.mockito.inline.extended.ExtendedMockito.doNothing;
 import static com.android.dx.mockito.inline.extended.ExtendedMockito.doReturn;
@@ -417,10 +416,11 @@ public class WindowStateTests extends WindowTestsBase {
         assertFalse(app.canAffectSystemUiFlags());
     }
 
+    @UseTestDisplay(addWindows = { W_ACTIVITY, W_STATUS_BAR })
     @Test
     public void testVisibleWithInsetsProvider() {
-        final WindowState statusBar = createWindow(null, TYPE_STATUS_BAR, "statusBar");
-        final WindowState app = createWindow(null, TYPE_APPLICATION, "app");
+        final WindowState statusBar = mStatusBarWindow;
+        final WindowState app = mAppWindow;
         statusBar.mHasSurface = true;
         assertTrue(statusBar.isVisible());
         mDisplayContent.getInsetsStateController().getSourceProvider(ITYPE_STATUS_BAR)
@@ -542,6 +542,7 @@ public class WindowStateTests extends WindowTestsBase {
         assertTrue(window.isVisibleByPolicy());
     }
 
+    @UseTestDisplay(addWindows = { W_ABOVE_ACTIVITY, W_NOTIFICATION_SHADE })
     @Test
     public void testRequestDrawIfNeeded() {
         final WindowState startingApp = createWindow(null /* parent */,
@@ -567,6 +568,7 @@ public class WindowStateTests extends WindowTestsBase {
         assertEquals(Arrays.asList(keyguardHostWindow, startingWindow), outWaitingForDrawn);
     }
 
+    @UseTestDisplay(addWindows = W_ABOVE_ACTIVITY)
     @Test
     public void testReportResizedWithRemoteException() {
         final WindowState win = mChildAppWindowAbove;
@@ -597,6 +599,7 @@ public class WindowStateTests extends WindowTestsBase {
         assertFalse(win.getOrientationChanging());
     }
 
+    @UseTestDisplay(addWindows = W_ABOVE_ACTIVITY)
     @Test
     public void testRequestResizeForBlastSync() {
         final WindowState win = mChildAppWindowAbove;
@@ -677,6 +680,7 @@ public class WindowStateTests extends WindowTestsBase {
         assertTrue(win0.cantReceiveTouchInput());
     }
 
+    @UseTestDisplay(addWindows = W_ACTIVITY)
     @Test
     public void testNeedsRelativeLayeringToIme_notAttached() {
         WindowState sameTokenWindow = createWindow(null, TYPE_BASE_APPLICATION, mAppWindow.mToken,
@@ -689,6 +693,7 @@ public class WindowStateTests extends WindowTestsBase {
         assertFalse(sameTokenWindow.needsRelativeLayeringToIme());
     }
 
+    @UseTestDisplay(addWindows = { W_ACTIVITY, W_INPUT_METHOD })
     @Test
     public void testNeedsRelativeLayeringToIme_startingWindow() {
         WindowState sameTokenWindow = createWindow(null, TYPE_APPLICATION_STARTING,
diff --git a/services/tests/wmtests/src/com/android/server/wm/WindowTestsBase.java b/services/tests/wmtests/src/com/android/server/wm/WindowTestsBase.java
index 4c18c86829f1e..0bbe0a04f6189 100644
--- a/services/tests/wmtests/src/com/android/server/wm/WindowTestsBase.java
+++ b/services/tests/wmtests/src/com/android/server/wm/WindowTestsBase.java
@@ -39,11 +39,11 @@ import static androidx.test.platform.app.InstrumentationRegistry.getInstrumentat
 
 import static org.mockito.Mockito.mock;
 
-import android.content.Context;
+import android.annotation.IntDef;
 import android.content.Intent;
+import android.hardware.display.DisplayManager;
 import android.os.RemoteException;
 import android.os.UserHandle;
-import android.util.Log;
 import android.view.Display;
 import android.view.DisplayInfo;
 import android.view.IDisplayWindowInsetsController;
@@ -54,23 +54,36 @@ import android.view.SurfaceControl.Transaction;
 import android.view.View;
 import android.view.WindowManager;
 
+import com.android.internal.util.ArrayUtils;
 import com.android.server.AttributeCache;
 
 import org.junit.Before;
 import org.junit.BeforeClass;
+import org.junit.runner.Description;
+
+import java.lang.annotation.ElementType;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+import java.lang.annotation.Target;
 
 /** Common base class for window manager unit test classes. */
 class WindowTestsBase extends SystemServiceTestsBase {
-    private static final String TAG = WindowTestsBase.class.getSimpleName();
 
     WindowManagerService mWm;
     private final IWindow mIWindow = new TestIWindow();
     private Session mMockSession;
     static int sNextStackId = 1000;
 
-    /** Non-default display. */
-    DisplayContent mDisplayContent;
     DisplayInfo mDisplayInfo = new DisplayInfo();
+    DisplayContent mDefaultDisplay;
+
+    /**
+     * It is {@link #mDefaultDisplay} by default. If the test class or method is annotated with
+     * {@link UseTestDisplay}, it will be an additional display.
+     */
+    DisplayContent mDisplayContent;
+
+    // The following fields are only available depending on the usage of annotation UseTestDisplay.
     WindowState mWallpaperWindow;
     WindowState mImeWindow;
     WindowState mImeDialogWindow;
@@ -97,48 +110,82 @@ class WindowTestsBase extends SystemServiceTestsBase {
         mWm = mSystemServicesTestRule.getWindowManagerService();
         SystemServicesTestRule.checkHoldsLock(mWm.mGlobalLock);
 
+        mDefaultDisplay = mWm.mRoot.getDefaultDisplay();
         mTransaction = mSystemServicesTestRule.mTransaction;
         mMockSession = mock(Session.class);
-        final Context context = getInstrumentation().getTargetContext();
-        // If @Before throws an exception, the error isn't logged. This will make sure any failures
-        // in the set up are clear. This can be removed when b/37850063 is fixed.
-        try {
-            beforeCreateDisplay();
 
-            context.getDisplay().getDisplayInfo(mDisplayInfo);
-            mDisplayContent = createNewDisplay(true /* supportIme */);
+        getInstrumentation().getTargetContext().getSystemService(DisplayManager.class)
+                .getDisplay(Display.DEFAULT_DISPLAY).getDisplayInfo(mDisplayInfo);
 
-            // Set-up some common windows.
-            mWallpaperWindow = createCommonWindow(null, TYPE_WALLPAPER, "wallpaperWindow");
-            mImeWindow = createCommonWindow(null, TYPE_INPUT_METHOD, "mImeWindow");
-            mDisplayContent.mInputMethodWindow = mImeWindow;
-            mImeDialogWindow = createCommonWindow(null, TYPE_INPUT_METHOD_DIALOG,
-                    "mImeDialogWindow");
-            mStatusBarWindow = createCommonWindow(null, TYPE_STATUS_BAR, "mStatusBarWindow");
-            mNotificationShadeWindow = createCommonWindow(null, TYPE_NOTIFICATION_SHADE,
-                    "mNotificationShadeWindow");
-            mNavBarWindow = createCommonWindow(null, TYPE_NAVIGATION_BAR, "mNavBarWindow");
-            mDockedDividerWindow = createCommonWindow(null, TYPE_DOCK_DIVIDER,
-                    "mDockedDividerWindow");
-            mAppWindow = createCommonWindow(null, TYPE_BASE_APPLICATION, "mAppWindow");
-            mChildAppWindowAbove = createCommonWindow(mAppWindow,
-                    TYPE_APPLICATION_ATTACHED_DIALOG,
-                    "mChildAppWindowAbove");
-            mChildAppWindowBelow = createCommonWindow(mAppWindow,
-                    TYPE_APPLICATION_MEDIA_OVERLAY,
-                    "mChildAppWindowBelow");
-            mDisplayContent.getInsetsPolicy().setRemoteInsetsControllerControlsSystemBars(false);
-
-            // Adding a display will cause freezing the display. Make sure to wait until it's
-            // unfrozen to not run into race conditions with the tests.
-            waitUntilHandlersIdle();
-        } catch (Exception e) {
-            Log.e(TAG, "Failed to set up test", e);
-            throw e;
+        // Only create an additional test display for annotated test class/method because it may
+        // significantly increase the execution time.
+        final Description description = mSystemServicesTestRule.getDescription();
+        UseTestDisplay testDisplayAnnotation = description.getAnnotation(UseTestDisplay.class);
+        if (testDisplayAnnotation == null) {
+            testDisplayAnnotation = description.getTestClass().getAnnotation(UseTestDisplay.class);
+        }
+        if (testDisplayAnnotation != null) {
+            createTestDisplay(testDisplayAnnotation);
+        } else {
+            mDisplayContent = mDefaultDisplay;
         }
     }
 
-    void beforeCreateDisplay() {
+    private void createTestDisplay(UseTestDisplay annotation) {
+        beforeCreateTestDisplay();
+        mDisplayContent = createNewDisplay(true /* supportIme */);
+
+        final boolean addAll = annotation.addAllCommonWindows();
+        final @CommonTypes int[] requestedWindows = annotation.addWindows();
+
+        if (addAll || ArrayUtils.contains(requestedWindows, W_WALLPAPER)) {
+            mWallpaperWindow = createCommonWindow(null, TYPE_WALLPAPER, "wallpaperWindow");
+        }
+        if (addAll || ArrayUtils.contains(requestedWindows, W_INPUT_METHOD)) {
+            mImeWindow = createCommonWindow(null, TYPE_INPUT_METHOD, "mImeWindow");
+            mDisplayContent.mInputMethodWindow = mImeWindow;
+        }
+        if (addAll || ArrayUtils.contains(requestedWindows, W_INPUT_METHOD_DIALOG)) {
+            mImeDialogWindow = createCommonWindow(null, TYPE_INPUT_METHOD_DIALOG,
+                    "mImeDialogWindow");
+        }
+        if (addAll || ArrayUtils.contains(requestedWindows, W_STATUS_BAR)) {
+            mStatusBarWindow = createCommonWindow(null, TYPE_STATUS_BAR, "mStatusBarWindow");
+        }
+        if (addAll || ArrayUtils.contains(requestedWindows, W_NOTIFICATION_SHADE)) {
+            mNotificationShadeWindow = createCommonWindow(null, TYPE_NOTIFICATION_SHADE,
+                    "mNotificationShadeWindow");
+        }
+        if (addAll || ArrayUtils.contains(requestedWindows, W_NAVIGATION_BAR)) {
+            mNavBarWindow = createCommonWindow(null, TYPE_NAVIGATION_BAR, "mNavBarWindow");
+        }
+        if (addAll || ArrayUtils.contains(requestedWindows, W_DOCK_DIVIDER)) {
+            mDockedDividerWindow = createCommonWindow(null, TYPE_DOCK_DIVIDER,
+                    "mDockedDividerWindow");
+        }
+        final boolean addAboveApp = ArrayUtils.contains(requestedWindows, W_ABOVE_ACTIVITY);
+        final boolean addBelowApp = ArrayUtils.contains(requestedWindows, W_BELOW_ACTIVITY);
+        if (addAll || addAboveApp || addBelowApp
+                || ArrayUtils.contains(requestedWindows, W_ACTIVITY)) {
+            mAppWindow = createCommonWindow(null, TYPE_BASE_APPLICATION, "mAppWindow");
+        }
+        if (addAll || addAboveApp) {
+            mChildAppWindowAbove = createCommonWindow(mAppWindow, TYPE_APPLICATION_ATTACHED_DIALOG,
+                    "mChildAppWindowAbove");
+        }
+        if (addAll || addBelowApp) {
+            mChildAppWindowBelow = createCommonWindow(mAppWindow, TYPE_APPLICATION_MEDIA_OVERLAY,
+                    "mChildAppWindowBelow");
+        }
+
+        mDisplayContent.getInsetsPolicy().setRemoteInsetsControllerControlsSystemBars(false);
+
+        // Adding a display will cause freezing the display. Make sure to wait until it's
+        // unfrozen to not run into race conditions with the tests.
+        waitUntilHandlersIdle();
+    }
+
+    void beforeCreateTestDisplay() {
         // Called before display is created.
     }
 
@@ -392,4 +439,66 @@ class WindowTestsBase extends SystemServiceTestsBase {
     void removeGlobalMinSizeRestriction() {
         mWm.mAtmService.mRootWindowContainer.mDefaultMinSizeOfResizeableTaskDp = 1;
     }
+
+    /**
+     * Avoids rotating screen disturbed by some conditions. It is usually used for the default
+     * display that is not the instance of {@link TestDisplayContent} (it bypasses the conditions).
+     *
+     * @see DisplayRotation#updateRotationUnchecked
+     */
+    void unblockDisplayRotation(DisplayContent dc) {
+        mWm.stopFreezingDisplayLocked();
+        // The rotation animation won't actually play, it needs to be cleared manually.
+        dc.setRotationAnimation(null);
+    }
+
+    // The window definition for UseTestDisplay#addWindows. The test can declare to add only
+    // necessary windows, that avoids adding unnecessary overhead of unused windows.
+    static final int W_NOTIFICATION_SHADE = TYPE_NOTIFICATION_SHADE;
+    static final int W_STATUS_BAR = TYPE_STATUS_BAR;
+    static final int W_NAVIGATION_BAR = TYPE_NAVIGATION_BAR;
+    static final int W_INPUT_METHOD_DIALOG = TYPE_INPUT_METHOD_DIALOG;
+    static final int W_INPUT_METHOD = TYPE_INPUT_METHOD;
+    static final int W_DOCK_DIVIDER = TYPE_DOCK_DIVIDER;
+    static final int W_ABOVE_ACTIVITY = TYPE_APPLICATION_ATTACHED_DIALOG;
+    static final int W_ACTIVITY = TYPE_BASE_APPLICATION;
+    static final int W_BELOW_ACTIVITY = TYPE_APPLICATION_MEDIA_OVERLAY;
+    static final int W_WALLPAPER = TYPE_WALLPAPER;
+
+    /** The common window types supported by {@link UseTestDisplay}. */
+    @Retention(RetentionPolicy.RUNTIME)
+    @IntDef(value = {
+            W_NOTIFICATION_SHADE,
+            W_STATUS_BAR,
+            W_NAVIGATION_BAR,
+            W_INPUT_METHOD_DIALOG,
+            W_INPUT_METHOD,
+            W_DOCK_DIVIDER,
+            W_ABOVE_ACTIVITY,
+            W_ACTIVITY,
+            W_BELOW_ACTIVITY,
+            W_WALLPAPER,
+    })
+    @interface CommonTypes {
+    }
+
+    /**
+     * The annotation for class and method (higher priority) to create a non-default display that
+     * will be assigned to {@link #mDisplayContent}. It is used if the test needs
+     * <ul>
+     * <li>Pure empty display.</li>
+     * <li>Configured common windows.</li>
+     * <li>Independent and customizable orientation.</li>
+     * <li>Cross display operation.</li>
+     * </ul>
+     *
+     * @see TestDisplayContent
+     * @see #createTestDisplay
+     **/
+    @Target({ ElementType.METHOD, ElementType.TYPE })
+    @Retention(RetentionPolicy.RUNTIME)
+    @interface UseTestDisplay {
+        boolean addAllCommonWindows() default false;
+        @CommonTypes int[] addWindows() default {};
+    }
 }
diff --git a/services/tests/wmtests/src/com/android/server/wm/WindowTokenTests.java b/services/tests/wmtests/src/com/android/server/wm/WindowTokenTests.java
index 5264e9aaba272..cef202cdc5d51 100644
--- a/services/tests/wmtests/src/com/android/server/wm/WindowTokenTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/WindowTokenTests.java
@@ -142,6 +142,7 @@ public class WindowTokenTests extends WindowTestsBase {
         assertEquals(0, token.getWindowsCount());
     }
 
+    @UseTestDisplay(addWindows = { W_ACTIVITY, W_WALLPAPER })
     @Test
     public void testFinishFixedRotationTransform() {
         final WindowToken appToken = mAppWindow.mToken;
diff --git a/services/tests/wmtests/src/com/android/server/wm/ZOrderingTests.java b/services/tests/wmtests/src/com/android/server/wm/ZOrderingTests.java
index e6b4e0f4baf88..dfb7280e7e540 100644
--- a/services/tests/wmtests/src/com/android/server/wm/ZOrderingTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/ZOrderingTests.java
@@ -62,6 +62,7 @@ import java.util.function.Function;
  */
 @SmallTest
 @Presubmit
+@WindowTestsBase.UseTestDisplay(addAllCommonWindows = true)
 @RunWith(WindowTestRunner.class)
 public class ZOrderingTests extends WindowTestsBase {
 
@@ -152,7 +153,7 @@ public class ZOrderingTests extends WindowTestsBase {
     private LayerRecordingTransaction mTransaction;
 
     @Override
-    void beforeCreateDisplay() {
+    void beforeCreateTestDisplay() {
         // We can't use @Before here because it may happen after WindowTestsBase @Before
         // which is after construction of the DisplayContent, meaning the HierarchyRecorder
         // would miss construction of the top-level layers.