Revert submission

Reason for revert: Potential culprit for b/145602966 - Verifying through Forrest before revert submission

Test: n/a
Change-Id: I5938fe60283bcf9e785c011125c9ddba66a71d68
diff --git a/tests/framework/base/windowmanager/backgroundactivity/src/android/server/wm/BackgroundActivityLaunchTest.java b/tests/framework/base/windowmanager/backgroundactivity/src/android/server/wm/BackgroundActivityLaunchTest.java
index 3f3ec3a..34e3dbd 100644
--- a/tests/framework/base/windowmanager/backgroundactivity/src/android/server/wm/BackgroundActivityLaunchTest.java
+++ b/tests/framework/base/windowmanager/backgroundactivity/src/android/server/wm/BackgroundActivityLaunchTest.java
@@ -64,6 +64,7 @@
 import android.server.wm.backgroundactivity.common.EventReceiver;
 
 import androidx.annotation.Nullable;
+import androidx.test.InstrumentationRegistry;
 import androidx.test.filters.FlakyTest;
 
 import com.android.compatibility.common.util.AppOpsUtils;
@@ -98,15 +99,20 @@
      */
     private static final int BROADCAST_DELIVERY_TIMEOUT_MS = 60000;
 
-    @Override
     @Before
     public void setUp() throws Exception {
+        mContext = InstrumentationRegistry.getContext();
+        mAm = mContext.getSystemService(ActivityManager.class);
+        mAtm = mContext.getSystemService(ActivityTaskManager.class);
+
         // disable SAW appopp for AppA (it's granted autonatically when installed in CTS)
         AppOpsUtils.setOpMode(APP_A_PACKAGE_NAME, "android:system_alert_window", MODE_ERRORED);
         assertEquals(AppOpsUtils.getOpMode(APP_A_PACKAGE_NAME, "android:system_alert_window"),
                 MODE_ERRORED);
 
-        super.setUp();
+        pressWakeupButton();
+        pressUnlockButton();
+        removeStacksWithActivityTypes(ALL_ACTIVITY_TYPE_BUT_HOME);
         assertNull(mAmWmState.getAmState().getTaskByActivity(APP_A_BACKGROUND_ACTIVITY));
         assertNull(mAmWmState.getAmState().getTaskByActivity(APP_A_FOREGROUND_ACTIVITY));
         assertNull(mAmWmState.getAmState().getTaskByActivity(APP_B_FOREGROUND_ACTIVITY));
@@ -126,8 +132,11 @@
                     + APP_A_SIMPLE_ADMIN_RECEIVER.flattenToString());
         });
 
+        super.tearDown();
+
         stopTestPackage(TEST_PACKAGE_APP_A);
         stopTestPackage(TEST_PACKAGE_APP_B);
+        launchHomeActivity();
         AppOpsUtils.reset(APP_A_PACKAGE_NAME);
     }
 
diff --git a/tests/framework/base/windowmanager/src/android/server/wm/ActivityViewTest.java b/tests/framework/base/windowmanager/src/android/server/wm/ActivityViewTest.java
index ecc5aa4..069bedb 100644
--- a/tests/framework/base/windowmanager/src/android/server/wm/ActivityViewTest.java
+++ b/tests/framework/base/windowmanager/src/android/server/wm/ActivityViewTest.java
@@ -59,6 +59,7 @@
 import com.android.cts.mockime.ImeCommand;
 import com.android.cts.mockime.ImeEvent;
 import com.android.cts.mockime.ImeEventStream;
+import com.android.cts.mockime.ImeSettings;
 import com.android.cts.mockime.MockImeSession;
 
 import org.junit.After;
@@ -99,11 +100,16 @@
     }
 
     @After
-    public void tearDown() throws Throwable {
+    public void tearDown() throws Exception {
+        super.tearDown();
         if (mActivityView != null) {
-            // Detach ActivityView before releasing to avoid accessing removed display.
-            mActivityRule.runOnUiThread(
-                    () -> ((ViewGroup) mActivityView.getParent()).removeView(mActivityView));
+            try {
+              // Detach ActivityView before releasing to avoid accessing removed display.
+              mActivityRule.runOnUiThread(
+                      () -> ((ViewGroup) mActivityView.getParent()).removeView(mActivityView));
+            } catch (Throwable e) {
+              throw new Exception(e);
+            }
             SystemUtil.runWithShellPermissionIdentity(() -> mActivityView.release());
         }
     }
@@ -220,40 +226,43 @@
         extras.putString(EXTRA_PRIVATE_IME_OPTIONS, privateImeOptions);
         extras.putParcelable(EXTRA_TEST_CURSOR_ANCHOR_INFO, mockResult);
 
-        final MockImeSession imeSession = MockImeHelper.createManagedMockImeSession(this);
-        final ImeEventStream stream = imeSession.openEventStream();
-        launchActivityInActivityView(INPUT_METHOD_TEST_ACTIVITY, extras);
+        try (final MockImeSession imeSession = MockImeSession.create(mContext,
+                mInstrumentation.getUiAutomation(), new ImeSettings.Builder())) {
+            final ImeEventStream stream = imeSession.openEventStream();
+            launchActivityInActivityView(INPUT_METHOD_TEST_ACTIVITY, extras);
 
-        // IME's seeing uniqueStringValue means that a valid connection is successfully
-        // established from INPUT_METHOD_TEST_ACTIVITY the MockIme.
-        expectEvent(stream, editorMatcher("onStartInput", privateImeOptions), IME_EVENT_TIMEOUT);
+            // IME's seeing uniqueStringValue means that a valid connection is successfully
+            // established from INPUT_METHOD_TEST_ACTIVITY the MockIme
+            expectEvent(stream, editorMatcher("onStartInput", privateImeOptions),
+                    IME_EVENT_TIMEOUT);
 
-        // Make sure that InputConnection#requestCursorUpdates() works.
-        final ImeCommand cursorUpdatesCommand = imeSession.callRequestCursorUpdates(
-                InputConnection.CURSOR_UPDATE_IMMEDIATE);
-        final ImeEvent cursorUpdatesEvent = expectCommand(
-                stream, cursorUpdatesCommand, IME_EVENT_TIMEOUT);
-        assertTrue(cursorUpdatesEvent.getReturnBooleanValue());
+            // Make sure that InputConnection#requestCursorUpdates() works.
+            final ImeCommand cursorUpdatesCommand = imeSession.callRequestCursorUpdates(
+                    InputConnection.CURSOR_UPDATE_IMMEDIATE);
+            final ImeEvent cursorUpdatesEvent = expectCommand(
+                    stream, cursorUpdatesCommand, IME_EVENT_TIMEOUT);
+            assertTrue(cursorUpdatesEvent.getReturnBooleanValue());
 
-        // Make sure that MockIme received the object sent above.
-        final CursorAnchorInfo receivedInfo = expectEvent(stream,
-                event -> "onUpdateCursorAnchorInfo".equals(event.getEventName()),
-                IME_EVENT_TIMEOUT).getArguments().getParcelable("cursorAnchorInfo");
-        assertNotNull(receivedInfo);
+            // Make sure that MockIme received the object sent above.
+            final CursorAnchorInfo receivedInfo = expectEvent(stream,
+                    event -> "onUpdateCursorAnchorInfo".equals(event.getEventName()),
+                    IME_EVENT_TIMEOUT).getArguments().getParcelable("cursorAnchorInfo");
+            assertNotNull(receivedInfo);
 
-        // Get the location of ActivityView in the default display's screen coordinates.
-        final AtomicReference<Point> offsetRef = new AtomicReference<>();
-        mInstrumentation.runOnMainSync(() -> {
-            final int[] xy = new int[2];
-            mActivityView.getLocationOnScreen(xy);
-            offsetRef.set(new Point(xy[0], xy[1]));
-        });
-        final Point offset = offsetRef.get();
+            // Get the location of ActivityView in the default display's screen coordinates.
+            final AtomicReference<Point> offsetRef = new AtomicReference<>();
+            mInstrumentation.runOnMainSync(() -> {
+                final int[] xy = new int[2];
+                mActivityView.getLocationOnScreen(xy);
+                offsetRef.set(new Point(xy[0], xy[1]));
+            });
+            final Point offset = offsetRef.get();
 
-        // Make sure that the received CursorAnchorInfo has an adjusted Matrix.
-        final Matrix expectedMatrix = mockResult.getMatrix();
-        expectedMatrix.postTranslate(offset.x, offset.y);
-        assertEquals(expectedMatrix, receivedInfo.getMatrix());
+            // Make sure that the received CursorAnchorInfo has an adjusted Matrix.
+            final Matrix expectedMatrix = mockResult.getMatrix();
+            expectedMatrix.postTranslate(offset.x, offset.y);
+            assertEquals(expectedMatrix, receivedInfo.getMatrix());
+        }
     }
 
     private void launchActivityInActivityView(ComponentName activity) {
diff --git a/tests/framework/base/windowmanager/src/android/server/wm/ActivityVisibilityTests.java b/tests/framework/base/windowmanager/src/android/server/wm/ActivityVisibilityTests.java
index f062954..226f48a 100644
--- a/tests/framework/base/windowmanager/src/android/server/wm/ActivityVisibilityTests.java
+++ b/tests/framework/base/windowmanager/src/android/server/wm/ActivityVisibilityTests.java
@@ -220,10 +220,11 @@
 
     @Test
     public void testTurnScreenOnActivity() {
-        final LockScreenSession lockScreenSession = createManagedLockScreenSession();
-        final ActivitySessionClient activityClient = createManagedActivityClientSession();
-        testTurnScreenOnActivity(lockScreenSession, activityClient, true /* useWindowFlags */);
-        testTurnScreenOnActivity(lockScreenSession, activityClient, false /* useWindowFlags */);
+        try (final LockScreenSession lockScreenSession = new LockScreenSession();
+             final ActivitySessionClient activityClient = new ActivitySessionClient(mContext)) {
+            testTurnScreenOnActivity(lockScreenSession, activityClient, true /* useWindowFlags */);
+            testTurnScreenOnActivity(lockScreenSession, activityClient, false /* useWindowFlags */);
+        }
     }
 
     private void testTurnScreenOnActivity(LockScreenSession lockScreenSession,
@@ -445,99 +446,108 @@
     }
 
     @Test
-    public void testTurnScreenOnAttrNoLockScreen() {
-        final LockScreenSession lockScreenSession = createManagedLockScreenSession();
-        lockScreenSession.disableLockScreen().sleepDevice();
-        separateTestJournal();
-        launchActivity(TURN_SCREEN_ON_ATTR_ACTIVITY);
-        mAmWmState.assertVisibility(TURN_SCREEN_ON_ATTR_ACTIVITY, true);
-        assertTrue("Display turns on", isDisplayOn(DEFAULT_DISPLAY));
-        assertSingleLaunch(TURN_SCREEN_ON_ATTR_ACTIVITY);
+    public void testTurnScreenOnAttrNoLockScreen() throws Exception {
+        try (final LockScreenSession lockScreenSession = new LockScreenSession()) {
+            lockScreenSession.disableLockScreen()
+                    .sleepDevice();
+            separateTestJournal();
+            launchActivity(TURN_SCREEN_ON_ATTR_ACTIVITY);
+            mAmWmState.assertVisibility(TURN_SCREEN_ON_ATTR_ACTIVITY, true);
+            assertTrue("Display turns on", isDisplayOn(DEFAULT_DISPLAY));
+            assertSingleLaunch(TURN_SCREEN_ON_ATTR_ACTIVITY);
+        }
     }
 
     @Test
-    public void testTurnScreenOnAttrWithLockScreen() {
+    public void testTurnScreenOnAttrWithLockScreen() throws Exception {
         assumeTrue(supportsSecureLock());
 
-        final LockScreenSession lockScreenSession = createManagedLockScreenSession();
-        lockScreenSession.setLockCredential().sleepDevice();
-        separateTestJournal();
-        launchActivityNoWait(TURN_SCREEN_ON_ATTR_ACTIVITY);
-        // Wait for the activity stopped because lock screen prevent showing the activity.
-        mAmWmState.waitForActivityState(TURN_SCREEN_ON_ATTR_ACTIVITY, STATE_STOPPED);
-        assertFalse("Display keeps off", isDisplayOn(DEFAULT_DISPLAY));
-        assertSingleLaunchAndStop(TURN_SCREEN_ON_ATTR_ACTIVITY);
+        try (final LockScreenSession lockScreenSession = new LockScreenSession()) {
+            lockScreenSession.setLockCredential()
+                    .sleepDevice();
+            separateTestJournal();
+            launchActivityNoWait(TURN_SCREEN_ON_ATTR_ACTIVITY);
+            // Wait for the activity stopped because lock screen prevent showing the activity.
+            mAmWmState.waitForActivityState(TURN_SCREEN_ON_ATTR_ACTIVITY, STATE_STOPPED);
+            assertFalse("Display keeps off", isDisplayOn(DEFAULT_DISPLAY));
+            assertSingleLaunchAndStop(TURN_SCREEN_ON_ATTR_ACTIVITY);
+        }
     }
 
     @Test
-    public void testTurnScreenOnShowOnLockAttr() {
-        final LockScreenSession lockScreenSession = createManagedLockScreenSession();
-        lockScreenSession.sleepDevice();
-        mAmWmState.waitForAllStoppedActivities();
-        separateTestJournal();
-        launchActivity(TURN_SCREEN_ON_SHOW_ON_LOCK_ACTIVITY);
-        mAmWmState.assertVisibility(TURN_SCREEN_ON_SHOW_ON_LOCK_ACTIVITY, true);
-        assertTrue("Display turns on", isDisplayOn(DEFAULT_DISPLAY));
-        assertSingleLaunch(TURN_SCREEN_ON_SHOW_ON_LOCK_ACTIVITY);
+    public void testTurnScreenOnShowOnLockAttr() throws Exception {
+        try (final LockScreenSession lockScreenSession = new LockScreenSession()) {
+            lockScreenSession.sleepDevice();
+            mAmWmState.waitForAllStoppedActivities();
+            separateTestJournal();
+            launchActivity(TURN_SCREEN_ON_SHOW_ON_LOCK_ACTIVITY);
+            mAmWmState.assertVisibility(TURN_SCREEN_ON_SHOW_ON_LOCK_ACTIVITY, true);
+            assertTrue("Display turns on", isDisplayOn(DEFAULT_DISPLAY));
+            assertSingleLaunch(TURN_SCREEN_ON_SHOW_ON_LOCK_ACTIVITY);
+        }
     }
 
     @Test
-    public void testTurnScreenOnAttrRemove() {
-        final LockScreenSession lockScreenSession = createManagedLockScreenSession();
-        lockScreenSession.sleepDevice();
-        mAmWmState.waitForAllStoppedActivities();
-        separateTestJournal();
-        launchActivity(TURN_SCREEN_ON_ATTR_REMOVE_ATTR_ACTIVITY);
-        assertTrue("Display turns on", isDisplayOn(DEFAULT_DISPLAY));
-        assertSingleLaunch(TURN_SCREEN_ON_ATTR_REMOVE_ATTR_ACTIVITY);
+    public void testTurnScreenOnAttrRemove() throws Exception {
+        try (final LockScreenSession lockScreenSession = new LockScreenSession()) {
+            lockScreenSession.sleepDevice();
+            mAmWmState.waitForAllStoppedActivities();
+            separateTestJournal();
+            launchActivity(TURN_SCREEN_ON_ATTR_REMOVE_ATTR_ACTIVITY);
+            assertTrue("Display turns on", isDisplayOn(DEFAULT_DISPLAY));
+            assertSingleLaunch(TURN_SCREEN_ON_ATTR_REMOVE_ATTR_ACTIVITY);
 
-        lockScreenSession.sleepDevice();
-        mAmWmState.waitForAllStoppedActivities();
-        separateTestJournal();
-        launchActivity(TURN_SCREEN_ON_ATTR_REMOVE_ATTR_ACTIVITY);
-        mAmWmState.waitForActivityState(TURN_SCREEN_ON_ATTR_REMOVE_ATTR_ACTIVITY, STATE_STOPPED);
-        // Display should keep off, because setTurnScreenOn(false) has been called at
-        // {@link TURN_SCREEN_ON_ATTR_REMOVE_ATTR_ACTIVITY}'s onStop.
-        assertFalse("Display keeps off", isDisplayOn(DEFAULT_DISPLAY));
-        assertSingleStartAndStop(TURN_SCREEN_ON_ATTR_REMOVE_ATTR_ACTIVITY);
+            lockScreenSession.sleepDevice();
+            mAmWmState.waitForAllStoppedActivities();
+            separateTestJournal();
+            launchActivity(TURN_SCREEN_ON_ATTR_REMOVE_ATTR_ACTIVITY);
+            mAmWmState.waitForActivityState(TURN_SCREEN_ON_ATTR_REMOVE_ATTR_ACTIVITY,
+                STATE_STOPPED);
+            // Display should keep off, because setTurnScreenOn(false) has been called at
+            // {@link TURN_SCREEN_ON_ATTR_REMOVE_ATTR_ACTIVITY}'s onStop.
+            assertFalse("Display keeps off", isDisplayOn(DEFAULT_DISPLAY));
+            assertSingleStartAndStop(TURN_SCREEN_ON_ATTR_REMOVE_ATTR_ACTIVITY);
+        }
     }
 
     @Test
-    public void testTurnScreenOnSingleTask() {
-        final LockScreenSession lockScreenSession = createManagedLockScreenSession();
-        lockScreenSession.sleepDevice();
-        separateTestJournal();
-        launchActivity(TURN_SCREEN_ON_SINGLE_TASK_ACTIVITY);
-        mAmWmState.assertVisibility(TURN_SCREEN_ON_SINGLE_TASK_ACTIVITY, true);
-        assertTrue("Display turns on", isDisplayOn(DEFAULT_DISPLAY));
-        assertSingleLaunch(TURN_SCREEN_ON_SINGLE_TASK_ACTIVITY);
+    public void testTurnScreenOnSingleTask() throws Exception {
+        try (final LockScreenSession lockScreenSession = new LockScreenSession()) {
+            lockScreenSession.sleepDevice();
+            separateTestJournal();
+            launchActivity(TURN_SCREEN_ON_SINGLE_TASK_ACTIVITY);
+            mAmWmState.assertVisibility(TURN_SCREEN_ON_SINGLE_TASK_ACTIVITY, true);
+            assertTrue("Display turns on", isDisplayOn(DEFAULT_DISPLAY));
+            assertSingleLaunch(TURN_SCREEN_ON_SINGLE_TASK_ACTIVITY);
 
-        lockScreenSession.sleepDevice();
-        // We should make sure test activity stopped to prevent a false alarm stop state
-        // included in the lifecycle count.
-        waitAndAssertActivityState(TURN_SCREEN_ON_SINGLE_TASK_ACTIVITY, STATE_STOPPED,
-                "Activity should be stopped");
-        separateTestJournal();
-        launchActivity(TURN_SCREEN_ON_SINGLE_TASK_ACTIVITY);
-        mAmWmState.assertVisibility(TURN_SCREEN_ON_SINGLE_TASK_ACTIVITY, true);
-        // Wait more for display state change since turning the display ON may take longer
-        // and reported after the activity launch.
-        waitForDefaultDisplayState(true /* wantOn */);
-        assertTrue("Display turns on", isDisplayOn(DEFAULT_DISPLAY));
-        assertSingleStart(TURN_SCREEN_ON_SINGLE_TASK_ACTIVITY);
+            lockScreenSession.sleepDevice();
+            // We should make sure test activity stopped to prevent a false alarm stop state
+            // included in the lifecycle count.
+            waitAndAssertActivityState(TURN_SCREEN_ON_SINGLE_TASK_ACTIVITY, STATE_STOPPED,
+                    "Activity should be stopped");
+            separateTestJournal();
+            launchActivity(TURN_SCREEN_ON_SINGLE_TASK_ACTIVITY);
+            mAmWmState.assertVisibility(TURN_SCREEN_ON_SINGLE_TASK_ACTIVITY, true);
+            // Wait more for display state change since turning the display ON may take longer
+            // and reported after the activity launch.
+            waitForDefaultDisplayState(true /* wantOn */);
+            assertTrue("Display turns on", isDisplayOn(DEFAULT_DISPLAY));
+            assertSingleStart(TURN_SCREEN_ON_SINGLE_TASK_ACTIVITY);
+        }
     }
 
     @Test
-    public void testTurnScreenOnActivity_withRelayout() {
-        final LockScreenSession lockScreenSession = createManagedLockScreenSession();
-        lockScreenSession.sleepDevice();
-        launchActivity(TURN_SCREEN_ON_WITH_RELAYOUT_ACTIVITY);
-        mAmWmState.assertVisibility(TURN_SCREEN_ON_WITH_RELAYOUT_ACTIVITY, true);
+    public void testTurnScreenOnActivity_withRelayout() throws Exception {
+        try (final LockScreenSession lockScreenSession = new LockScreenSession()) {
+            lockScreenSession.sleepDevice();
+            launchActivity(TURN_SCREEN_ON_WITH_RELAYOUT_ACTIVITY);
+            mAmWmState.assertVisibility(TURN_SCREEN_ON_WITH_RELAYOUT_ACTIVITY, true);
 
-        lockScreenSession.sleepDevice();
-        waitAndAssertActivityState(TURN_SCREEN_ON_WITH_RELAYOUT_ACTIVITY, STATE_STOPPED,
-                "Activity should be stopped");
-        assertFalse("Display keeps off", isDisplayOn(DEFAULT_DISPLAY));
+            lockScreenSession.sleepDevice();
+            waitAndAssertActivityState(TURN_SCREEN_ON_WITH_RELAYOUT_ACTIVITY, STATE_STOPPED,
+                    "Activity should be stopped");
+            assertFalse("Display keeps off", isDisplayOn(DEFAULT_DISPLAY));
+        }
     }
 
     @Test
@@ -596,69 +606,73 @@
 
     @Test
     public void testConvertTranslucentOnTranslucentActivity() {
-        final ActivitySessionClient activityClient = createManagedActivityClientSession();
-        // Start CONVERT_TRANSLUCENT_DIALOG_ACTIVITY on top of LAUNCHING_ACTIVITY
-        final ActivitySession activity = activityClient.startActivity(
-                getLaunchActivityBuilder().setTargetActivity(TRANSLUCENT_TOP_ACTIVITY));
-        verifyActivityVisibilities(TRANSLUCENT_TOP_ACTIVITY, false);
-        verifyActivityVisibilities(LAUNCHING_ACTIVITY, false);
+        try (final ActivitySessionClient activityClient = new ActivitySessionClient(mContext)) {
+            // Start CONVERT_TRANSLUCENT_DIALOG_ACTIVITY on top of LAUNCHING_ACTIVITY
+            final ActivitySession activity = activityClient.startActivity(
+                    getLaunchActivityBuilder().setTargetActivity(TRANSLUCENT_TOP_ACTIVITY));
+            verifyActivityVisibilities(TRANSLUCENT_TOP_ACTIVITY, false);
+            verifyActivityVisibilities(LAUNCHING_ACTIVITY, false);
 
-        activity.sendCommand(ACTION_CONVERT_FROM_TRANSLUCENT);
-        verifyActivityVisibilities(LAUNCHING_ACTIVITY, true);
+            activity.sendCommand(ACTION_CONVERT_FROM_TRANSLUCENT);
+            verifyActivityVisibilities(LAUNCHING_ACTIVITY, true);
 
-        activity.sendCommand(ACTION_CONVERT_TO_TRANSLUCENT);
-        verifyActivityVisibilities(LAUNCHING_ACTIVITY, false);
+            activity.sendCommand(ACTION_CONVERT_TO_TRANSLUCENT);
+            verifyActivityVisibilities(LAUNCHING_ACTIVITY, false);
+        }
     }
 
     @Test
     public void testConvertTranslucentOnNonTopTranslucentActivity() {
-        final ActivitySessionClient activityClient = createManagedActivityClientSession();
-        final ActivitySession activity = activityClient.startActivity(
-                getLaunchActivityBuilder().setTargetActivity(TRANSLUCENT_TOP_ACTIVITY));
-        getLaunchActivityBuilder().setTargetActivity(SHOW_WHEN_LOCKED_DIALOG_ACTIVITY)
-                .setUseInstrumentation().execute();
-        verifyActivityVisibilities(SHOW_WHEN_LOCKED_DIALOG_ACTIVITY, false);
-        verifyActivityVisibilities(TRANSLUCENT_TOP_ACTIVITY, false);
-        verifyActivityVisibilities(LAUNCHING_ACTIVITY, false);
+        try (final ActivitySessionClient activityClient = new ActivitySessionClient(mContext)) {
+            final ActivitySession activity = activityClient.startActivity(
+                    getLaunchActivityBuilder().setTargetActivity(TRANSLUCENT_TOP_ACTIVITY));
+            getLaunchActivityBuilder().setTargetActivity(SHOW_WHEN_LOCKED_DIALOG_ACTIVITY)
+                    .setUseInstrumentation().execute();
+            verifyActivityVisibilities(SHOW_WHEN_LOCKED_DIALOG_ACTIVITY, false);
+            verifyActivityVisibilities(TRANSLUCENT_TOP_ACTIVITY, false);
+            verifyActivityVisibilities(LAUNCHING_ACTIVITY, false);
 
-        activity.sendCommand(ACTION_CONVERT_FROM_TRANSLUCENT);
-        verifyActivityVisibilities(LAUNCHING_ACTIVITY, true);
+            activity.sendCommand(ACTION_CONVERT_FROM_TRANSLUCENT);
+            verifyActivityVisibilities(LAUNCHING_ACTIVITY, true);
 
-        activity.sendCommand(ACTION_CONVERT_TO_TRANSLUCENT);
-        verifyActivityVisibilities(LAUNCHING_ACTIVITY, false);
+            activity.sendCommand(ACTION_CONVERT_TO_TRANSLUCENT);
+            verifyActivityVisibilities(LAUNCHING_ACTIVITY, false);
+        }
     }
 
     @Test
     public void testConvertTranslucentOnOpaqueActivity() {
-        final ActivitySessionClient activityClient = createManagedActivityClientSession();
-        final ActivitySession activity = activityClient.startActivity(
-                getLaunchActivityBuilder().setTargetActivity(TOP_ACTIVITY));
-        verifyActivityVisibilities(TOP_ACTIVITY, false);
-        verifyActivityVisibilities(LAUNCHING_ACTIVITY, true);
+        try (final ActivitySessionClient activityClient = new ActivitySessionClient(mContext)) {
+            final ActivitySession activity = activityClient.startActivity(
+                    getLaunchActivityBuilder().setTargetActivity(TOP_ACTIVITY));
+            verifyActivityVisibilities(TOP_ACTIVITY, false);
+            verifyActivityVisibilities(LAUNCHING_ACTIVITY, true);
 
-        activity.sendCommand(ACTION_CONVERT_TO_TRANSLUCENT);
-        verifyActivityVisibilities(LAUNCHING_ACTIVITY, false);
+            activity.sendCommand(ACTION_CONVERT_TO_TRANSLUCENT);
+            verifyActivityVisibilities(LAUNCHING_ACTIVITY, false);
 
-        activity.sendCommand(ACTION_CONVERT_FROM_TRANSLUCENT);
-        verifyActivityVisibilities(LAUNCHING_ACTIVITY, true);
+            activity.sendCommand(ACTION_CONVERT_FROM_TRANSLUCENT);
+            verifyActivityVisibilities(LAUNCHING_ACTIVITY, true);
+        }
     }
 
     @Test
     public void testConvertTranslucentOnNonTopOpaqueActivity() {
-        final ActivitySessionClient activityClient = createManagedActivityClientSession();
-        final ActivitySession activity = activityClient.startActivity(
-                getLaunchActivityBuilder().setTargetActivity(TOP_ACTIVITY));
-        getLaunchActivityBuilder().setTargetActivity(SHOW_WHEN_LOCKED_DIALOG_ACTIVITY)
-                .setUseInstrumentation().execute();
-        verifyActivityVisibilities(SHOW_WHEN_LOCKED_DIALOG_ACTIVITY, false);
-        verifyActivityVisibilities(TOP_ACTIVITY, false);
-        verifyActivityVisibilities(LAUNCHING_ACTIVITY, true);
+        try (final ActivitySessionClient activityClient = new ActivitySessionClient(mContext)) {
+            final ActivitySession activity = activityClient.startActivity(
+                    getLaunchActivityBuilder().setTargetActivity(TOP_ACTIVITY));
+            getLaunchActivityBuilder().setTargetActivity(SHOW_WHEN_LOCKED_DIALOG_ACTIVITY)
+                    .setUseInstrumentation().execute();
+            verifyActivityVisibilities(SHOW_WHEN_LOCKED_DIALOG_ACTIVITY, false);
+            verifyActivityVisibilities(TOP_ACTIVITY, false);
+            verifyActivityVisibilities(LAUNCHING_ACTIVITY, true);
 
-        activity.sendCommand(ACTION_CONVERT_TO_TRANSLUCENT);
-        verifyActivityVisibilities(LAUNCHING_ACTIVITY, false);
+            activity.sendCommand(ACTION_CONVERT_TO_TRANSLUCENT);
+            verifyActivityVisibilities(LAUNCHING_ACTIVITY, false);
 
-        activity.sendCommand(ACTION_CONVERT_FROM_TRANSLUCENT);
-        verifyActivityVisibilities(LAUNCHING_ACTIVITY, true);
+            activity.sendCommand(ACTION_CONVERT_FROM_TRANSLUCENT);
+            verifyActivityVisibilities(LAUNCHING_ACTIVITY, true);
+        }
     }
 
     private void verifyActivityVisibilities(ComponentName activityBehind,
diff --git a/tests/framework/base/windowmanager/src/android/server/wm/AlertWindowsTests.java b/tests/framework/base/windowmanager/src/android/server/wm/AlertWindowsTests.java
index 62fa17e..0befd92 100644
--- a/tests/framework/base/windowmanager/src/android/server/wm/AlertWindowsTests.java
+++ b/tests/framework/base/windowmanager/src/android/server/wm/AlertWindowsTests.java
@@ -89,7 +89,9 @@
     }
 
     @After
+    @Override
     public void tearDown() throws Exception {
+        super.tearDown();
         resetPermissionState(ALERT_WINDOW_TEST_ACTIVITY);
         resetPermissionState(SDK25_ALERT_WINDOW_TEST_ACTIVITY);
         stopTestPackage(ALERT_WINDOW_TEST_ACTIVITY.getPackageName());
diff --git a/tests/framework/base/windowmanager/src/android/server/wm/AnrTests.java b/tests/framework/base/windowmanager/src/android/server/wm/AnrTests.java
index a5d9bcd..e561321 100644
--- a/tests/framework/base/windowmanager/src/android/server/wm/AnrTests.java
+++ b/tests/framework/base/windowmanager/src/android/server/wm/AnrTests.java
@@ -72,9 +72,10 @@
     }
 
     @After
-    public void teardown() {
+    public void teardown() throws Exception {
         mHideDialogSetting.close();
         stopTestPackage(UNRESPONSIVE_ACTIVITY.getPackageName());
+        super.tearDown();
     }
 
     @Test
diff --git a/tests/framework/base/windowmanager/src/android/server/wm/AppConfigurationTests.java b/tests/framework/base/windowmanager/src/android/server/wm/AppConfigurationTests.java
index 479fe42..ce2a776 100644
--- a/tests/framework/base/windowmanager/src/android/server/wm/AppConfigurationTests.java
+++ b/tests/framework/base/windowmanager/src/android/server/wm/AppConfigurationTests.java
@@ -120,17 +120,19 @@
      * Tests whether the Display sizes change when rotating the device.
      */
     @Test
-    public void testConfigurationUpdatesWhenRotatingWhileFullscreen() {
+    public void testConfigurationUpdatesWhenRotatingWhileFullscreen() throws Exception {
         assumeTrue("Skipping test: no rotation support", supportsRotation());
 
-        final RotationSession rotationSession = createManagedRotationSession();
-        rotationSession.set(ROTATION_0);
+        try (final RotationSession rotationSession = new RotationSession()) {
+            rotationSession.set(ROTATION_0);
 
-        separateTestJournal();
-        launchActivity(RESIZEABLE_ACTIVITY, WINDOWING_MODE_FULLSCREEN_OR_SPLIT_SCREEN_SECONDARY);
-        final SizeInfo initialSizes = getActivityDisplaySize(RESIZEABLE_ACTIVITY);
+            separateTestJournal();
+            launchActivity(RESIZEABLE_ACTIVITY,
+                    WINDOWING_MODE_FULLSCREEN_OR_SPLIT_SCREEN_SECONDARY);
+            final SizeInfo initialSizes = getActivityDisplaySize(RESIZEABLE_ACTIVITY);
 
-        rotateAndCheckSizes(rotationSession, initialSizes);
+            rotateAndCheckSizes(rotationSession, initialSizes);
+        }
     }
 
     /**
@@ -138,23 +140,24 @@
      * is in the docked stack.
      */
     @Test
-    public void testConfigurationUpdatesWhenRotatingWhileDocked() {
+    public void testConfigurationUpdatesWhenRotatingWhileDocked() throws Exception {
         assumeTrue("Skipping test: no multi-window support", supportsSplitScreenMultiWindow());
 
-        final RotationSession rotationSession = createManagedRotationSession();
-        rotationSession.set(ROTATION_0);
+        try (final RotationSession rotationSession = new RotationSession()) {
+            rotationSession.set(ROTATION_0);
 
-        separateTestJournal();
-        // Launch our own activity to side in case Recents (or other activity to side) doesn't
-        // support rotation.
-        launchActivitiesInSplitScreen(
-                getLaunchActivityBuilder().setTargetActivity(LAUNCHING_ACTIVITY),
-                getLaunchActivityBuilder().setTargetActivity(TEST_ACTIVITY));
-        // Launch target activity in docked stack.
-        getLaunchActivityBuilder().setTargetActivity(RESIZEABLE_ACTIVITY).execute();
-        final SizeInfo initialSizes = getActivityDisplaySize(RESIZEABLE_ACTIVITY);
+            separateTestJournal();
+            // Launch our own activity to side in case Recents (or other activity to side) doesn't
+            // support rotation.
+            launchActivitiesInSplitScreen(
+                    getLaunchActivityBuilder().setTargetActivity(LAUNCHING_ACTIVITY),
+                    getLaunchActivityBuilder().setTargetActivity(TEST_ACTIVITY));
+            // Launch target activity in docked stack.
+            getLaunchActivityBuilder().setTargetActivity(RESIZEABLE_ACTIVITY).execute();
+            final SizeInfo initialSizes = getActivityDisplaySize(RESIZEABLE_ACTIVITY);
 
-        rotateAndCheckSizes(rotationSession, initialSizes);
+            rotateAndCheckSizes(rotationSession, initialSizes);
+        }
     }
 
     /**
@@ -162,22 +165,24 @@
      * is launched to side from docked stack.
      */
     @Test
-    public void testConfigurationUpdatesWhenRotatingToSideFromDocked() {
+    public void testConfigurationUpdatesWhenRotatingToSideFromDocked() throws Exception {
         assumeTrue("Skipping test: no multi-window support", supportsSplitScreenMultiWindow());
 
-        final RotationSession rotationSession = createManagedRotationSession();
-        rotationSession.set(ROTATION_0);
+        try (final RotationSession rotationSession = new RotationSession()) {
+            rotationSession.set(ROTATION_0);
 
-        separateTestJournal();
-        launchActivitiesInSplitScreen(
-                getLaunchActivityBuilder().setTargetActivity(LAUNCHING_ACTIVITY),
-                getLaunchActivityBuilder().setTargetActivity(RESIZEABLE_ACTIVITY));
-        final SizeInfo initialSizes = getActivityDisplaySize(RESIZEABLE_ACTIVITY);
+            separateTestJournal();
+            launchActivitiesInSplitScreen(
+                    getLaunchActivityBuilder().setTargetActivity(LAUNCHING_ACTIVITY),
+                    getLaunchActivityBuilder().setTargetActivity(RESIZEABLE_ACTIVITY));
+            final SizeInfo initialSizes = getActivityDisplaySize(RESIZEABLE_ACTIVITY);
 
-        rotateAndCheckSizes(rotationSession, initialSizes);
+            rotateAndCheckSizes(rotationSession, initialSizes);
+        }
     }
 
-    private void rotateAndCheckSizes(RotationSession rotationSession, SizeInfo prevSizes) {
+    private void rotateAndCheckSizes(RotationSession rotationSession, SizeInfo prevSizes)
+            throws Exception {
         final ActivityManagerState.ActivityTask task =
                 mAmWmState.getAmState().getTaskByActivity(RESIZEABLE_ACTIVITY);
         final int displayId = mAmWmState.getAmState().getStackById(task.mStackId).mDisplayId;
@@ -459,16 +464,16 @@
             //cannot physically rotate the screen on automotive device, skip
             return;
         }
+        try (final RotationSession rotationSession = new RotationSession()) {
+            rotationSession.set(ROTATION_0);
 
-        final RotationSession rotationSession = createManagedRotationSession();
-        rotationSession.set(ROTATION_0);
-
-        launchActivity(SDK26_TRANSLUCENT_LANDSCAPE_ACTIVITY);
-        mAmWmState.assertResumedActivity(
-                "target SDK <= 26 non-fullscreen activity should be allowed to launch",
-                SDK26_TRANSLUCENT_LANDSCAPE_ACTIVITY);
-        assertEquals("non-fullscreen activity requested landscape orientation",
-                0 /* landscape */, mAmWmState.getWmState().getLastOrientation());
+            launchActivity(SDK26_TRANSLUCENT_LANDSCAPE_ACTIVITY);
+            mAmWmState.assertResumedActivity(
+                    "target SDK <= 26 non-fullscreen activity should be allowed to launch",
+                    SDK26_TRANSLUCENT_LANDSCAPE_ACTIVITY);
+            assertEquals("non-fullscreen activity requested landscape orientation",
+                    0 /* landscape */, mAmWmState.getWmState().getLastOrientation());
+        }
     }
 
     /**
@@ -569,28 +574,29 @@
 
         // Rotate the activity and check that it receives configuration changes with a different
         // orientation each time.
-        final RotationSession rotationSession = createManagedRotationSession();
-        assumeTrue("Skipping test: no locked user rotation mode support.",
-                supportsLockedUserRotation(rotationSession, displayId));
+        try (final RotationSession rotationSession = new RotationSession()) {
+            assumeTrue("Skipping test: no locked user rotation mode support.",
+                    supportsLockedUserRotation(rotationSession, displayId));
 
-        rotationSession.set(ROTATION_0);
-        SizeInfo reportedSizes = getLastReportedSizesForActivity(RESIZEABLE_ACTIVITY);
-        int prevOrientation = reportedSizes.orientation;
+            rotationSession.set(ROTATION_0);
+            SizeInfo reportedSizes = getLastReportedSizesForActivity(RESIZEABLE_ACTIVITY);
+            int prevOrientation = reportedSizes.orientation;
 
-        final int[] rotations = { ROTATION_270, ROTATION_180, ROTATION_90, ROTATION_0 };
-        for (final int rotation : rotations) {
-            separateTestJournal();
-            rotationSession.set(rotation);
+            final int[] rotations = { ROTATION_270, ROTATION_180, ROTATION_90, ROTATION_0 };
+            for (final int rotation : rotations) {
+                separateTestJournal();
+                rotationSession.set(rotation);
 
-            // Verify lifecycle count and orientation changes.
-            assertRelaunchOrConfigChanged(RESIZEABLE_ACTIVITY, 0 /* numRelaunch */,
-                    1 /* numConfigChange */);
-            reportedSizes = getLastReportedSizesForActivity(RESIZEABLE_ACTIVITY);
-            assertNotEquals(prevOrientation, reportedSizes.orientation);
-            assertRelaunchOrConfigChanged(TEST_ACTIVITY, 0 /* numRelaunch */,
-                    0 /* numConfigChange */);
+                // Verify lifecycle count and orientation changes.
+                assertRelaunchOrConfigChanged(RESIZEABLE_ACTIVITY, 0 /* numRelaunch */,
+                        1 /* numConfigChange */);
+                reportedSizes = getLastReportedSizesForActivity(RESIZEABLE_ACTIVITY);
+                assertNotEquals(prevOrientation, reportedSizes.orientation);
+                assertRelaunchOrConfigChanged(TEST_ACTIVITY, 0 /* numRelaunch */,
+                        0 /* numConfigChange */);
 
-            prevOrientation = reportedSizes.orientation;
+                prevOrientation = reportedSizes.orientation;
+            }
         }
     }
 
@@ -616,25 +622,26 @@
                 .getDisplayByActivity(PORTRAIT_ORIENTATION_ACTIVITY);
 
         // Rotate the activity and check that the orientation doesn't change
-        final RotationSession rotationSession = createManagedRotationSession();
-        assumeTrue("Skipping test: no user locked rotation support.",
-                supportsLockedUserRotation(rotationSession, displayId));
+        try (final RotationSession rotationSession = new RotationSession()) {
+            assumeTrue("Skipping test: no user locked rotation support.",
+                    supportsLockedUserRotation(rotationSession, displayId));
 
-        rotationSession.set(ROTATION_0);
+            rotationSession.set(ROTATION_0);
 
-        final int[] rotations = { ROTATION_270, ROTATION_180, ROTATION_90, ROTATION_0 };
-        for (final int rotation : rotations) {
-            separateTestJournal();
-            rotationSession.set(rotation);
+            final int[] rotations = { ROTATION_270, ROTATION_180, ROTATION_90, ROTATION_0 };
+            for (final int rotation : rotations) {
+                separateTestJournal();
+                rotationSession.set(rotation);
 
-            // Verify lifecycle count and orientation changes.
-            assertRelaunchOrConfigChanged(PORTRAIT_ORIENTATION_ACTIVITY, 0 /* numRelaunch */,
-                    0 /* numConfigChange */);
-            final SizeInfo reportedSizes = getLastReportedSizesForActivity(
-                    PORTRAIT_ORIENTATION_ACTIVITY);
-            assertNull("No new sizes must be reported", reportedSizes);
-            assertRelaunchOrConfigChanged(RESIZEABLE_ACTIVITY, 0 /* numRelaunch */,
-                    0 /* numConfigChange */);
+                // Verify lifecycle count and orientation changes.
+                assertRelaunchOrConfigChanged(PORTRAIT_ORIENTATION_ACTIVITY, 0 /* numRelaunch */,
+                        0 /* numConfigChange */);
+                final SizeInfo reportedSizes = getLastReportedSizesForActivity(
+                        PORTRAIT_ORIENTATION_ACTIVITY);
+                assertNull("No new sizes must be reported", reportedSizes);
+                assertRelaunchOrConfigChanged(RESIZEABLE_ACTIVITY, 0 /* numRelaunch */,
+                        0 /* numConfigChange */);
+            }
         }
     }
 
@@ -676,8 +683,10 @@
         assumeTrue("Skipping test: no rotation support", supportsRotation());
         assumeTrue("Skipping test: no multi-window support", supportsSplitScreenMultiWindow());
 
-        requestOrientationInSplitScreen(createManagedRotationSession(),
-                ROTATION_90 /* portrait */, LANDSCAPE_ORIENTATION_ACTIVITY);
+        try (final RotationSession rotationSession = new RotationSession()) {
+            requestOrientationInSplitScreen(rotationSession,
+                    ROTATION_90 /* portrait */, LANDSCAPE_ORIENTATION_ACTIVITY);
+        }
     }
 
     /**
@@ -687,8 +696,10 @@
     public void testSplitscreenLandscapeAppOrientationRequests() throws Exception {
         assumeTrue("Skipping test: no multi-window support", supportsSplitScreenMultiWindow());
 
-        requestOrientationInSplitScreen(createManagedRotationSession(),
-                ROTATION_0 /* landscape */, PORTRAIT_ORIENTATION_ACTIVITY);
+        try (final RotationSession rotationSession = new RotationSession()) {
+            requestOrientationInSplitScreen(rotationSession,
+                    ROTATION_0 /* landscape */, PORTRAIT_ORIENTATION_ACTIVITY);
+        }
     }
 
     /**
diff --git a/tests/framework/base/windowmanager/src/android/server/wm/ConfigChangeTests.java b/tests/framework/base/windowmanager/src/android/server/wm/ConfigChangeTests.java
index b1a31df..8fc36e8 100644
--- a/tests/framework/base/windowmanager/src/android/server/wm/ConfigChangeTests.java
+++ b/tests/framework/base/windowmanager/src/android/server/wm/ConfigChangeTests.java
@@ -82,7 +82,7 @@
     private @interface TestMode {}
 
     @Test
-    public void testRotation90Relaunch() {
+    public void testRotation90Relaunch() throws Exception{
         assumeTrue("Skipping test: no rotation support", supportsRotation());
 
         // Should relaunch on every rotation and receive no onConfigurationChanged()
@@ -90,7 +90,7 @@
     }
 
     @Test
-    public void testRotation90NoRelaunch() {
+    public void testRotation90NoRelaunch() throws Exception {
         assumeTrue("Skipping test: no rotation support", supportsRotation());
 
         // Should receive onConfigurationChanged() on every rotation and no relaunch
@@ -98,7 +98,7 @@
     }
 
     @Test
-    public void testRotation180_RegularActivity() {
+    public void testRotation180_RegularActivity() throws Exception {
         assumeTrue("Skipping test: no rotation support", supportsRotation());
         assumeFalse("Skipping test: display cutout present, can't predict exact lifecycle",
                 hasDisplayCutout());
@@ -108,7 +108,7 @@
     }
 
     @Test
-    public void testRotation180_NoRelaunchActivity() {
+    public void testRotation180_NoRelaunchActivity() throws Exception {
         assumeTrue("Skipping test: no rotation support", supportsRotation());
         assumeFalse("Skipping test: display cutout present, can't predict exact lifecycle",
                 hasDisplayCutout());
@@ -150,39 +150,40 @@
         launchActivity(activityName);
         mAmWmState.computeState(activityName);
 
-        final RotationSession rotationSession = createManagedRotationSession();
-        final StateCount count1 = getStateCountForRotation(activityName, rotationSession,
-                ROTATION_0 /* before */, ROTATION_180 /* after */);
-        final StateCount count2 = getStateCountForRotation(activityName, rotationSession,
-                ROTATION_90 /* before */, ROTATION_270 /* after */);
+        try(final RotationSession rotationSession = new RotationSession()) {
+            final StateCount count1 = getStateCountForRotation(activityName, rotationSession,
+                    ROTATION_0 /* before */, ROTATION_180 /* after */);
+            final StateCount count2 = getStateCountForRotation(activityName, rotationSession,
+                    ROTATION_90 /* before */, ROTATION_270 /* after */);
 
-        final int configChange = count1.mConfigChangeCount + count2.mConfigChangeCount;
-        final int relaunch = count1.mRelaunchCount + count2.mRelaunchCount;
-        // There should at least one 180 rotation without resize.
-        final boolean sameSize = !count1.mResize || !count2.mResize;
+            final int configChange = count1.mConfigChangeCount + count2.mConfigChangeCount;
+            final int relaunch = count1.mRelaunchCount + count2.mRelaunchCount;
+            // There should at least one 180 rotation without resize.
+            final boolean sameSize = !count1.mResize || !count2.mResize;
 
-        switch (testMode) {
-            case TEST_MODE_CONFIGURATION_CHANGE: {
-                assertTrue("There must be at most one 180 degree rotation that results in the"
-                        + " same configuration on device with cutout", configChange <= 1);
-                assertEquals("There must be no relaunch during test", 0, relaunch);
-                break;
-            }
-            case TEST_MODE_RELAUNCH_OR_CONFIG_CHANGE: {
-                // If the size change does not cross the threshold, the activity will receive
-                // onConfigurationChanged instead of relaunching.
-                assertTrue("There must be at most one 180 degree rotation that results in"
-                        + " relaunch or a configuration change on device with cutout",
-                        relaunch + configChange <= 1);
-                break;
-            }
-            case TEST_MODE_RESIZE: {
-                assertTrue("A device with cutout should have the same available screen space"
-                        + " in landscape and reverse-landscape", sameSize);
-                break;
-            }
-            default: {
-                fail("unrecognized test mode: " + testMode);
+            switch(testMode) {
+                case TEST_MODE_CONFIGURATION_CHANGE: {
+                    assertTrue("There must be at most one 180 degree rotation that results in the"
+                            + " same configuration on device with cutout", configChange <= 1);
+                    assertEquals("There must be no relaunch during test", 0, relaunch);
+                    break;
+                }
+                case TEST_MODE_RELAUNCH_OR_CONFIG_CHANGE: {
+                    // If the size change does not cross the threshold, the activity will receive
+                    // onConfigurationChanged instead of relaunching.
+                    assertTrue("There must be at most one 180 degree rotation that results in"
+                            + " relaunch or a configuration change on device with cutout",
+                            relaunch + configChange <= 1);
+                    break;
+                }
+                case TEST_MODE_RESIZE: {
+                    assertTrue("A device with cutout should have the same available screen space"
+                            + " in landscape and reverse-landscape", sameSize);
+                    break;
+                }
+                default: {
+                    fail("unrecognized test mode: " + testMode);
+                }
             }
         }
     }
@@ -215,45 +216,46 @@
     }
 
     @Test
-    public void testChangeFontScaleRelaunch() {
+    public void testChangeFontScaleRelaunch() throws Exception {
         // Should relaunch and receive no onConfigurationChanged()
         testChangeFontScale(FONT_SCALE_ACTIVITY, true /* relaunch */);
     }
 
     @Test
-    public void testChangeFontScaleNoRelaunch() {
+    public void testChangeFontScaleNoRelaunch() throws Exception {
         // Should receive onConfigurationChanged() and no relaunch
         testChangeFontScale(FONT_SCALE_NO_RELAUNCH_ACTIVITY, false /* relaunch */);
     }
 
     private void testRotation(ComponentName activityName, int rotationStep, int numRelaunch,
-            int numConfigChange) {
+            int numConfigChange) throws Exception {
         launchActivity(activityName);
 
         mAmWmState.computeState(activityName);
 
         final int initialRotation = 4 - rotationStep;
-        final RotationSession rotationSession = createManagedRotationSession();
-        rotationSession.set(initialRotation);
-        mAmWmState.computeState(activityName);
-        final int actualStackId =
-                mAmWmState.getAmState().getTaskByActivity(activityName).mStackId;
-        final int displayId = mAmWmState.getAmState().getStackById(actualStackId).mDisplayId;
-        final int newDeviceRotation = getDeviceRotation(displayId);
-        if (newDeviceRotation == INVALID_DEVICE_ROTATION) {
-            logE("Got an invalid device rotation value. "
-                    + "Continuing the test despite of that, but it is likely to fail.");
-        } else if (newDeviceRotation != initialRotation) {
-            log("This device doesn't support user rotation "
-                    + "mode. Not continuing the rotation checks.");
-            return;
-        }
-
-        for (int rotation = 0; rotation < 4; rotation += rotationStep) {
-            separateTestJournal();
-            rotationSession.set(rotation);
+        try (final RotationSession rotationSession = new RotationSession()) {
+            rotationSession.set(initialRotation);
             mAmWmState.computeState(activityName);
-            assertRelaunchOrConfigChanged(activityName, numRelaunch, numConfigChange);
+            final int actualStackId =
+                    mAmWmState.getAmState().getTaskByActivity(activityName).mStackId;
+            final int displayId = mAmWmState.getAmState().getStackById(actualStackId).mDisplayId;
+            final int newDeviceRotation = getDeviceRotation(displayId);
+            if (newDeviceRotation == INVALID_DEVICE_ROTATION) {
+                logE("Got an invalid device rotation value. "
+                        + "Continuing the test despite of that, but it is likely to fail.");
+            } else if (newDeviceRotation != initialRotation) {
+                log("This device doesn't support user rotation "
+                        + "mode. Not continuing the rotation checks.");
+                return;
+            }
+
+            for (int rotation = 0; rotation < 4; rotation += rotationStep) {
+                separateTestJournal();
+                rotationSession.set(rotation);
+                mAmWmState.computeState(activityName);
+                assertRelaunchOrConfigChanged(activityName, numRelaunch, numConfigChange);
+            }
         }
     }
 
@@ -266,34 +268,27 @@
         }
     }
 
-    private void testChangeFontScale(ComponentName activityName, boolean relaunch) {
-        final FontScaleSession fontScaleSession = mObjectTracker.manage(new FontScaleSession());
-        fontScaleSession.set(1.0f);
-        separateTestJournal();
-        launchActivity(activityName);
-        mAmWmState.computeState(activityName);
-
-        final Bundle extras = TestJournalContainer.get(activityName).extras;
-        if (!extras.containsKey(EXTRA_FONT_ACTIVITY_DPI)) {
-            fail("No fontActivityDpi reported from activity " + activityName);
-        }
-        final int densityDpi = extras.getInt(EXTRA_FONT_ACTIVITY_DPI);
-
-        for (float fontScale = 0.85f; fontScale <= 1.3f; fontScale += 0.15f) {
+    private void testChangeFontScale(
+            ComponentName activityName, boolean relaunch) throws Exception {
+        try (final FontScaleSession fontScaleSession = new FontScaleSession()) {
+            fontScaleSession.set(1.0f);
             separateTestJournal();
-            fontScaleSession.set(fontScale);
+            launchActivity(activityName);
             mAmWmState.computeState(activityName);
-            assertRelaunchOrConfigChanged(activityName, relaunch ? 1 : 0, relaunch ? 0 : 1);
 
-            // Verify that the display metrics are updated, and therefore the text size is also
-            // updated accordingly.
-            final Bundle changedExtras = TestJournalContainer.get(activityName).extras;
-            waitForOrFail("reported fontPixelSize from " + activityName,
-                    () -> changedExtras.containsKey(EXTRA_FONT_PIXEL_SIZE));
-            final int expectedFontPixelSize =
-                    scaledPixelsToPixels(EXPECTED_FONT_SIZE_SP, fontScale, densityDpi);
-            assertEquals("Expected font pixel size should match", expectedFontPixelSize,
-                    changedExtras.getInt(EXTRA_FONT_PIXEL_SIZE));
+            final int densityDpi = getActivityDensityDpi(activityName);
+
+            for (float fontScale = 0.85f; fontScale <= 1.3f; fontScale += 0.15f) {
+                separateTestJournal();
+                fontScaleSession.set(fontScale);
+                mAmWmState.computeState(activityName);
+                assertRelaunchOrConfigChanged(activityName, relaunch ? 1 : 0, relaunch ? 0 : 1);
+
+                // Verify that the display metrics are updated, and therefore the text size is also
+                // updated accordingly.
+                assertExpectedFontPixelSize(activityName,
+                        scaledPixelsToPixels(EXPECTED_FONT_SIZE_SP, fontScale, densityDpi));
+            }
         }
     }
 
@@ -341,6 +336,26 @@
         return (int) ((f >= 0) ? (f + 0.5f) : (f - 0.5f));
     }
 
+    private static int getActivityDensityDpi(ComponentName activityName)
+            throws Exception {
+        final Bundle extras = TestJournalContainer.get(activityName).extras;
+        if (!extras.containsKey(EXTRA_FONT_ACTIVITY_DPI)) {
+            fail("No fontActivityDpi reported from activity " + activityName);
+            return -1;
+        }
+        return extras.getInt(EXTRA_FONT_ACTIVITY_DPI);
+    }
+
+    private void assertExpectedFontPixelSize(ComponentName activityName, int fontPixelSize)
+            throws Exception {
+        final Bundle extras = TestJournalContainer.get(activityName).extras;
+        if (!extras.containsKey(EXTRA_FONT_PIXEL_SIZE)) {
+            fail("No fontPixelSize reported from activity " + activityName);
+        }
+        assertEquals("Expected font pixel size does not match", fontPixelSize,
+                extras.getInt(EXTRA_FONT_PIXEL_SIZE));
+    }
+
     private void updateApplicationInfo(List<String> packages) {
         SystemUtil.runWithShellPermissionIdentity(
                 () -> mAm.scheduleApplicationInfoChanged(packages,
diff --git a/tests/framework/base/windowmanager/src/android/server/wm/CrossAppDragAndDropTests.java b/tests/framework/base/windowmanager/src/android/server/wm/CrossAppDragAndDropTests.java
index 67cabaa..a3a6503 100644
--- a/tests/framework/base/windowmanager/src/android/server/wm/CrossAppDragAndDropTests.java
+++ b/tests/framework/base/windowmanager/src/android/server/wm/CrossAppDragAndDropTests.java
@@ -117,7 +117,9 @@
     }
 
     @After
-    public void tearDown() {
+    @Override
+    public void tearDown() throws Exception {
+        super.tearDown();
         cleanupState();
     }
 
@@ -125,7 +127,7 @@
      * Make sure that the special activity stacks are removed and the ActivityManager/WindowManager
      * is in a good state.
      */
-    private void cleanupState() {
+    private void cleanupState() throws Exception {
         stopTestPackage(DRAG_SOURCE.getPackageName());
         stopTestPackage(DROP_TARGET.getPackageName());
         stopTestPackage(DROP_TARGET_SDK23.getPackageName());
diff --git a/tests/framework/base/windowmanager/src/android/server/wm/DeprecatedTargetSdkTest.java b/tests/framework/base/windowmanager/src/android/server/wm/DeprecatedTargetSdkTest.java
index 4f37a3c..8690b73 100644
--- a/tests/framework/base/windowmanager/src/android/server/wm/DeprecatedTargetSdkTest.java
+++ b/tests/framework/base/windowmanager/src/android/server/wm/DeprecatedTargetSdkTest.java
@@ -38,7 +38,10 @@
             "DeprecatedTargetSdkVersionDialog";
 
     @After
-    public void tearDown() {
+    @Override
+    public void tearDown() throws Exception {
+        super.tearDown();
+
         // Ensure app process is stopped.
         stopTestPackage(MAIN_ACTIVITY.getPackageName());
     }
diff --git a/tests/framework/base/windowmanager/src/android/server/wm/DisplaySizeTest.java b/tests/framework/base/windowmanager/src/android/server/wm/DisplaySizeTest.java
index e5b9a04..e9baa7a 100644
--- a/tests/framework/base/windowmanager/src/android/server/wm/DisplaySizeTest.java
+++ b/tests/framework/base/windowmanager/src/android/server/wm/DisplaySizeTest.java
@@ -54,38 +54,45 @@
             "UnsupportedDisplaySizeDialog";
 
     @After
-    public void tearDown() {
+    @Override
+    public void tearDown() throws Exception {
+        super.tearDown();
+
         // Ensure app process is stopped.
         stopTestPackage(SMALLEST_WIDTH_ACTIVITY.getPackageName());
         stopTestPackage(TEST_ACTIVITY.getPackageName());
     }
 
     @Test
-    public void testCompatibilityDialog() {
+    public void testCompatibilityDialog() throws Exception {
         // Launch some other app (not to perform density change on launcher).
         launchActivity(TEST_ACTIVITY);
         mAmWmState.assertActivityDisplayed(TEST_ACTIVITY);
 
-        createManagedScreenDensitySession().setUnsupportedDensity();
+        try (final ScreenDensitySession screenDensitySession = new ScreenDensitySession()) {
+            screenDensitySession.setUnsupportedDensity();
 
-        // Launch target app.
-        launchActivity(SMALLEST_WIDTH_ACTIVITY);
-        mAmWmState.assertActivityDisplayed(SMALLEST_WIDTH_ACTIVITY);
-        mAmWmState.assertWindowDisplayed(UNSUPPORTED_DISPLAY_SIZE_DIALOG_NAME);
+            // Launch target app.
+            launchActivity(SMALLEST_WIDTH_ACTIVITY);
+            mAmWmState.assertActivityDisplayed(SMALLEST_WIDTH_ACTIVITY);
+            mAmWmState.assertWindowDisplayed(UNSUPPORTED_DISPLAY_SIZE_DIALOG_NAME);
+        }
     }
 
     @Test
-    public void testCompatibilityDialogWhenFocused() {
+    public void testCompatibilityDialogWhenFocused() throws Exception {
         launchActivity(SMALLEST_WIDTH_ACTIVITY);
         mAmWmState.assertActivityDisplayed(SMALLEST_WIDTH_ACTIVITY);
 
-        createManagedScreenDensitySession().setUnsupportedDensity();
+        try (final ScreenDensitySession screenDensitySession = new ScreenDensitySession()) {
+            screenDensitySession.setUnsupportedDensity();
 
-        mAmWmState.assertWindowDisplayed(UNSUPPORTED_DISPLAY_SIZE_DIALOG_NAME);
+            mAmWmState.assertWindowDisplayed(UNSUPPORTED_DISPLAY_SIZE_DIALOG_NAME);
+        }
     }
 
     @Test
-    public void testCompatibilityDialogAfterReturn() {
+    public void testCompatibilityDialogAfterReturn() throws Exception {
         // Launch target app.
         launchActivity(SMALLEST_WIDTH_ACTIVITY);
         mAmWmState.assertActivityDisplayed(SMALLEST_WIDTH_ACTIVITY);
@@ -97,13 +104,15 @@
         mAmWmState.assertActivityDisplayed(TEST_ACTIVITY);
         separateTestJournal();
 
-        createManagedScreenDensitySession().setUnsupportedDensity();
+        try (final ScreenDensitySession screenDensitySession = new ScreenDensitySession()) {
+            screenDensitySession.setUnsupportedDensity();
 
-        assertActivityLifecycle(TEST_ACTIVITY, true /* relaunched */);
-        pressBackButton();
+            assertActivityLifecycle(TEST_ACTIVITY, true /* relaunched */);
+            pressBackButton();
 
-        mAmWmState.assertActivityDisplayed(SMALLEST_WIDTH_ACTIVITY);
-        mAmWmState.assertWindowDisplayed(UNSUPPORTED_DISPLAY_SIZE_DIALOG_NAME);
+            mAmWmState.assertActivityDisplayed(SMALLEST_WIDTH_ACTIVITY);
+            mAmWmState.assertWindowDisplayed(UNSUPPORTED_DISPLAY_SIZE_DIALOG_NAME);
+        }
     }
 
     @Test
@@ -158,10 +167,6 @@
         }
     }
 
-    protected ScreenDensitySession createManagedScreenDensitySession() {
-        return mObjectTracker.manage(new ScreenDensitySession());
-    }
-
     private static class ScreenDensitySession implements AutoCloseable {
         private static final String DENSITY_PROP_DEVICE = "ro.sf.lcd_density";
         private static final String DENSITY_PROP_EMULATOR = "qemu.sf.lcd_density";
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 01b48b5..adce698 100644
--- a/tests/framework/base/windowmanager/src/android/server/wm/DisplayTests.java
+++ b/tests/framework/base/windowmanager/src/android/server/wm/DisplayTests.java
@@ -63,11 +63,13 @@
      */
     @Test
     public void testCreateVirtualDisplayWithCustomConfig() throws Exception {
-        final ActivityDisplay newDisplay = createManagedVirtualDisplaySession().createDisplay();
+        try (final VirtualDisplaySession virtualDisplaySession = new VirtualDisplaySession()) {
+            final ActivityDisplay newDisplay = virtualDisplaySession.createDisplay();
 
-        // Find the density of created display.
-        final int newDensityDpi = newDisplay.mFullConfiguration.densityDpi;
-        assertEquals(CUSTOM_DENSITY_DPI, newDensityDpi);
+            // Find the density of created display.
+            final int newDensityDpi = newDisplay.mFullConfiguration.densityDpi;
+            assertEquals(CUSTOM_DENSITY_DPI, newDensityDpi);
+        }
     }
 
     @Test
@@ -98,25 +100,27 @@
         // Only check devices with the feature disabled.
         assumeFalse(supportsMultiDisplay());
 
-        // Create new virtual display.
-        final ActivityDisplay newDisplay = createManagedVirtualDisplaySession().createDisplay();
+        try (final VirtualDisplaySession virtualDisplaySession = new VirtualDisplaySession()) {
+            // Create new virtual display.
+            final ActivityDisplay newDisplay = virtualDisplaySession.createDisplay();
 
-        // Launch activity on new secondary display.
-        launchActivityOnDisplay(TEST_ACTIVITY, newDisplay.mId);
-        mAmWmState.computeState(TEST_ACTIVITY);
+            // Launch activity on new secondary display.
+            launchActivityOnDisplay(TEST_ACTIVITY, newDisplay.mId);
+            mAmWmState.computeState(TEST_ACTIVITY);
 
-        mAmWmState.assertFocusedActivity("Launched activity must be focused",
-                TEST_ACTIVITY);
+            mAmWmState.assertFocusedActivity("Launched activity must be focused",
+                    TEST_ACTIVITY);
 
-        // Check that activity is on the right display.
-        final int frontStackId = mAmWmState.getAmState().getFrontStackId(DEFAULT_DISPLAY);
-        final ActivityManagerState.ActivityStack frontStack =
-                mAmWmState.getAmState().getStackById(frontStackId);
-        assertEquals("Launched activity must be resumed",
-                getActivityName(TEST_ACTIVITY), frontStack.mResumedActivity);
-        assertEquals("Front stack must be on the default display",
-                DEFAULT_DISPLAY, frontStack.mDisplayId);
-        mAmWmState.assertFocusedStack("Focus must be on the default display", frontStackId);
+            // Check that activity is on the right display.
+            final int frontStackId = mAmWmState.getAmState().getFrontStackId(DEFAULT_DISPLAY);
+            final ActivityManagerState.ActivityStack frontStack =
+                    mAmWmState.getAmState().getStackById(frontStackId);
+            assertEquals("Launched activity must be resumed",
+                    getActivityName(TEST_ACTIVITY), frontStack.mResumedActivity);
+            assertEquals("Front stack must be on the default display",
+                    DEFAULT_DISPLAY, frontStack.mDisplayId);
+            mAmWmState.assertFocusedStack("Focus must be on the default display", frontStackId);
+        }
     }
 
     @Test
@@ -136,40 +140,40 @@
      * and unlocking the phone and verifies that overrides are kept.
      */
     @Test
-    public void testForceDisplayMetrics() {
+    public void testForceDisplayMetrics() throws Exception {
         launchHomeActivity();
 
-        final DisplayMetricsSession displayMetricsSession =
-                mObjectTracker.manage(new DisplayMetricsSession(DEFAULT_DISPLAY));
-        final LockScreenSession lockScreenSession = createManagedLockScreenSession();
+        try (final DisplayMetricsSession displayMetricsSession =
+                     new DisplayMetricsSession(DEFAULT_DISPLAY);
+             final LockScreenSession lockScreenSession = new LockScreenSession()) {
+            // Read initial sizes.
+            final ReportedDisplayMetrics originalDisplayMetrics =
+                    displayMetricsSession.getInitialDisplayMetrics();
 
-        // Read initial sizes.
-        final ReportedDisplayMetrics originalDisplayMetrics =
-                displayMetricsSession.getInitialDisplayMetrics();
+            // Apply new override values that don't match the physical metrics.
+            final Size overrideSize = new Size(
+                    (int) (originalDisplayMetrics.physicalSize.getWidth() * 1.5),
+                    (int) (originalDisplayMetrics.physicalSize.getHeight() * 1.5));
+            final Integer overrideDensity = (int) (originalDisplayMetrics.physicalDensity * 1.1);
+            displayMetricsSession.overrideDisplayMetrics(overrideSize, overrideDensity);
 
-        // Apply new override values that don't match the physical metrics.
-        final Size overrideSize = new Size(
-                (int) (originalDisplayMetrics.physicalSize.getWidth() * 1.5),
-                (int) (originalDisplayMetrics.physicalSize.getHeight() * 1.5));
-        final Integer overrideDensity = (int) (originalDisplayMetrics.physicalDensity * 1.1);
-        displayMetricsSession.overrideDisplayMetrics(overrideSize, overrideDensity);
+            // Check if overrides applied correctly.
+            ReportedDisplayMetrics displayMetrics = displayMetricsSession.getDisplayMetrics();
+            assertEquals(overrideSize, displayMetrics.overrideSize);
+            assertEquals(overrideDensity, displayMetrics.overrideDensity);
 
-        // Check if overrides applied correctly.
-        ReportedDisplayMetrics displayMetrics = displayMetricsSession.getDisplayMetrics();
-        assertEquals(overrideSize, displayMetrics.overrideSize);
-        assertEquals(overrideDensity, displayMetrics.overrideDensity);
+            // Lock and unlock device. This will cause a DISPLAY_CHANGED event to be triggered and
+            // might update the metrics.
+            lockScreenSession.sleepDevice()
+                    .wakeUpDevice()
+                    .unlockDevice();
+            mAmWmState.waitForHomeActivityVisible();
 
-        // Lock and unlock device. This will cause a DISPLAY_CHANGED event to be triggered and
-        // might update the metrics.
-        lockScreenSession.sleepDevice()
-                .wakeUpDevice()
-                .unlockDevice();
-        mAmWmState.waitForHomeActivityVisible();
-
-        // Check if overrides are still applied.
-        displayMetrics = displayMetricsSession.getDisplayMetrics();
-        assertEquals(overrideSize, displayMetrics.overrideSize);
-        assertEquals(overrideDensity, displayMetrics.overrideDensity);
+            // Check if overrides are still applied.
+            displayMetrics = displayMetricsSession.getDisplayMetrics();
+            assertEquals(overrideSize, displayMetrics.overrideSize);
+            assertEquals(overrideDensity, displayMetrics.overrideDensity);
+        }
     }
 
     private Configuration getDisplayResourcesConfiguration(int displayWidth, int displayHeight)
@@ -195,7 +199,7 @@
         private final ReportedDisplayMetrics mInitialDisplayMetrics;
         private final int mDisplayId;
 
-        DisplayMetricsSession(int displayId) {
+        DisplayMetricsSession(int displayId) throws Exception {
             mDisplayId = displayId;
             mInitialDisplayMetrics = ReportedDisplayMetrics.getDisplayMetrics(mDisplayId);
         }
@@ -204,7 +208,7 @@
             return mInitialDisplayMetrics;
         }
 
-        ReportedDisplayMetrics getDisplayMetrics() {
+        ReportedDisplayMetrics getDisplayMetrics() throws Exception {
             return ReportedDisplayMetrics.getDisplayMetrics(mDisplayId);
         }
 
@@ -213,7 +217,7 @@
         }
 
         @Override
-        public void close() {
+        public void close() throws Exception {
             mInitialDisplayMetrics.restoreDisplayMetrics();
         }
     }
diff --git a/tests/framework/base/windowmanager/src/android/server/wm/FreeformWindowingModeTests.java b/tests/framework/base/windowmanager/src/android/server/wm/FreeformWindowingModeTests.java
index 7d380f2..d2d73c5 100644
--- a/tests/framework/base/windowmanager/src/android/server/wm/FreeformWindowingModeTests.java
+++ b/tests/framework/base/windowmanager/src/android/server/wm/FreeformWindowingModeTests.java
@@ -52,32 +52,35 @@
     // with bounds (0, 0, 900, 900)
 
     @Test
-    public void testFreeformWindowManagementSupport() {
-        int displayId = Display.DEFAULT_DISPLAY;
-        if (supportsMultiDisplay()) {
-            displayId = createManagedVirtualDisplaySession()
-                    .setSimulateDisplay(true)
-                    .setSimulationDisplaySize(1920 /* width */, 1080 /* height */)
-                    .createDisplay().mId;
+    public void testFreeformWindowManagementSupport() throws Exception {
+        try (VirtualDisplaySession virtualDisplaySession = new VirtualDisplaySession()) {
+            int displayId = Display.DEFAULT_DISPLAY;
+            if (supportsMultiDisplay()) {
+                final ActivityManagerState.ActivityDisplay display = virtualDisplaySession
+                        .setSimulateDisplay(true)
+                        .setSimulationDisplaySize(1920 /* width */, 1080 /* height */)
+                        .createDisplay();
+                displayId = display.mId;
+            }
+            launchActivityOnDisplay(FREEFORM_ACTIVITY, WINDOWING_MODE_FREEFORM, displayId);
+
+            mAmWmState.computeState(FREEFORM_ACTIVITY, TEST_ACTIVITY);
+
+            if (!supportsFreeform()) {
+                mAmWmState.assertDoesNotContainStack("Must not contain freeform stack.",
+                        WINDOWING_MODE_FREEFORM, ACTIVITY_TYPE_STANDARD);
+                return;
+            }
+
+            mAmWmState.assertFrontStackOnDisplay("Freeform stack must be the front stack.",
+                    WINDOWING_MODE_FREEFORM, ACTIVITY_TYPE_STANDARD, displayId);
+            mAmWmState.assertVisibility(FREEFORM_ACTIVITY, true);
+            mAmWmState.assertVisibility(TEST_ACTIVITY, true);
+            mAmWmState.assertFocusedActivity(
+                    TEST_ACTIVITY + " must be focused Activity", TEST_ACTIVITY);
+            assertEquals(new Rect(0, 0, TEST_TASK_SIZE_1, TEST_TASK_SIZE_1),
+                    mAmWmState.getAmState().getTaskByActivity(TEST_ACTIVITY).getBounds());
         }
-        launchActivityOnDisplay(FREEFORM_ACTIVITY, WINDOWING_MODE_FREEFORM, displayId);
-
-        mAmWmState.computeState(FREEFORM_ACTIVITY, TEST_ACTIVITY);
-
-        if (!supportsFreeform()) {
-            mAmWmState.assertDoesNotContainStack("Must not contain freeform stack.",
-                    WINDOWING_MODE_FREEFORM, ACTIVITY_TYPE_STANDARD);
-            return;
-        }
-
-        mAmWmState.assertFrontStackOnDisplay("Freeform stack must be the front stack.",
-                WINDOWING_MODE_FREEFORM, ACTIVITY_TYPE_STANDARD, displayId);
-        mAmWmState.assertVisibility(FREEFORM_ACTIVITY, true);
-        mAmWmState.assertVisibility(TEST_ACTIVITY, true);
-        mAmWmState.assertFocusedActivity(
-                TEST_ACTIVITY + " must be focused Activity", TEST_ACTIVITY);
-        assertEquals(new Rect(0, 0, TEST_TASK_SIZE_1, TEST_TASK_SIZE_1),
-                mAmWmState.getAmState().getTaskByActivity(TEST_ACTIVITY).getBounds());
     }
 
     @Test
diff --git a/tests/framework/base/windowmanager/src/android/server/wm/KeyguardLockedTests.java b/tests/framework/base/windowmanager/src/android/server/wm/KeyguardLockedTests.java
index 5ab75e7..81510f8 100644
--- a/tests/framework/base/windowmanager/src/android/server/wm/KeyguardLockedTests.java
+++ b/tests/framework/base/windowmanager/src/android/server/wm/KeyguardLockedTests.java
@@ -18,7 +18,6 @@
 
 import static android.app.WindowConfiguration.ACTIVITY_TYPE_STANDARD;
 import static android.app.WindowConfiguration.WINDOWING_MODE_PINNED;
-import static android.server.wm.MockImeHelper.createManagedMockImeSession;
 import static android.server.wm.UiDeviceUtils.pressBackButton;
 import static android.server.wm.app.Components.DISMISS_KEYGUARD_ACTIVITY;
 import static android.server.wm.app.Components.DISMISS_KEYGUARD_METHOD_ACTIVITY;
@@ -32,6 +31,7 @@
 import static android.server.wm.app.Components.TURN_SCREEN_ON_ATTR_DISMISS_KEYGUARD_ACTIVITY;
 import static android.view.Display.DEFAULT_DISPLAY;
 import static android.view.WindowManager.LayoutParams.SOFT_INPUT_STATE_ALWAYS_VISIBLE;
+import static androidx.test.InstrumentationRegistry.getInstrumentation;
 import static com.android.cts.mockime.ImeEventStreamTestUtils.expectEvent;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertTrue;
@@ -46,8 +46,8 @@
 import android.widget.EditText;
 import android.widget.LinearLayout;
 import com.android.cts.mockime.ImeEventStream;
+import com.android.cts.mockime.ImeSettings;
 import com.android.cts.mockime.MockImeSession;
-
 import java.util.concurrent.TimeUnit;
 import org.junit.Before;
 import org.junit.Test;
@@ -67,22 +67,22 @@
     }
 
     @Test
-    public void testLockAndUnlock() {
-        final LockScreenSession lockScreenSession = createManagedLockScreenSession();
-        lockScreenSession.setLockCredential().gotoKeyguard();
-
-        assertTrue(mKeyguardManager.isKeyguardLocked());
-        assertTrue(mKeyguardManager.isDeviceLocked());
-        assertTrue(mKeyguardManager.isDeviceSecure());
-        assertTrue(mKeyguardManager.isKeyguardSecure());
-        mAmWmState.assertKeyguardShowingAndNotOccluded();
-
-        lockScreenSession.unlockDevice().enterAndConfirmLockCredential();
-
-        mAmWmState.waitForKeyguardGone();
-        mAmWmState.assertKeyguardGone();
-        assertFalse(mKeyguardManager.isDeviceLocked());
-        assertFalse(mKeyguardManager.isKeyguardLocked());
+    public void testLockAndUnlock() throws Exception {
+        try (final LockScreenSession lockScreenSession = new LockScreenSession()) {
+            lockScreenSession.setLockCredential()
+                    .gotoKeyguard();
+            assertTrue(mKeyguardManager.isKeyguardLocked());
+            assertTrue(mKeyguardManager.isDeviceLocked());
+            assertTrue(mKeyguardManager.isDeviceSecure());
+            assertTrue(mKeyguardManager.isKeyguardSecure());
+            mAmWmState.assertKeyguardShowingAndNotOccluded();
+            lockScreenSession.unlockDevice()
+                    .enterAndConfirmLockCredential();
+            mAmWmState.waitForKeyguardGone();
+            mAmWmState.assertKeyguardGone();
+            assertFalse(mKeyguardManager.isDeviceLocked());
+            assertFalse(mKeyguardManager.isKeyguardLocked());
+        }
     }
 
     @Test
@@ -103,246 +103,255 @@
     }
 
     @Test
-    public void testDismissKeyguard() {
-        final LockScreenSession lockScreenSession = createManagedLockScreenSession();
-        lockScreenSession.setLockCredential().gotoKeyguard();
-
-        mAmWmState.assertKeyguardShowingAndNotOccluded();
-        launchActivity(DISMISS_KEYGUARD_ACTIVITY);
-        lockScreenSession.enterAndConfirmLockCredential();
-
-        mAmWmState.waitForKeyguardGone();
-        mAmWmState.assertKeyguardGone();
-        mAmWmState.assertVisibility(DISMISS_KEYGUARD_ACTIVITY, true);
+    public void testDismissKeyguard() throws Exception {
+        try (final LockScreenSession lockScreenSession = new LockScreenSession()) {
+            lockScreenSession.setLockCredential()
+                    .gotoKeyguard();
+            mAmWmState.assertKeyguardShowingAndNotOccluded();
+            launchActivity(DISMISS_KEYGUARD_ACTIVITY);
+            lockScreenSession.enterAndConfirmLockCredential();
+            mAmWmState.waitForKeyguardGone();
+            mAmWmState.assertKeyguardGone();
+            mAmWmState.assertVisibility(DISMISS_KEYGUARD_ACTIVITY, true);
+        }
     }
 
     @Test
-    public void testDismissKeyguard_whileOccluded() {
-        final LockScreenSession lockScreenSession = createManagedLockScreenSession();
-        lockScreenSession.setLockCredential().gotoKeyguard();
-
-        mAmWmState.assertKeyguardShowingAndNotOccluded();
-        launchActivity(SHOW_WHEN_LOCKED_ACTIVITY);
-        mAmWmState.computeState(SHOW_WHEN_LOCKED_ACTIVITY);
-        mAmWmState.assertVisibility(SHOW_WHEN_LOCKED_ACTIVITY, true);
-
-        launchActivity(DISMISS_KEYGUARD_ACTIVITY);
-        lockScreenSession.enterAndConfirmLockCredential();
-        mAmWmState.waitForKeyguardGone();
-        mAmWmState.assertKeyguardGone();
-        mAmWmState.computeState(DISMISS_KEYGUARD_ACTIVITY);
-
-        final boolean isDismissTranslucent = mAmWmState.getAmState()
-                .isActivityTranslucent(DISMISS_KEYGUARD_ACTIVITY);
-        mAmWmState.assertVisibility(DISMISS_KEYGUARD_ACTIVITY, true);
-        mAmWmState.assertVisibility(SHOW_WHEN_LOCKED_ACTIVITY, isDismissTranslucent);
+    public void testDismissKeyguard_whileOccluded() throws Exception {
+        try (final LockScreenSession lockScreenSession = new LockScreenSession()) {
+            lockScreenSession.setLockCredential()
+                    .gotoKeyguard();
+            mAmWmState.assertKeyguardShowingAndNotOccluded();
+            launchActivity(SHOW_WHEN_LOCKED_ACTIVITY);
+            mAmWmState.computeState(SHOW_WHEN_LOCKED_ACTIVITY);
+            mAmWmState.assertVisibility(SHOW_WHEN_LOCKED_ACTIVITY, true);
+            launchActivity(DISMISS_KEYGUARD_ACTIVITY);
+            lockScreenSession.enterAndConfirmLockCredential();
+            mAmWmState.waitForKeyguardGone();
+            mAmWmState.assertKeyguardGone();
+            mAmWmState.computeState(DISMISS_KEYGUARD_ACTIVITY);
+            boolean isDismissTranslucent =
+                    mAmWmState.getAmState().isActivityTranslucent(DISMISS_KEYGUARD_ACTIVITY);
+            mAmWmState.assertVisibility(DISMISS_KEYGUARD_ACTIVITY, true);
+            mAmWmState.assertVisibility(SHOW_WHEN_LOCKED_ACTIVITY, isDismissTranslucent);
+        }
     }
 
     @Test
-    public void testDismissKeyguard_fromShowWhenLocked_notAllowed() {
-        final LockScreenSession lockScreenSession = createManagedLockScreenSession();
-        lockScreenSession.setLockCredential().gotoKeyguard();
+    public void testDismissKeyguard_fromShowWhenLocked_notAllowed() throws Exception {
+        try (final LockScreenSession lockScreenSession = new LockScreenSession()) {
+            lockScreenSession.setLockCredential()
+                    .gotoKeyguard();
+            mAmWmState.assertKeyguardShowingAndNotOccluded();
+            launchActivity(SHOW_WHEN_LOCKED_ACTIVITY);
+            mAmWmState.computeState(SHOW_WHEN_LOCKED_ACTIVITY);
+            mAmWmState.assertVisibility(SHOW_WHEN_LOCKED_ACTIVITY, true);
+            mBroadcastActionTrigger.dismissKeyguardByFlag();
+            lockScreenSession.enterAndConfirmLockCredential();
 
-        mAmWmState.assertKeyguardShowingAndNotOccluded();
-        launchActivity(SHOW_WHEN_LOCKED_ACTIVITY);
-        mAmWmState.computeState(SHOW_WHEN_LOCKED_ACTIVITY);
-        mAmWmState.assertVisibility(SHOW_WHEN_LOCKED_ACTIVITY, true);
-        mBroadcastActionTrigger.dismissKeyguardByFlag();
-        lockScreenSession.enterAndConfirmLockCredential();
-
-        // Make sure we stay on Keyguard.
-        mAmWmState.assertKeyguardShowingAndOccluded();
-        mAmWmState.assertVisibility(SHOW_WHEN_LOCKED_ACTIVITY, true);
+            // Make sure we stay on Keyguard.
+            mAmWmState.assertKeyguardShowingAndOccluded();
+            mAmWmState.assertVisibility(SHOW_WHEN_LOCKED_ACTIVITY, true);
+        }
     }
 
     @Test
-    public void testDismissKeyguardActivity_method() {
-        final LockScreenSession lockScreenSession = createManagedLockScreenSession();
-        lockScreenSession.setLockCredential();
-        separateTestJournal();
-
-        lockScreenSession.gotoKeyguard();
-        mAmWmState.computeState(true);
-        assertTrue(mAmWmState.getAmState().getKeyguardControllerState().keyguardShowing);
-
-        launchActivity(DISMISS_KEYGUARD_METHOD_ACTIVITY);
-        lockScreenSession.enterAndConfirmLockCredential();
-        mAmWmState.waitForKeyguardGone();
-        mAmWmState.computeState(DISMISS_KEYGUARD_METHOD_ACTIVITY);
-        mAmWmState.assertVisibility(DISMISS_KEYGUARD_METHOD_ACTIVITY, true);
-        assertFalse(mAmWmState.getAmState().getKeyguardControllerState().keyguardShowing);
-        assertOnDismissSucceeded(DISMISS_KEYGUARD_METHOD_ACTIVITY);
+    public void testDismissKeyguardActivity_method() throws Exception {
+        try (final LockScreenSession lockScreenSession = new LockScreenSession()) {
+            lockScreenSession.setLockCredential();
+            separateTestJournal();
+            lockScreenSession.gotoKeyguard();
+            mAmWmState.computeState(true);
+            assertTrue(mAmWmState.getAmState().getKeyguardControllerState().keyguardShowing);
+            launchActivity(DISMISS_KEYGUARD_METHOD_ACTIVITY);
+            lockScreenSession.enterAndConfirmLockCredential();
+            mAmWmState.waitForKeyguardGone();
+            mAmWmState.computeState(DISMISS_KEYGUARD_METHOD_ACTIVITY);
+            mAmWmState.assertVisibility(DISMISS_KEYGUARD_METHOD_ACTIVITY, true);
+            assertFalse(mAmWmState.getAmState().getKeyguardControllerState().keyguardShowing);
+            assertOnDismissSucceeded(DISMISS_KEYGUARD_METHOD_ACTIVITY);
+        }
     }
 
     @Test
-    public void testDismissKeyguardActivity_method_cancelled() {
-        final LockScreenSession lockScreenSession = createManagedLockScreenSession();
-        lockScreenSession.setLockCredential();
-        separateTestJournal();
-
-        lockScreenSession.gotoKeyguard();
-        mAmWmState.computeState(true);
-        assertTrue(mAmWmState.getAmState().getKeyguardControllerState().keyguardShowing);
-
-        launchActivity(DISMISS_KEYGUARD_METHOD_ACTIVITY);
-        pressBackButton();
-        assertOnDismissCancelled(DISMISS_KEYGUARD_METHOD_ACTIVITY);
-        mAmWmState.computeState(true);
-        mAmWmState.assertVisibility(DISMISS_KEYGUARD_METHOD_ACTIVITY, false);
-        assertTrue(mAmWmState.getAmState().getKeyguardControllerState().keyguardShowing);
+    public void testDismissKeyguardActivity_method_cancelled() throws Exception {
+        try (final LockScreenSession lockScreenSession = new LockScreenSession()) {
+            lockScreenSession.setLockCredential();
+            separateTestJournal();
+            lockScreenSession.gotoKeyguard();
+            mAmWmState.computeState(true);
+            assertTrue(mAmWmState.getAmState().getKeyguardControllerState().keyguardShowing);
+            launchActivity(DISMISS_KEYGUARD_METHOD_ACTIVITY);
+            pressBackButton();
+            assertOnDismissCancelled(DISMISS_KEYGUARD_METHOD_ACTIVITY);
+            mAmWmState.computeState(true);
+            mAmWmState.assertVisibility(DISMISS_KEYGUARD_METHOD_ACTIVITY, false);
+            assertTrue(mAmWmState.getAmState().getKeyguardControllerState().keyguardShowing);
+        }
     }
 
     @Test
-    public void testDismissKeyguardAttrActivity_method_turnScreenOn_withSecureKeyguard() {
-        final LockScreenSession lockScreenSession = createManagedLockScreenSession();
-        lockScreenSession.setLockCredential().sleepDevice();
-        mAmWmState.computeState(true);
-        assertTrue(mAmWmState.getAmState().getKeyguardControllerState().keyguardShowing);
+    public void testDismissKeyguardAttrActivity_method_turnScreenOn_withSecureKeyguard()
+            throws Exception {
+        try (final LockScreenSession lockScreenSession = new LockScreenSession()) {
+            lockScreenSession.setLockCredential().sleepDevice();
 
-        launchActivity(TURN_SCREEN_ON_ATTR_DISMISS_KEYGUARD_ACTIVITY);
-        mAmWmState.waitForKeyguardShowingAndNotOccluded();
-        mAmWmState.assertVisibility(TURN_SCREEN_ON_ATTR_DISMISS_KEYGUARD_ACTIVITY, false);
-        assertTrue(mAmWmState.getAmState().getKeyguardControllerState().keyguardShowing);
-        assertTrue(isDisplayOn(DEFAULT_DISPLAY));
+            mAmWmState.computeState(true);
+            assertTrue(mAmWmState.getAmState().getKeyguardControllerState().keyguardShowing);
+            launchActivity(TURN_SCREEN_ON_ATTR_DISMISS_KEYGUARD_ACTIVITY);
+            mAmWmState.waitForKeyguardShowingAndNotOccluded();
+            mAmWmState.assertVisibility(TURN_SCREEN_ON_ATTR_DISMISS_KEYGUARD_ACTIVITY, false);
+            assertTrue(mAmWmState.getAmState().getKeyguardControllerState().keyguardShowing);
+            assertTrue(isDisplayOn(DEFAULT_DISPLAY));
+        }
     }
 
     @Test
-    public void testEnterPipOverKeyguard() {
+    public void testEnterPipOverKeyguard() throws Exception {
         assumeTrue(supportsPip());
 
-        final LockScreenSession lockScreenSession = createManagedLockScreenSession();
-        lockScreenSession.setLockCredential();
+        try (final LockScreenSession lockScreenSession = new LockScreenSession()) {
+            lockScreenSession.setLockCredential();
 
-        // Show the PiP activity in fullscreen.
-        launchActivity(PIP_ACTIVITY, EXTRA_SHOW_OVER_KEYGUARD, "true");
+            // Show the PiP activity in fullscreen
+            launchActivity(PIP_ACTIVITY, EXTRA_SHOW_OVER_KEYGUARD, "true");
 
-        // Lock the screen and ensure that the PiP activity showing over the LockScreen.
-        lockScreenSession.gotoKeyguard(PIP_ACTIVITY);
-        mAmWmState.waitForKeyguardShowingAndOccluded();
-        mAmWmState.assertKeyguardShowingAndOccluded();
+            // Lock the screen and ensure that the PiP activity showing over the LockScreen.
+            lockScreenSession.gotoKeyguard(PIP_ACTIVITY);
+            mAmWmState.waitForKeyguardShowingAndOccluded();
+            mAmWmState.assertKeyguardShowingAndOccluded();
 
-        // Request that the PiP activity enter picture-in-picture mode (ensure it does not).
-        mBroadcastActionTrigger.doAction(ACTION_ENTER_PIP);
-        waitForEnterPip(PIP_ACTIVITY);
-        mAmWmState.assertDoesNotContainStack("Must not contain pinned stack.",
-                WINDOWING_MODE_PINNED, ACTIVITY_TYPE_STANDARD);
+            // Request that the PiP activity enter picture-in-picture mode (ensure it does not)
+            mBroadcastActionTrigger.doAction(ACTION_ENTER_PIP);
+            waitForEnterPip(PIP_ACTIVITY);
+            mAmWmState.assertDoesNotContainStack("Must not contain pinned stack.",
+                    WINDOWING_MODE_PINNED, ACTIVITY_TYPE_STANDARD);
 
-        // Enter the credentials and ensure that the activity actually entered picture-in-picture.
-        lockScreenSession.enterAndConfirmLockCredential();
-        mAmWmState.waitForKeyguardGone();
-        mAmWmState.assertKeyguardGone();
-        waitForEnterPip(PIP_ACTIVITY);
-        mAmWmState.assertContainsStack("Must contain pinned stack.", WINDOWING_MODE_PINNED,
-                ACTIVITY_TYPE_STANDARD);
+            // Enter the credentials and ensure that the activity actually entered picture-in
+            // -picture
+            lockScreenSession.enterAndConfirmLockCredential();
+            mAmWmState.waitForKeyguardGone();
+            mAmWmState.assertKeyguardGone();
+            waitForEnterPip(PIP_ACTIVITY);
+            mAmWmState.assertContainsStack("Must contain pinned stack.", WINDOWING_MODE_PINNED,
+                    ACTIVITY_TYPE_STANDARD);
+        }
     }
 
     @Test
-    public void testShowWhenLockedActivityAndPipActivity() {
+    public void testShowWhenLockedActivityAndPipActivity() throws Exception {
         assumeTrue(supportsPip());
 
-        final LockScreenSession lockScreenSession = createManagedLockScreenSession();
-        lockScreenSession.setLockCredential();
+        try (final LockScreenSession lockScreenSession = new LockScreenSession()) {
+            lockScreenSession.setLockCredential();
 
-        // Show an activity in PIP.
-        launchActivity(PIP_ACTIVITY, EXTRA_ENTER_PIP, "true");
-        waitForEnterPip(PIP_ACTIVITY);
-        mAmWmState.assertContainsStack("Must contain pinned stack.", WINDOWING_MODE_PINNED,
-                ACTIVITY_TYPE_STANDARD);
-        mAmWmState.assertVisibility(PIP_ACTIVITY, true);
+            // Show an activity in PIP
+            launchActivity(PIP_ACTIVITY, EXTRA_ENTER_PIP, "true");
+            waitForEnterPip(PIP_ACTIVITY);
+            mAmWmState.assertContainsStack("Must contain pinned stack.", WINDOWING_MODE_PINNED,
+                    ACTIVITY_TYPE_STANDARD);
+            mAmWmState.assertVisibility(PIP_ACTIVITY, true);
 
-        // Show an activity that will keep above the keyguard.
-        launchActivity(SHOW_WHEN_LOCKED_ACTIVITY);
-        mAmWmState.computeState(SHOW_WHEN_LOCKED_ACTIVITY);
-        mAmWmState.assertVisibility(SHOW_WHEN_LOCKED_ACTIVITY, true);
+            // Show an activity that will keep above the keyguard
+            launchActivity(SHOW_WHEN_LOCKED_ACTIVITY);
+            mAmWmState.computeState(SHOW_WHEN_LOCKED_ACTIVITY);
+            mAmWmState.assertVisibility(SHOW_WHEN_LOCKED_ACTIVITY, true);
 
-        // Lock the screen and ensure that the fullscreen activity showing over the lockscreen
-        // is visible, but not the PiP activity.
-        lockScreenSession.gotoKeyguard(SHOW_WHEN_LOCKED_ACTIVITY);
-        mAmWmState.computeState(true);
-        mAmWmState.assertKeyguardShowingAndOccluded();
-        mAmWmState.assertVisibility(SHOW_WHEN_LOCKED_ACTIVITY, true);
-        mAmWmState.assertVisibility(PIP_ACTIVITY, false);
+            // Lock the screen and ensure that the fullscreen activity showing over the lockscreen
+            // is visible, but not the PiP activity
+            lockScreenSession.gotoKeyguard(SHOW_WHEN_LOCKED_ACTIVITY);
+            mAmWmState.computeState(true);
+            mAmWmState.assertKeyguardShowingAndOccluded();
+            mAmWmState.assertVisibility(SHOW_WHEN_LOCKED_ACTIVITY, true);
+            mAmWmState.assertVisibility(PIP_ACTIVITY, false);
+        }
     }
 
     @Test
-    public void testShowWhenLockedPipActivity() {
+    public void testShowWhenLockedPipActivity() throws Exception {
         assumeTrue(supportsPip());
 
-        final LockScreenSession lockScreenSession = createManagedLockScreenSession();
-        lockScreenSession.setLockCredential();
+        try (final LockScreenSession lockScreenSession = new LockScreenSession()) {
+            lockScreenSession.setLockCredential();
 
-        // Show an activity in PIP.
-        launchActivity(PIP_ACTIVITY, EXTRA_ENTER_PIP, "true", EXTRA_SHOW_OVER_KEYGUARD, "true");
-        waitForEnterPip(PIP_ACTIVITY);
-        mAmWmState.assertContainsStack("Must contain pinned stack.", WINDOWING_MODE_PINNED,
-                ACTIVITY_TYPE_STANDARD);
-        mAmWmState.assertVisibility(PIP_ACTIVITY, true);
+            // Show an activity in PIP
+            launchActivity(PIP_ACTIVITY, EXTRA_ENTER_PIP, "true",
+                    EXTRA_SHOW_OVER_KEYGUARD, "true");
+            waitForEnterPip(PIP_ACTIVITY);
+            mAmWmState.assertContainsStack("Must contain pinned stack.", WINDOWING_MODE_PINNED,
+                    ACTIVITY_TYPE_STANDARD);
+            mAmWmState.assertVisibility(PIP_ACTIVITY, true);
 
-        // Lock the screen and ensure the PiP activity is not visible on the lockscreen even
-        // though it's marked as showing over the lockscreen itself.
-        lockScreenSession.gotoKeyguard();
-        mAmWmState.assertKeyguardShowingAndNotOccluded();
-        mAmWmState.assertVisibility(PIP_ACTIVITY, false);
+            // Lock the screen and ensure the PiP activity is not visible on the lockscreen even
+            // though it's marked as showing over the lockscreen itself
+            lockScreenSession.gotoKeyguard();
+            mAmWmState.assertKeyguardShowingAndNotOccluded();
+            mAmWmState.assertVisibility(PIP_ACTIVITY, false);
+        }
     }
 
     @Test
-    public void testDismissKeyguardPipActivity() {
+    public void testDismissKeyguardPipActivity() throws Exception {
         assumeTrue(supportsPip());
 
-        final LockScreenSession lockScreenSession = createManagedLockScreenSession();
-        // Show an activity in PIP.
-        launchActivity(PIP_ACTIVITY, EXTRA_ENTER_PIP, "true", EXTRA_DISMISS_KEYGUARD, "true");
-        waitForEnterPip(PIP_ACTIVITY);
-        mAmWmState.assertContainsStack("Must contain pinned stack.", WINDOWING_MODE_PINNED,
+        try (final LockScreenSession lockScreenSession = new LockScreenSession()) {
+            // Show an activity in PIP
+            launchActivity(PIP_ACTIVITY, EXTRA_ENTER_PIP, "true", EXTRA_DISMISS_KEYGUARD, "true");
+            waitForEnterPip(PIP_ACTIVITY);
+            mAmWmState.assertContainsStack("Must contain pinned stack.", WINDOWING_MODE_PINNED,
                 ACTIVITY_TYPE_STANDARD);
-        mAmWmState.assertVisibility(PIP_ACTIVITY, true);
+            mAmWmState.assertVisibility(PIP_ACTIVITY, true);
 
-        // Lock the screen and ensure the PiP activity is not visible on the lockscreen even
-        // though it's marked as dismiss keyguard.
-        lockScreenSession.gotoKeyguard();
-        mAmWmState.computeState(true);
-        mAmWmState.assertKeyguardShowingAndNotOccluded();
-        mAmWmState.assertVisibility(PIP_ACTIVITY, false);
+            // Lock the screen and ensure the PiP activity is not visible on the lockscreen even
+            // though it's marked as dismiss keyguard.
+            lockScreenSession.gotoKeyguard();
+            mAmWmState.computeState(true);
+            mAmWmState.assertKeyguardShowingAndNotOccluded();
+            mAmWmState.assertVisibility(PIP_ACTIVITY, false);
+        }
     }
 
     @Test
     public void testShowWhenLockedAttrImeActivityAndShowSoftInput() throws Exception {
-        final LockScreenSession lockScreenSession = createManagedLockScreenSession();
-        final MockImeSession mockImeSession = createManagedMockImeSession(this);
+        try (final LockScreenSession lockScreenSession = new LockScreenSession();
+             // Leverage MockImeSession to ensure at least an IME exists as default.
+             final MockImeSession mockImeSession = MockImeSession.create(mContext,
+                     getInstrumentation().getUiAutomation(), new ImeSettings.Builder())) {
+            lockScreenSession.setLockCredential().gotoKeyguard();
+            mAmWmState.assertKeyguardShowingAndNotOccluded();
+            launchActivity(SHOW_WHEN_LOCKED_ATTR_IME_ACTIVITY);
+            mAmWmState.computeState(SHOW_WHEN_LOCKED_ATTR_IME_ACTIVITY);
+            mAmWmState.assertVisibility(SHOW_WHEN_LOCKED_ATTR_IME_ACTIVITY, true);
 
-        lockScreenSession.setLockCredential().gotoKeyguard();
-        mAmWmState.assertKeyguardShowingAndNotOccluded();
-        launchActivity(SHOW_WHEN_LOCKED_ATTR_IME_ACTIVITY);
-        mAmWmState.computeState(SHOW_WHEN_LOCKED_ATTR_IME_ACTIVITY);
-        mAmWmState.assertVisibility(SHOW_WHEN_LOCKED_ATTR_IME_ACTIVITY, true);
-
-        // Make sure the activity has been called showSoftInput & IME window is visible.
-        final ImeEventStream stream = mockImeSession.openEventStream();
-        expectEvent(stream, event -> "showSoftInput".equals(event.getEventName()),
-                TimeUnit.SECONDS.toMillis(5) /* eventTimeout */);
-        // Assert the IME is shown on the expected display.
-        mAmWmState.waitAndAssertImeWindowShownOnDisplay(DEFAULT_DISPLAY);
+            // Make sure the activity has been called showSoftInput & IME window is visible.
+            final ImeEventStream stream = mockImeSession.openEventStream();
+            expectEvent(stream, event -> "showSoftInput".equals(event.getEventName()),
+                    TimeUnit.SECONDS.toMillis(5) /* eventTimeout */);
+            // Assert the IME is shown on the expected display.
+            mAmWmState.waitAndAssertImeWindowShownOnDisplay(DEFAULT_DISPLAY);
+        }
     }
 
     @Test
     public void testShowWhenLockedImeActivityAndShowSoftInput() throws Exception {
-        final LockScreenSession lockScreenSession = createManagedLockScreenSession();
-        final MockImeSession mockImeSession = createManagedMockImeSession(this);
-        final TestActivitySession<ShowWhenLockedImeActivity> imeTestActivitySession =
-                createManagedTestActivitySession();
+        try (final LockScreenSession lockScreenSession = new LockScreenSession();
+             final TestActivitySession<ShowWhenLockedImeActivity> imeTestActivitySession = new
+                     TestActivitySession<>();
+             // Leverage MockImeSession to ensure at least an IME exists as default.
+             final MockImeSession mockImeSession = MockImeSession.create(mContext,
+                     getInstrumentation().getUiAutomation(), new ImeSettings.Builder())) {
+            lockScreenSession.setLockCredential().gotoKeyguard();
+            mAmWmState.assertKeyguardShowingAndNotOccluded();
+            imeTestActivitySession.launchTestActivityOnDisplaySync(ShowWhenLockedImeActivity.class,
+                    DEFAULT_DISPLAY);
 
-        lockScreenSession.setLockCredential().gotoKeyguard();
-        mAmWmState.assertKeyguardShowingAndNotOccluded();
-        imeTestActivitySession.launchTestActivityOnDisplaySync(ShowWhenLockedImeActivity.class,
-                DEFAULT_DISPLAY);
-
-        // Make sure the activity has been called showSoftInput & IME window is visible.
-        final ImeEventStream stream = mockImeSession.openEventStream();
-        expectEvent(stream, event -> "showSoftInput".equals(event.getEventName()),
-                TimeUnit.SECONDS.toMillis(5) /* eventTimeout */);
-        // Assert the IME is shown on the expected display.
-        mAmWmState.waitAndAssertImeWindowShownOnDisplay(DEFAULT_DISPLAY);
-
+            // Make sure the activity has been called showSoftInput & IME window is visible.
+            final ImeEventStream stream = mockImeSession.openEventStream();
+            expectEvent(stream, event -> "showSoftInput".equals(event.getEventName()),
+                    TimeUnit.SECONDS.toMillis(5) /* eventTimeout */);
+            // Assert the IME is shown on the expected display.
+            mAmWmState.waitAndAssertImeWindowShownOnDisplay(DEFAULT_DISPLAY);
+        }
     }
 
     public static class ShowWhenLockedImeActivity extends Activity {
diff --git a/tests/framework/base/windowmanager/src/android/server/wm/KeyguardTests.java b/tests/framework/base/windowmanager/src/android/server/wm/KeyguardTests.java
index 0004b78..527dc31 100755
--- a/tests/framework/base/windowmanager/src/android/server/wm/KeyguardTests.java
+++ b/tests/framework/base/windowmanager/src/android/server/wm/KeyguardTests.java
@@ -86,7 +86,7 @@
             return mConfig.alwaysOnAvailable();
         }
 
-        void setAodEnabled(boolean enabled) {
+        void setAodEnabled(boolean enabled) throws Exception {
             set(enabled ? 1 : 0);
         }
     }
@@ -100,32 +100,32 @@
     }
 
     @Test
-    public void testKeyguardHidesActivity() {
-        final LockScreenSession lockScreenSession = createManagedLockScreenSession();
-        launchActivity(TEST_ACTIVITY);
-        mAmWmState.computeState(TEST_ACTIVITY);
-        mAmWmState.assertVisibility(TEST_ACTIVITY, true);
-        lockScreenSession.gotoKeyguard();
-        mAmWmState.computeState(true);
-        mAmWmState.assertKeyguardShowingAndNotOccluded();
-        assertTrue(mKeyguardManager.isKeyguardLocked());
-        mAmWmState.assertVisibility(TEST_ACTIVITY, false);
-
-        mObjectTracker.close(lockScreenSession);
+    public void testKeyguardHidesActivity() throws Exception {
+        try (final LockScreenSession lockScreenSession = new LockScreenSession()) {
+            launchActivity(TEST_ACTIVITY);
+            mAmWmState.computeState(TEST_ACTIVITY);
+            mAmWmState.assertVisibility(TEST_ACTIVITY, true);
+            lockScreenSession.gotoKeyguard();
+            mAmWmState.computeState(true);
+            mAmWmState.assertKeyguardShowingAndNotOccluded();
+            assertTrue(mKeyguardManager.isKeyguardLocked());
+            mAmWmState.assertVisibility(TEST_ACTIVITY, false);
+        }
         assertFalse(mKeyguardManager.isKeyguardLocked());
     }
 
     @Test
     @FlakyTest(bugId = 110276714)
-    public void testShowWhenLockedActivity() {
-        final LockScreenSession lockScreenSession = createManagedLockScreenSession();
-        launchActivity(SHOW_WHEN_LOCKED_ACTIVITY);
-        mAmWmState.computeState(SHOW_WHEN_LOCKED_ACTIVITY);
-        mAmWmState.assertVisibility(SHOW_WHEN_LOCKED_ACTIVITY, true);
-        lockScreenSession.gotoKeyguard(SHOW_WHEN_LOCKED_ACTIVITY);
-        mAmWmState.computeState(true);
-        mAmWmState.assertVisibility(SHOW_WHEN_LOCKED_ACTIVITY, true);
-        mAmWmState.assertKeyguardShowingAndOccluded();
+    public void testShowWhenLockedActivity() throws Exception {
+        try (final LockScreenSession lockScreenSession = new LockScreenSession()) {
+            launchActivity(SHOW_WHEN_LOCKED_ACTIVITY);
+            mAmWmState.computeState(SHOW_WHEN_LOCKED_ACTIVITY);
+            mAmWmState.assertVisibility(SHOW_WHEN_LOCKED_ACTIVITY, true);
+            lockScreenSession.gotoKeyguard(SHOW_WHEN_LOCKED_ACTIVITY);
+            mAmWmState.computeState(true);
+            mAmWmState.assertVisibility(SHOW_WHEN_LOCKED_ACTIVITY, true);
+            mAmWmState.assertKeyguardShowingAndOccluded();
+        }
     }
 
     /**
@@ -133,53 +133,56 @@
      * showing.
      */
     @Test
-    public void testShowWhenLockedActivity_withDialog() {
-        final LockScreenSession lockScreenSession = createManagedLockScreenSession();
-        launchActivity(SHOW_WHEN_LOCKED_WITH_DIALOG_ACTIVITY);
-        mAmWmState.computeState(SHOW_WHEN_LOCKED_WITH_DIALOG_ACTIVITY);
-        mAmWmState.assertVisibility(SHOW_WHEN_LOCKED_WITH_DIALOG_ACTIVITY, true);
-        lockScreenSession.gotoKeyguard(SHOW_WHEN_LOCKED_WITH_DIALOG_ACTIVITY);
-        mAmWmState.computeState(true);
-        mAmWmState.assertVisibility(SHOW_WHEN_LOCKED_WITH_DIALOG_ACTIVITY, true);
-        assertTrue(mAmWmState.getWmState().allWindowsVisible(
-                getWindowName(SHOW_WHEN_LOCKED_WITH_DIALOG_ACTIVITY)));
-        mAmWmState.assertKeyguardShowingAndOccluded();
+    public void testShowWhenLockedActivity_withDialog() throws Exception {
+        try (final LockScreenSession lockScreenSession = new LockScreenSession()) {
+            launchActivity(SHOW_WHEN_LOCKED_WITH_DIALOG_ACTIVITY);
+            mAmWmState.computeState(SHOW_WHEN_LOCKED_WITH_DIALOG_ACTIVITY);
+            mAmWmState.assertVisibility(SHOW_WHEN_LOCKED_WITH_DIALOG_ACTIVITY, true);
+            lockScreenSession.gotoKeyguard(SHOW_WHEN_LOCKED_WITH_DIALOG_ACTIVITY);
+            mAmWmState.computeState(true);
+            mAmWmState.assertVisibility(SHOW_WHEN_LOCKED_WITH_DIALOG_ACTIVITY, true);
+            assertTrue(mAmWmState.getWmState().allWindowsVisible(
+                    getWindowName(SHOW_WHEN_LOCKED_WITH_DIALOG_ACTIVITY)));
+            mAmWmState.assertKeyguardShowingAndOccluded();
+        }
     }
 
     /**
      * Tests whether multiple SHOW_WHEN_LOCKED activities are shown if the topmost is translucent.
      */
     @Test
-    public void testMultipleShowWhenLockedActivities() {
-        final LockScreenSession lockScreenSession = createManagedLockScreenSession();
-        launchActivity(SHOW_WHEN_LOCKED_ACTIVITY);
-        launchActivity(SHOW_WHEN_LOCKED_TRANSLUCENT_ACTIVITY);
-        mAmWmState.computeState(SHOW_WHEN_LOCKED_ACTIVITY,
-                SHOW_WHEN_LOCKED_TRANSLUCENT_ACTIVITY);
-        mAmWmState.assertVisibility(SHOW_WHEN_LOCKED_ACTIVITY, true);
-        mAmWmState.assertVisibility(SHOW_WHEN_LOCKED_TRANSLUCENT_ACTIVITY, true);
-        lockScreenSession.gotoKeyguard(
-                SHOW_WHEN_LOCKED_ACTIVITY, SHOW_WHEN_LOCKED_TRANSLUCENT_ACTIVITY);
-        mAmWmState.computeState(true);
-        mAmWmState.assertVisibility(SHOW_WHEN_LOCKED_ACTIVITY, true);
-        mAmWmState.assertVisibility(SHOW_WHEN_LOCKED_TRANSLUCENT_ACTIVITY, true);
-        mAmWmState.assertKeyguardShowingAndOccluded();
+    public void testMultipleShowWhenLockedActivities() throws Exception {
+        try (final LockScreenSession lockScreenSession = new LockScreenSession()) {
+            launchActivity(SHOW_WHEN_LOCKED_ACTIVITY);
+            launchActivity(SHOW_WHEN_LOCKED_TRANSLUCENT_ACTIVITY);
+            mAmWmState.computeState(SHOW_WHEN_LOCKED_ACTIVITY,
+                    SHOW_WHEN_LOCKED_TRANSLUCENT_ACTIVITY);
+            mAmWmState.assertVisibility(SHOW_WHEN_LOCKED_ACTIVITY, true);
+            mAmWmState.assertVisibility(SHOW_WHEN_LOCKED_TRANSLUCENT_ACTIVITY, true);
+            lockScreenSession.gotoKeyguard(
+                    SHOW_WHEN_LOCKED_ACTIVITY, SHOW_WHEN_LOCKED_TRANSLUCENT_ACTIVITY);
+            mAmWmState.computeState(true);
+            mAmWmState.assertVisibility(SHOW_WHEN_LOCKED_ACTIVITY, true);
+            mAmWmState.assertVisibility(SHOW_WHEN_LOCKED_TRANSLUCENT_ACTIVITY, true);
+            mAmWmState.assertKeyguardShowingAndOccluded();
+        }
     }
 
     /**
      * If we have a translucent SHOW_WHEN_LOCKED_ACTIVITY, the wallpaper should also be showing.
      */
     @Test
-    public void testTranslucentShowWhenLockedActivity() {
-        final LockScreenSession lockScreenSession = createManagedLockScreenSession();
-        launchActivity(SHOW_WHEN_LOCKED_TRANSLUCENT_ACTIVITY);
-        mAmWmState.computeState(SHOW_WHEN_LOCKED_TRANSLUCENT_ACTIVITY);
-        mAmWmState.assertVisibility(SHOW_WHEN_LOCKED_TRANSLUCENT_ACTIVITY, true);
-        lockScreenSession.gotoKeyguard(SHOW_WHEN_LOCKED_TRANSLUCENT_ACTIVITY);
-        mAmWmState.computeState(true);
-        mAmWmState.assertVisibility(SHOW_WHEN_LOCKED_TRANSLUCENT_ACTIVITY, true);
-        assertWallpaperShowing();
-        mAmWmState.assertKeyguardShowingAndOccluded();
+    public void testTranslucentShowWhenLockedActivity() throws Exception {
+        try (final LockScreenSession lockScreenSession = new LockScreenSession()) {
+            launchActivity(SHOW_WHEN_LOCKED_TRANSLUCENT_ACTIVITY);
+            mAmWmState.computeState(SHOW_WHEN_LOCKED_TRANSLUCENT_ACTIVITY);
+            mAmWmState.assertVisibility(SHOW_WHEN_LOCKED_TRANSLUCENT_ACTIVITY, true);
+            lockScreenSession.gotoKeyguard(SHOW_WHEN_LOCKED_TRANSLUCENT_ACTIVITY);
+            mAmWmState.computeState(true);
+            mAmWmState.assertVisibility(SHOW_WHEN_LOCKED_TRANSLUCENT_ACTIVITY, true);
+            assertWallpaperShowing();
+            mAmWmState.assertKeyguardShowingAndOccluded();
+        }
     }
 
     /**
@@ -187,31 +190,33 @@
      */
     @Test
     @FlakyTest
-    public void testTranslucentDoesntRevealBehind() {
-        final LockScreenSession lockScreenSession = createManagedLockScreenSession();
-        launchActivity(TEST_ACTIVITY);
-        launchActivity(SHOW_WHEN_LOCKED_TRANSLUCENT_ACTIVITY);
-        mAmWmState.computeState(TEST_ACTIVITY, SHOW_WHEN_LOCKED_TRANSLUCENT_ACTIVITY);
-        mAmWmState.assertVisibility(TEST_ACTIVITY, true);
-        mAmWmState.assertVisibility(SHOW_WHEN_LOCKED_TRANSLUCENT_ACTIVITY, true);
-        lockScreenSession.gotoKeyguard(SHOW_WHEN_LOCKED_TRANSLUCENT_ACTIVITY);
-        mAmWmState.computeState(true);
-        mAmWmState.assertVisibility(SHOW_WHEN_LOCKED_TRANSLUCENT_ACTIVITY, true);
-        mAmWmState.assertVisibility(TEST_ACTIVITY, false);
-        mAmWmState.assertKeyguardShowingAndOccluded();
+    public void testTranslucentDoesntRevealBehind() throws Exception {
+        try (final LockScreenSession lockScreenSession = new LockScreenSession()) {
+            launchActivity(TEST_ACTIVITY);
+            launchActivity(SHOW_WHEN_LOCKED_TRANSLUCENT_ACTIVITY);
+            mAmWmState.computeState(TEST_ACTIVITY, SHOW_WHEN_LOCKED_TRANSLUCENT_ACTIVITY);
+            mAmWmState.assertVisibility(TEST_ACTIVITY, true);
+            mAmWmState.assertVisibility(SHOW_WHEN_LOCKED_TRANSLUCENT_ACTIVITY, true);
+            lockScreenSession.gotoKeyguard(SHOW_WHEN_LOCKED_TRANSLUCENT_ACTIVITY);
+            mAmWmState.computeState(true);
+            mAmWmState.assertVisibility(SHOW_WHEN_LOCKED_TRANSLUCENT_ACTIVITY, true);
+            mAmWmState.assertVisibility(TEST_ACTIVITY, false);
+            mAmWmState.assertKeyguardShowingAndOccluded();
+        }
     }
 
     @Test
-    public void testDialogShowWhenLockedActivity() {
-        final LockScreenSession lockScreenSession = createManagedLockScreenSession();
-        launchActivity(SHOW_WHEN_LOCKED_DIALOG_ACTIVITY);
-        mAmWmState.computeState(SHOW_WHEN_LOCKED_DIALOG_ACTIVITY);
-        mAmWmState.assertVisibility(SHOW_WHEN_LOCKED_DIALOG_ACTIVITY, true);
-        lockScreenSession.gotoKeyguard();
-        mAmWmState.computeState(true);
-        mAmWmState.assertVisibility(SHOW_WHEN_LOCKED_DIALOG_ACTIVITY, true);
-        assertWallpaperShowing();
-        mAmWmState.assertKeyguardShowingAndOccluded();
+    public void testDialogShowWhenLockedActivity() throws Exception {
+        try (final LockScreenSession lockScreenSession = new LockScreenSession()) {
+            launchActivity(SHOW_WHEN_LOCKED_DIALOG_ACTIVITY);
+            mAmWmState.computeState(SHOW_WHEN_LOCKED_DIALOG_ACTIVITY);
+            mAmWmState.assertVisibility(SHOW_WHEN_LOCKED_DIALOG_ACTIVITY, true);
+            lockScreenSession.gotoKeyguard();
+            mAmWmState.computeState(true);
+            mAmWmState.assertVisibility(SHOW_WHEN_LOCKED_DIALOG_ACTIVITY, true);
+            assertWallpaperShowing();
+            mAmWmState.assertKeyguardShowingAndOccluded();
+        }
     }
 
     /**
@@ -219,22 +224,24 @@
      */
     @Test
     @Presubmit
-    public void testShowWhenLockedActivityWhileSplit() {
+    public void testShowWhenLockedActivityWhileSplit() throws Exception {
         assumeTrue(supportsSplitScreenMultiWindow());
 
-        final LockScreenSession lockScreenSession = createManagedLockScreenSession();
-        launchActivitiesInSplitScreen(
-                getLaunchActivityBuilder().setTargetActivity(LAUNCHING_ACTIVITY),
-                getLaunchActivityBuilder().setTargetActivity(SHOW_WHEN_LOCKED_ACTIVITY)
-                        .setRandomData(true)
-                        .setMultipleTask(false));
-        mAmWmState.assertVisibility(SHOW_WHEN_LOCKED_ACTIVITY, true);
-        lockScreenSession.gotoKeyguard(SHOW_WHEN_LOCKED_ACTIVITY);
-        mAmWmState.computeState(SHOW_WHEN_LOCKED_ACTIVITY);
-        mAmWmState.assertVisibility(SHOW_WHEN_LOCKED_ACTIVITY, true);
-        mAmWmState.assertKeyguardShowingAndOccluded();
-        mAmWmState.assertDoesNotContainStack("Activity must be full screen.",
-                WINDOWING_MODE_SPLIT_SCREEN_PRIMARY, ACTIVITY_TYPE_STANDARD);
+        try (final LockScreenSession lockScreenSession = new LockScreenSession()) {
+            launchActivitiesInSplitScreen(
+                    getLaunchActivityBuilder().setTargetActivity(LAUNCHING_ACTIVITY),
+                    getLaunchActivityBuilder().setTargetActivity(SHOW_WHEN_LOCKED_ACTIVITY)
+                            .setRandomData(true)
+                            .setMultipleTask(false)
+            );
+            mAmWmState.assertVisibility(SHOW_WHEN_LOCKED_ACTIVITY, true);
+            lockScreenSession.gotoKeyguard(SHOW_WHEN_LOCKED_ACTIVITY);
+            mAmWmState.computeState(SHOW_WHEN_LOCKED_ACTIVITY);
+            mAmWmState.assertVisibility(SHOW_WHEN_LOCKED_ACTIVITY, true);
+            mAmWmState.assertKeyguardShowingAndOccluded();
+            mAmWmState.assertDoesNotContainStack("Activity must be full screen.",
+                    WINDOWING_MODE_SPLIT_SCREEN_PRIMARY, ACTIVITY_TYPE_STANDARD);
+        }
     }
 
     /**
@@ -243,22 +250,23 @@
      */
     @Test
     @FlakyTest
-    public void testInheritShowWhenLockedAdd() {
-        final LockScreenSession lockScreenSession = createManagedLockScreenSession();
-        launchActivity(SHOW_WHEN_LOCKED_ATTR_ACTIVITY);
-        mAmWmState.assertVisibility(SHOW_WHEN_LOCKED_ATTR_ACTIVITY, true);
+    public void testInheritShowWhenLockedAdd() throws Exception {
+        try (final LockScreenSession lockScreenSession = new LockScreenSession()) {
+            launchActivity(SHOW_WHEN_LOCKED_ATTR_ACTIVITY);
+            mAmWmState.assertVisibility(SHOW_WHEN_LOCKED_ATTR_ACTIVITY, true);
 
-        launchActivity(INHERIT_SHOW_WHEN_LOCKED_ADD_ACTIVITY);
-        mAmWmState.computeState(
-                SHOW_WHEN_LOCKED_ATTR_ACTIVITY, INHERIT_SHOW_WHEN_LOCKED_ADD_ACTIVITY);
-        mAmWmState.assertVisibility(SHOW_WHEN_LOCKED_ATTR_ACTIVITY, false);
-        mAmWmState.assertVisibility(INHERIT_SHOW_WHEN_LOCKED_ADD_ACTIVITY, true);
+            launchActivity(INHERIT_SHOW_WHEN_LOCKED_ADD_ACTIVITY);
+            mAmWmState.computeState(
+                    SHOW_WHEN_LOCKED_ATTR_ACTIVITY, INHERIT_SHOW_WHEN_LOCKED_ADD_ACTIVITY);
+            mAmWmState.assertVisibility(SHOW_WHEN_LOCKED_ATTR_ACTIVITY, false);
+            mAmWmState.assertVisibility(INHERIT_SHOW_WHEN_LOCKED_ADD_ACTIVITY, true);
 
-        lockScreenSession.gotoKeyguard();
-        mAmWmState.computeState(true);
-        mAmWmState.assertKeyguardShowingAndOccluded();
-        mAmWmState.assertVisibility(SHOW_WHEN_LOCKED_ATTR_ACTIVITY, false);
-        mAmWmState.assertVisibility(INHERIT_SHOW_WHEN_LOCKED_ADD_ACTIVITY, true);
+            lockScreenSession.gotoKeyguard();
+            mAmWmState.computeState(true);
+            mAmWmState.assertKeyguardShowingAndOccluded();
+            mAmWmState.assertVisibility(SHOW_WHEN_LOCKED_ATTR_ACTIVITY, false);
+            mAmWmState.assertVisibility(INHERIT_SHOW_WHEN_LOCKED_ADD_ACTIVITY, true);
+        }
     }
 
     /**
@@ -268,23 +276,24 @@
      */
     @Test
     @FlakyTest
-    public void testInheritShowWhenLockedRemove() {
-        final LockScreenSession lockScreenSession = createManagedLockScreenSession();
-        launchActivity(SHOW_WHEN_LOCKED_ATTR_ACTIVITY);
-        mAmWmState.assertVisibility(SHOW_WHEN_LOCKED_ATTR_ACTIVITY, true);
+    public void testInheritShowWhenLockedRemove() throws Exception {
+        try (final LockScreenSession lockScreenSession = new LockScreenSession()) {
+            launchActivity(SHOW_WHEN_LOCKED_ATTR_ACTIVITY);
+            mAmWmState.assertVisibility(SHOW_WHEN_LOCKED_ATTR_ACTIVITY, true);
 
-        launchActivity(INHERIT_SHOW_WHEN_LOCKED_REMOVE_ACTIVITY);
-        mAmWmState.computeState(
-                SHOW_WHEN_LOCKED_ATTR_ACTIVITY, INHERIT_SHOW_WHEN_LOCKED_REMOVE_ACTIVITY);
-        mAmWmState.assertVisibility(SHOW_WHEN_LOCKED_ATTR_ACTIVITY, false);
-        mAmWmState.assertVisibility(INHERIT_SHOW_WHEN_LOCKED_REMOVE_ACTIVITY, true);
+            launchActivity(INHERIT_SHOW_WHEN_LOCKED_REMOVE_ACTIVITY);
+            mAmWmState.computeState(
+                    SHOW_WHEN_LOCKED_ATTR_ACTIVITY, INHERIT_SHOW_WHEN_LOCKED_REMOVE_ACTIVITY);
+            mAmWmState.assertVisibility(SHOW_WHEN_LOCKED_ATTR_ACTIVITY, false);
+            mAmWmState.assertVisibility(INHERIT_SHOW_WHEN_LOCKED_REMOVE_ACTIVITY, true);
 
-        lockScreenSession.gotoKeyguard();
-        mAmWmState.computeState(true);
-        mAmWmState.assertKeyguardShowingAndNotOccluded();
-        assertTrue(mKeyguardManager.isKeyguardLocked());
-        mAmWmState.assertVisibility(SHOW_WHEN_LOCKED_ATTR_ACTIVITY, false);
-        mAmWmState.assertVisibility(INHERIT_SHOW_WHEN_LOCKED_REMOVE_ACTIVITY, false);
+            lockScreenSession.gotoKeyguard();
+            mAmWmState.computeState(true);
+            mAmWmState.assertKeyguardShowingAndNotOccluded();
+            assertTrue(mKeyguardManager.isKeyguardLocked());
+            mAmWmState.assertVisibility(SHOW_WHEN_LOCKED_ATTR_ACTIVITY, false);
+            mAmWmState.assertVisibility(INHERIT_SHOW_WHEN_LOCKED_REMOVE_ACTIVITY, false);
+        }
     }
 
     /**
@@ -293,22 +302,23 @@
      * */
     @Test
     @FlakyTest
-    public void testInheritShowWhenLockedAttr() {
-        final LockScreenSession lockScreenSession = createManagedLockScreenSession();
-        launchActivity(SHOW_WHEN_LOCKED_ATTR_ACTIVITY);
-        mAmWmState.assertVisibility(SHOW_WHEN_LOCKED_ATTR_ACTIVITY, true);
+    public void testInheritShowWhenLockedAttr() throws Exception {
+        try (final LockScreenSession lockScreenSession = new LockScreenSession()) {
+            launchActivity(SHOW_WHEN_LOCKED_ATTR_ACTIVITY);
+            mAmWmState.assertVisibility(SHOW_WHEN_LOCKED_ATTR_ACTIVITY, true);
 
-        launchActivity(INHERIT_SHOW_WHEN_LOCKED_ATTR_ACTIVITY);
-        mAmWmState.computeState(
-                SHOW_WHEN_LOCKED_ATTR_ACTIVITY, INHERIT_SHOW_WHEN_LOCKED_ATTR_ACTIVITY);
-        mAmWmState.assertVisibility(SHOW_WHEN_LOCKED_ATTR_ACTIVITY, false);
-        mAmWmState.assertVisibility(INHERIT_SHOW_WHEN_LOCKED_ATTR_ACTIVITY, true);
+            launchActivity(INHERIT_SHOW_WHEN_LOCKED_ATTR_ACTIVITY);
+            mAmWmState.computeState(
+                    SHOW_WHEN_LOCKED_ATTR_ACTIVITY, INHERIT_SHOW_WHEN_LOCKED_ATTR_ACTIVITY);
+            mAmWmState.assertVisibility(SHOW_WHEN_LOCKED_ATTR_ACTIVITY, false);
+            mAmWmState.assertVisibility(INHERIT_SHOW_WHEN_LOCKED_ATTR_ACTIVITY, true);
 
-        lockScreenSession.gotoKeyguard();
-        mAmWmState.computeState(true);
-        mAmWmState.assertKeyguardShowingAndOccluded();
-        mAmWmState.assertVisibility(SHOW_WHEN_LOCKED_ATTR_ACTIVITY, false);
-        mAmWmState.assertVisibility(INHERIT_SHOW_WHEN_LOCKED_ATTR_ACTIVITY, true);
+            lockScreenSession.gotoKeyguard();
+            mAmWmState.computeState(true);
+            mAmWmState.assertKeyguardShowingAndOccluded();
+            mAmWmState.assertVisibility(SHOW_WHEN_LOCKED_ATTR_ACTIVITY, false);
+            mAmWmState.assertVisibility(INHERIT_SHOW_WHEN_LOCKED_ATTR_ACTIVITY, true);
+        }
     }
 
     /**
@@ -317,62 +327,63 @@
      * */
     @Test
     @FlakyTest
-    public void testNoInheritShowWhenLocked() {
-        final LockScreenSession lockScreenSession = createManagedLockScreenSession();
-        launchActivity(SHOW_WHEN_LOCKED_ATTR_ACTIVITY);
-        mAmWmState.assertVisibility(SHOW_WHEN_LOCKED_ATTR_ACTIVITY, true);
+    public void testNoInheritShowWhenLocked() throws Exception {
+        try (final LockScreenSession lockScreenSession = new LockScreenSession()) {
+            launchActivity(SHOW_WHEN_LOCKED_ATTR_ACTIVITY);
+            mAmWmState.assertVisibility(SHOW_WHEN_LOCKED_ATTR_ACTIVITY, true);
 
-        launchActivity(NO_INHERIT_SHOW_WHEN_LOCKED_ATTR_ACTIVITY);
-        mAmWmState.computeState(
-                SHOW_WHEN_LOCKED_ATTR_ACTIVITY, NO_INHERIT_SHOW_WHEN_LOCKED_ATTR_ACTIVITY);
-        mAmWmState.assertVisibility(SHOW_WHEN_LOCKED_ATTR_ACTIVITY, false);
-        mAmWmState.assertVisibility(NO_INHERIT_SHOW_WHEN_LOCKED_ATTR_ACTIVITY, true);
+            launchActivity(NO_INHERIT_SHOW_WHEN_LOCKED_ATTR_ACTIVITY);
+            mAmWmState.computeState(
+                    SHOW_WHEN_LOCKED_ATTR_ACTIVITY, NO_INHERIT_SHOW_WHEN_LOCKED_ATTR_ACTIVITY);
+            mAmWmState.assertVisibility(SHOW_WHEN_LOCKED_ATTR_ACTIVITY, false);
+            mAmWmState.assertVisibility(NO_INHERIT_SHOW_WHEN_LOCKED_ATTR_ACTIVITY, true);
 
-        lockScreenSession.gotoKeyguard();
-        mAmWmState.computeState(true);
-        mAmWmState.assertKeyguardShowingAndNotOccluded();
-        assertTrue(mKeyguardManager.isKeyguardLocked());
-        mAmWmState.assertVisibility(SHOW_WHEN_LOCKED_ATTR_ACTIVITY, false);
-        mAmWmState.assertVisibility(NO_INHERIT_SHOW_WHEN_LOCKED_ATTR_ACTIVITY, false);
+            lockScreenSession.gotoKeyguard();
+            mAmWmState.computeState(true);
+            mAmWmState.assertKeyguardShowingAndNotOccluded();
+            assertTrue(mKeyguardManager.isKeyguardLocked());
+            mAmWmState.assertVisibility(SHOW_WHEN_LOCKED_ATTR_ACTIVITY, false);
+            mAmWmState.assertVisibility(NO_INHERIT_SHOW_WHEN_LOCKED_ATTR_ACTIVITY, false);
+        }
     }
 
     @Test
     public void testNoTransientConfigurationWhenShowWhenLockedRequestsOrientation() {
-        final LockScreenSession lockScreenSession = createManagedLockScreenSession();
-        final ActivitySessionClient activitySession = createManagedActivityClientSession();
+        try (final LockScreenSession lockScreenSession = new LockScreenSession();
+                final ActivitySessionClient activitySession = new ActivitySessionClient(mContext)) {
+            final ActivitySession showWhenLockedActivitySession =
+                    activitySession.startActivity(getLaunchActivityBuilder()
+                            .setUseInstrumentation()
+                            .setTargetActivity(SHOW_WHEN_LOCKED_ATTR_ROTATION_ACTIVITY));
+            mAmWmState.assertVisibility(SHOW_WHEN_LOCKED_ATTR_ROTATION_ACTIVITY, true);
 
-        final ActivitySession showWhenLockedActivitySession =
-                activitySession.startActivity(getLaunchActivityBuilder()
-                        .setUseInstrumentation()
-                        .setTargetActivity(SHOW_WHEN_LOCKED_ATTR_ROTATION_ACTIVITY));
-        mAmWmState.assertVisibility(SHOW_WHEN_LOCKED_ATTR_ROTATION_ACTIVITY, true);
+            lockScreenSession.gotoKeyguard(SHOW_WHEN_LOCKED_ATTR_ROTATION_ACTIVITY);
 
-        lockScreenSession.gotoKeyguard(SHOW_WHEN_LOCKED_ATTR_ROTATION_ACTIVITY);
+            separateTestJournal();
 
-        separateTestJournal();
+            final int displayId = mAmWmState.getAmState()
+                    .getDisplayByActivity(SHOW_WHEN_LOCKED_ATTR_ROTATION_ACTIVITY);
+            ActivityManagerState.ActivityDisplay display = mAmWmState.getAmState()
+                    .getDisplay(displayId);
+            final int origDisplayOrientation = display.mFullConfiguration.orientation;
+            final int orientation = origDisplayOrientation == Configuration.ORIENTATION_LANDSCAPE
+                    ? SCREEN_ORIENTATION_PORTRAIT
+                    : SCREEN_ORIENTATION_LANDSCAPE;
+            showWhenLockedActivitySession.requestOrientation(orientation);
 
-        final int displayId = mAmWmState.getAmState()
-                .getDisplayByActivity(SHOW_WHEN_LOCKED_ATTR_ROTATION_ACTIVITY);
-        ActivityManagerState.ActivityDisplay display = mAmWmState.getAmState()
-                .getDisplay(displayId);
-        final int origDisplayOrientation = display.mFullConfiguration.orientation;
-        final int orientation = origDisplayOrientation == Configuration.ORIENTATION_LANDSCAPE
-                ? SCREEN_ORIENTATION_PORTRAIT
-                : SCREEN_ORIENTATION_LANDSCAPE;
-        showWhenLockedActivitySession.requestOrientation(orientation);
+            mAmWmState.waitForActivityOrientation(SHOW_WHEN_LOCKED_ATTR_ROTATION_ACTIVITY,
+                    orientation == SCREEN_ORIENTATION_LANDSCAPE
+                            ? Configuration.ORIENTATION_LANDSCAPE
+                            : Configuration.ORIENTATION_PORTRAIT);
 
-        mAmWmState.waitForActivityOrientation(SHOW_WHEN_LOCKED_ATTR_ROTATION_ACTIVITY,
-                orientation == SCREEN_ORIENTATION_LANDSCAPE
-                        ? Configuration.ORIENTATION_LANDSCAPE
-                        : Configuration.ORIENTATION_PORTRAIT);
+            display = mAmWmState.getAmState().getDisplay(displayId);
 
-        display = mAmWmState.getAmState().getDisplay(displayId);
-
-        // If the window is a non-fullscreen window (e.g. a freeform window) or the display is
-        // squared, there won't be activity lifecycle.
-        if (display.mFullConfiguration.orientation != origDisplayOrientation) {
-            assertActivityLifecycle(SHOW_WHEN_LOCKED_ATTR_ROTATION_ACTIVITY,
-                    false /* relaunched */);
+            // If the window is a non-fullscreen window (e.g. a freeform window) or the display is
+            // squared, there won't be activity lifecycle.
+            if (display.mFullConfiguration.orientation != origDisplayOrientation) {
+                assertActivityLifecycle(SHOW_WHEN_LOCKED_ATTR_ROTATION_ACTIVITY,
+                        false /* relaunched */);
+            }
         }
     }
 
@@ -396,147 +407,156 @@
     }
 
     private void testResumeOccludingActivityFromBackground(ComponentName occludingActivity) {
-        final LockScreenSession lockScreenSession = createManagedLockScreenSession();
-        lockScreenSession.gotoKeyguard();
-        mAmWmState.assertKeyguardShowingAndNotOccluded();
+        try (final LockScreenSession lockScreenSession = new LockScreenSession()) {
+            lockScreenSession.gotoKeyguard();
+            mAmWmState.assertKeyguardShowingAndNotOccluded();
 
-        // Launch an activity which is able to occlude keyguard.
-        getLaunchActivityBuilder().setUseInstrumentation()
-                .setTargetActivity(occludingActivity).execute();
+            // Launch an activity which is able to occlude keyguard.
+            getLaunchActivityBuilder().setUseInstrumentation()
+                    .setTargetActivity(occludingActivity).execute();
 
-        // Launch an activity without SHOW_WHEN_LOCKED and finish it.
-        getLaunchActivityBuilder().setUseInstrumentation()
-                .setMultipleTask(true)
-                // Don't wait for activity visible because keyguard will show.
-                .setWaitForLaunched(false)
-                .setTargetActivity(BROADCAST_RECEIVER_ACTIVITY).execute();
-        mAmWmState.waitForKeyguardShowingAndNotOccluded();
-        mAmWmState.assertKeyguardShowingAndNotOccluded();
+            // Launch an activity without SHOW_WHEN_LOCKED and finish it.
+            getLaunchActivityBuilder().setUseInstrumentation()
+                    .setMultipleTask(true)
+                    // Don't wait for activity visible because keyguard will show.
+                    .setWaitForLaunched(false)
+                    .setTargetActivity(BROADCAST_RECEIVER_ACTIVITY).execute();
+            mAmWmState.waitForKeyguardShowingAndNotOccluded();
+            mAmWmState.assertKeyguardShowingAndNotOccluded();
 
-        mBroadcastActionTrigger.finishBroadcastReceiverActivity();
-        mAmWmState.waitForKeyguardShowingAndOccluded();
+            mBroadcastActionTrigger.finishBroadcastReceiverActivity();
+            mAmWmState.waitForKeyguardShowingAndOccluded();
 
-        // The occluding activity should be resumed because it becomes the top activity.
-        mAmWmState.computeState(occludingActivity);
-        mAmWmState.assertVisibility(occludingActivity, true);
-        assertTrue(occludingActivity + " must be resumed.",
-                mAmWmState.getAmState().hasActivityState(occludingActivity,
-                        ActivityManagerState.STATE_RESUMED));
+            // The occluding activity should be resumed because it becomes the top activity.
+            mAmWmState.computeState(occludingActivity);
+            mAmWmState.assertVisibility(occludingActivity, true);
+            assertTrue(occludingActivity + " must be resumed.",
+                    mAmWmState.getAmState().hasActivityState(occludingActivity,
+                            ActivityManagerState.STATE_RESUMED));
+        }
     }
 
     /**
      * Tests whether a FLAG_DISMISS_KEYGUARD activity occludes Keyguard.
      */
     @Test
-    public void testDismissKeyguardActivity() {
-        final LockScreenSession lockScreenSession = createManagedLockScreenSession();
-        lockScreenSession.gotoKeyguard();
-        mAmWmState.computeState(true);
-        assertTrue(mAmWmState.getAmState().getKeyguardControllerState().keyguardShowing);
-        launchActivity(DISMISS_KEYGUARD_ACTIVITY);
-        mAmWmState.waitForKeyguardShowingAndOccluded();
-        mAmWmState.computeState(DISMISS_KEYGUARD_ACTIVITY);
-        mAmWmState.assertVisibility(DISMISS_KEYGUARD_ACTIVITY, true);
-        mAmWmState.assertKeyguardShowingAndOccluded();
+    public void testDismissKeyguardActivity() throws Exception {
+        try (final LockScreenSession lockScreenSession = new LockScreenSession()) {
+            lockScreenSession.gotoKeyguard();
+            mAmWmState.computeState(true);
+            assertTrue(mAmWmState.getAmState().getKeyguardControllerState().keyguardShowing);
+            launchActivity(DISMISS_KEYGUARD_ACTIVITY);
+            mAmWmState.waitForKeyguardShowingAndOccluded();
+            mAmWmState.computeState(DISMISS_KEYGUARD_ACTIVITY);
+            mAmWmState.assertVisibility(DISMISS_KEYGUARD_ACTIVITY, true);
+            mAmWmState.assertKeyguardShowingAndOccluded();
+        }
     }
 
     @Test
-    public void testDismissKeyguardActivity_method() {
-        final LockScreenSession lockScreenSession = createManagedLockScreenSession();
-        separateTestJournal();
-        lockScreenSession.gotoKeyguard();
-        mAmWmState.computeState(true);
-        assertTrue(mAmWmState.getAmState().getKeyguardControllerState().keyguardShowing);
-        launchActivity(DISMISS_KEYGUARD_METHOD_ACTIVITY);
-        mAmWmState.waitForKeyguardGone();
-        mAmWmState.computeState(DISMISS_KEYGUARD_METHOD_ACTIVITY);
-        mAmWmState.assertVisibility(DISMISS_KEYGUARD_METHOD_ACTIVITY, true);
-        assertFalse(mAmWmState.getAmState().getKeyguardControllerState().keyguardShowing);
-        assertOnDismissSucceeded(DISMISS_KEYGUARD_METHOD_ACTIVITY);
+    public void testDismissKeyguardActivity_method() throws Exception {
+        try (final LockScreenSession lockScreenSession = new LockScreenSession()) {
+            separateTestJournal();
+            lockScreenSession.gotoKeyguard();
+            mAmWmState.computeState(true);
+            assertTrue(mAmWmState.getAmState().getKeyguardControllerState().keyguardShowing);
+            launchActivity(DISMISS_KEYGUARD_METHOD_ACTIVITY);
+            mAmWmState.waitForKeyguardGone();
+            mAmWmState.computeState(DISMISS_KEYGUARD_METHOD_ACTIVITY);
+            mAmWmState.assertVisibility(DISMISS_KEYGUARD_METHOD_ACTIVITY, true);
+            assertFalse(mAmWmState.getAmState().getKeyguardControllerState().keyguardShowing);
+            assertOnDismissSucceeded(DISMISS_KEYGUARD_METHOD_ACTIVITY);
+        }
     }
 
     @Test
-    public void testDismissKeyguardActivity_method_notTop() {
-        final LockScreenSession lockScreenSession = createManagedLockScreenSession();
-        separateTestJournal();
-        lockScreenSession.gotoKeyguard();
-        mAmWmState.computeState(true);
-        assertTrue(mAmWmState.getAmState().getKeyguardControllerState().keyguardShowing);
-        launchActivity(BROADCAST_RECEIVER_ACTIVITY);
-        launchActivity(TEST_ACTIVITY);
-        mBroadcastActionTrigger.dismissKeyguardByMethod();
-        assertOnDismissError(BROADCAST_RECEIVER_ACTIVITY);
+    public void testDismissKeyguardActivity_method_notTop() throws Exception {
+        try (final LockScreenSession lockScreenSession = new LockScreenSession()) {
+            separateTestJournal();
+            lockScreenSession.gotoKeyguard();
+            mAmWmState.computeState(true);
+            assertTrue(mAmWmState.getAmState().getKeyguardControllerState().keyguardShowing);
+            launchActivity(BROADCAST_RECEIVER_ACTIVITY);
+            launchActivity(TEST_ACTIVITY);
+            mBroadcastActionTrigger.dismissKeyguardByMethod();
+            assertOnDismissError(BROADCAST_RECEIVER_ACTIVITY);
+        }
     }
 
     @Test
-    public void testDismissKeyguardActivity_method_turnScreenOn() {
-        final LockScreenSession lockScreenSession = createManagedLockScreenSession();
-        separateTestJournal();
-        lockScreenSession.sleepDevice();
-        mAmWmState.computeState(true);
-        assertTrue(mAmWmState.getAmState().getKeyguardControllerState().keyguardShowing);
-        launchActivity(TURN_SCREEN_ON_DISMISS_KEYGUARD_ACTIVITY);
-        mAmWmState.waitForKeyguardGone();
-        mAmWmState.computeState(TURN_SCREEN_ON_DISMISS_KEYGUARD_ACTIVITY);
-        mAmWmState.assertVisibility(TURN_SCREEN_ON_DISMISS_KEYGUARD_ACTIVITY, true);
-        assertFalse(mAmWmState.getAmState().getKeyguardControllerState().keyguardShowing);
-        assertOnDismissSucceeded(TURN_SCREEN_ON_DISMISS_KEYGUARD_ACTIVITY);
-        assertTrue(isDisplayOn(DEFAULT_DISPLAY));
+    public void testDismissKeyguardActivity_method_turnScreenOn() throws Exception {
+        try (final LockScreenSession lockScreenSession = new LockScreenSession()) {
+            separateTestJournal();
+            lockScreenSession.sleepDevice();
+            mAmWmState.computeState(true);
+            assertTrue(mAmWmState.getAmState().getKeyguardControllerState().keyguardShowing);
+            launchActivity(TURN_SCREEN_ON_DISMISS_KEYGUARD_ACTIVITY);
+            mAmWmState.waitForKeyguardGone();
+            mAmWmState.computeState(TURN_SCREEN_ON_DISMISS_KEYGUARD_ACTIVITY);
+            mAmWmState.assertVisibility(TURN_SCREEN_ON_DISMISS_KEYGUARD_ACTIVITY, true);
+            assertFalse(mAmWmState.getAmState().getKeyguardControllerState().keyguardShowing);
+            assertOnDismissSucceeded(TURN_SCREEN_ON_DISMISS_KEYGUARD_ACTIVITY);
+            assertTrue(isDisplayOn(DEFAULT_DISPLAY));
+        }
     }
 
     @Test
-    public void testDismissKeyguard_fromShowWhenLocked_notAllowed() {
-        final LockScreenSession lockScreenSession = createManagedLockScreenSession();
-        lockScreenSession.gotoKeyguard();
-        mAmWmState.assertKeyguardShowingAndNotOccluded();
-        launchActivity(SHOW_WHEN_LOCKED_ACTIVITY);
-        mAmWmState.computeState(SHOW_WHEN_LOCKED_ACTIVITY);
-        mAmWmState.assertVisibility(SHOW_WHEN_LOCKED_ACTIVITY, true);
-        mAmWmState.assertKeyguardShowingAndOccluded();
-        mBroadcastActionTrigger.dismissKeyguardByFlag();
-        mAmWmState.assertKeyguardShowingAndOccluded();
-        mAmWmState.assertVisibility(SHOW_WHEN_LOCKED_ACTIVITY, true);
+    public void testDismissKeyguard_fromShowWhenLocked_notAllowed() throws Exception {
+        try (final LockScreenSession lockScreenSession = new LockScreenSession()) {
+            lockScreenSession.gotoKeyguard();
+            mAmWmState.assertKeyguardShowingAndNotOccluded();
+            launchActivity(SHOW_WHEN_LOCKED_ACTIVITY);
+            mAmWmState.computeState(SHOW_WHEN_LOCKED_ACTIVITY);
+            mAmWmState.assertVisibility(SHOW_WHEN_LOCKED_ACTIVITY, true);
+            mAmWmState.assertKeyguardShowingAndOccluded();
+            mBroadcastActionTrigger.dismissKeyguardByFlag();
+            mAmWmState.assertKeyguardShowingAndOccluded();
+            mAmWmState.assertVisibility(SHOW_WHEN_LOCKED_ACTIVITY, true);
+        }
     }
 
     @Test
-    public void testKeyguardLock() {
-        final LockScreenSession lockScreenSession = createManagedLockScreenSession();
-        lockScreenSession.gotoKeyguard();
-        mAmWmState.assertKeyguardShowingAndNotOccluded();
-        launchActivity(KEYGUARD_LOCK_ACTIVITY);
-        mAmWmState.computeState(KEYGUARD_LOCK_ACTIVITY);
-        mAmWmState.assertVisibility(KEYGUARD_LOCK_ACTIVITY, true);
-        mBroadcastActionTrigger.finishBroadcastReceiverActivity();
-        mAmWmState.waitForKeyguardShowingAndNotOccluded();
-        mAmWmState.assertKeyguardShowingAndNotOccluded();
+    public void testKeyguardLock() throws Exception {
+        try (final LockScreenSession lockScreenSession = new LockScreenSession()) {
+            lockScreenSession.gotoKeyguard();
+            mAmWmState.assertKeyguardShowingAndNotOccluded();
+            launchActivity(KEYGUARD_LOCK_ACTIVITY);
+            mAmWmState.computeState(KEYGUARD_LOCK_ACTIVITY);
+            mAmWmState.assertVisibility(KEYGUARD_LOCK_ACTIVITY, true);
+            mBroadcastActionTrigger.finishBroadcastReceiverActivity();
+            mAmWmState.waitForKeyguardShowingAndNotOccluded();
+            mAmWmState.assertKeyguardShowingAndNotOccluded();
+        }
     }
 
     @Test
-    public void testUnoccludeRotationChange() {
+    public void testUnoccludeRotationChange() throws Exception {
+
         // Go home now to make sure Home is behind Keyguard.
         launchHomeActivity();
-        final LockScreenSession lockScreenSession = createManagedLockScreenSession();
-        final RotationSession rotationSession = createManagedRotationSession();
-        lockScreenSession.gotoKeyguard();
-        mAmWmState.assertKeyguardShowingAndNotOccluded();
-        launchActivity(SHOW_WHEN_LOCKED_ACTIVITY);
-        mAmWmState.computeState(SHOW_WHEN_LOCKED_ACTIVITY);
-        mAmWmState.assertVisibility(SHOW_WHEN_LOCKED_ACTIVITY, true);
+        try (final LockScreenSession lockScreenSession = new LockScreenSession();
+             final RotationSession rotationSession = new RotationSession()) {
+            lockScreenSession.gotoKeyguard();
+            mAmWmState.assertKeyguardShowingAndNotOccluded();
+            launchActivity(SHOW_WHEN_LOCKED_ACTIVITY);
+            mAmWmState.computeState(SHOW_WHEN_LOCKED_ACTIVITY);
+            mAmWmState.assertVisibility(SHOW_WHEN_LOCKED_ACTIVITY, true);
 
-        rotationSession.set(ROTATION_90);
-        pressBackButton();
-        mAmWmState.waitForKeyguardShowingAndNotOccluded();
-        mAmWmState.waitForDisplayUnfrozen();
-        mAmWmState.waitForAppTransitionIdleOnDisplay(DEFAULT_DISPLAY);
-        mAmWmState.assertSanity();
-        mAmWmState.assertHomeActivityVisible(false);
-        mAmWmState.assertKeyguardShowingAndNotOccluded();
-        // The activity may not be destroyed immediately.
-        mAmWmState.waitForWithWmState(
-                wmState -> !wmState.containsWindow(getWindowName(SHOW_WHEN_LOCKED_ACTIVITY)),
-                getActivityName(SHOW_WHEN_LOCKED_ACTIVITY) + " to be removed");
-        // The {@link SHOW_WHEN_LOCKED_ACTIVITY} has gone because of {@link pressBackButton()}.
-        mAmWmState.assertNotExist(SHOW_WHEN_LOCKED_ACTIVITY);
+            rotationSession.set(ROTATION_90);
+            pressBackButton();
+            mAmWmState.waitForKeyguardShowingAndNotOccluded();
+            mAmWmState.waitForDisplayUnfrozen();
+            mAmWmState.waitForAppTransitionIdleOnDisplay(DEFAULT_DISPLAY);
+            mAmWmState.assertSanity();
+            mAmWmState.assertHomeActivityVisible(false);
+            mAmWmState.assertKeyguardShowingAndNotOccluded();
+            // The activity may not be destroyed immediately.
+            mAmWmState.waitForWithWmState(
+                    wmState -> !wmState.containsWindow(getWindowName(SHOW_WHEN_LOCKED_ACTIVITY)),
+                    getActivityName(SHOW_WHEN_LOCKED_ACTIVITY) + " to be removed");
+            // The {@link SHOW_WHEN_LOCKED_ACTIVITY} has gone because of {@link pressBackButton()}.
+            mAmWmState.assertNotExist(SHOW_WHEN_LOCKED_ACTIVITY);
+        }
     }
 
     private void assertWallpaperShowing() {
@@ -547,23 +567,24 @@
     }
 
     @Test
-    public void testDismissKeyguardAttrActivity_method_turnScreenOn() {
-        final LockScreenSession lockScreenSession = createManagedLockScreenSession();
-        lockScreenSession.sleepDevice();
+    public void testDismissKeyguardAttrActivity_method_turnScreenOn() throws Exception {
+        try (final LockScreenSession lockScreenSession = new LockScreenSession()) {
+            lockScreenSession.sleepDevice();
 
-        separateTestJournal();
-        mAmWmState.computeState(true);
-        assertTrue(mAmWmState.getAmState().getKeyguardControllerState().keyguardShowing);
-        launchActivity(TURN_SCREEN_ON_ATTR_DISMISS_KEYGUARD_ACTIVITY);
-        mAmWmState.waitForKeyguardGone();
-        mAmWmState.assertVisibility(TURN_SCREEN_ON_ATTR_DISMISS_KEYGUARD_ACTIVITY, true);
-        assertFalse(mAmWmState.getAmState().getKeyguardControllerState().keyguardShowing);
-        assertOnDismissSucceeded(TURN_SCREEN_ON_ATTR_DISMISS_KEYGUARD_ACTIVITY);
-        assertTrue(isDisplayOn(DEFAULT_DISPLAY));
+            separateTestJournal();
+            mAmWmState.computeState(true);
+            assertTrue(mAmWmState.getAmState().getKeyguardControllerState().keyguardShowing);
+            launchActivity(TURN_SCREEN_ON_ATTR_DISMISS_KEYGUARD_ACTIVITY);
+            mAmWmState.waitForKeyguardGone();
+            mAmWmState.assertVisibility(TURN_SCREEN_ON_ATTR_DISMISS_KEYGUARD_ACTIVITY, true);
+            assertFalse(mAmWmState.getAmState().getKeyguardControllerState().keyguardShowing);
+            assertOnDismissSucceeded(TURN_SCREEN_ON_ATTR_DISMISS_KEYGUARD_ACTIVITY);
+            assertTrue(isDisplayOn(DEFAULT_DISPLAY));
+        }
     }
 
     @Test
-    public void testScreenOffWhileOccludedStopsActivityNoAod() {
+    public void testScreenOffWhileOccludedStopsActivityNoAod() throws Exception {
         try (final AodSession aodSession = new AodSession()) {
             aodSession.setAodEnabled(false);
             testScreenOffWhileOccludedStopsActivity(false /* assertAod */);
@@ -571,7 +592,7 @@
     }
 
     @Test
-    public void testScreenOffWhileOccludedStopsActivityAod() {
+    public void testScreenOffWhileOccludedStopsActivityAod() throws Exception {
         try (final AodSession aodSession = new AodSession()) {
             assumeTrue(aodSession.isAodAvailable());
             aodSession.setAodEnabled(true);
@@ -605,7 +626,7 @@
     }
 
     @Test
-    public void testScreenOffCausesSingleStopNoAod() {
+    public void testScreenOffCausesSingleStopNoAod() throws Exception {
         try (final AodSession aodSession = new AodSession()) {
             aodSession.setAodEnabled(false);
             testScreenOffCausesSingleStop();
@@ -613,7 +634,7 @@
     }
 
     @Test
-    public void testScreenOffCausesSingleStopAod() {
+    public void testScreenOffCausesSingleStopAod() throws Exception {
         try (final AodSession aodSession = new AodSession()) {
             assumeTrue(aodSession.isAodAvailable());
             aodSession.setAodEnabled(true);
diff --git a/tests/framework/base/windowmanager/src/android/server/wm/KeyguardTransitionTests.java b/tests/framework/base/windowmanager/src/android/server/wm/KeyguardTransitionTests.java
index b941913..88c0f26 100644
--- a/tests/framework/base/windowmanager/src/android/server/wm/KeyguardTransitionTests.java
+++ b/tests/framework/base/windowmanager/src/android/server/wm/KeyguardTransitionTests.java
@@ -56,99 +56,112 @@
     }
 
     @Test
-    public void testUnlock() {
-        final LockScreenSession lockScreenSession = createManagedLockScreenSession();
-        launchActivity(TEST_ACTIVITY);
-        lockScreenSession.gotoKeyguard().unlockDevice();
-        mAmWmState.computeState(TEST_ACTIVITY);
-        assertEquals("Picked wrong transition", TRANSIT_KEYGUARD_GOING_AWAY,
-                mAmWmState.getWmState().getDefaultDisplayLastTransition());
+    public void testUnlock() throws Exception {
+        try (final LockScreenSession lockScreenSession = new LockScreenSession()) {
+            launchActivity(TEST_ACTIVITY);
+            lockScreenSession.gotoKeyguard()
+                    .unlockDevice();
+            mAmWmState.computeState(TEST_ACTIVITY);
+            assertEquals("Picked wrong transition", TRANSIT_KEYGUARD_GOING_AWAY,
+                    mAmWmState.getWmState().getDefaultDisplayLastTransition());
+        }
     }
 
     @Test
-    public void testUnlockWallpaper() {
-        final LockScreenSession lockScreenSession = createManagedLockScreenSession();
-        launchActivity(WALLPAPAER_ACTIVITY);
-        lockScreenSession.gotoKeyguard().unlockDevice();
-        mAmWmState.computeState(WALLPAPAER_ACTIVITY);
-        assertEquals("Picked wrong transition", TRANSIT_KEYGUARD_GOING_AWAY_ON_WALLPAPER,
-                mAmWmState.getWmState().getDefaultDisplayLastTransition());
+    public void testUnlockWallpaper() throws Exception {
+        try (final LockScreenSession lockScreenSession = new LockScreenSession()) {
+            launchActivity(WALLPAPAER_ACTIVITY);
+            lockScreenSession.gotoKeyguard()
+                    .unlockDevice();
+            mAmWmState.computeState(WALLPAPAER_ACTIVITY);
+            assertEquals("Picked wrong transition", TRANSIT_KEYGUARD_GOING_AWAY_ON_WALLPAPER,
+                    mAmWmState.getWmState().getDefaultDisplayLastTransition());
+        }
     }
 
     @Test
-    public void testOcclude() {
-        createManagedLockScreenSession().gotoKeyguard();
-        launchActivity(SHOW_WHEN_LOCKED_ACTIVITY);
-        mAmWmState.computeState(SHOW_WHEN_LOCKED_ACTIVITY);
-        assertEquals("Picked wrong transition", TRANSIT_KEYGUARD_OCCLUDE,
-                mAmWmState.getWmState().getDefaultDisplayLastTransition());
+    public void testOcclude() throws Exception {
+        try (final LockScreenSession lockScreenSession = new LockScreenSession()) {
+            lockScreenSession.gotoKeyguard();
+            launchActivity(SHOW_WHEN_LOCKED_ACTIVITY);
+            mAmWmState.computeState(SHOW_WHEN_LOCKED_ACTIVITY);
+            assertEquals("Picked wrong transition", TRANSIT_KEYGUARD_OCCLUDE,
+                    mAmWmState.getWmState().getDefaultDisplayLastTransition());
+        }
     }
 
     @Test
-    public void testUnocclude() {
-        createManagedLockScreenSession().gotoKeyguard();
-        launchActivity(SHOW_WHEN_LOCKED_ACTIVITY);
-        launchActivity(TEST_ACTIVITY);
-        mAmWmState.waitForKeyguardShowingAndNotOccluded();
-        mAmWmState.computeState(true);
-        assertEquals("Picked wrong transition", TRANSIT_KEYGUARD_UNOCCLUDE,
-                mAmWmState.getWmState().getDefaultDisplayLastTransition());
+    public void testUnocclude() throws Exception {
+        try (final LockScreenSession lockScreenSession = new LockScreenSession()) {
+            lockScreenSession.gotoKeyguard();
+            launchActivity(SHOW_WHEN_LOCKED_ACTIVITY);
+            launchActivity(TEST_ACTIVITY);
+            mAmWmState.waitForKeyguardShowingAndNotOccluded();
+            mAmWmState.computeState(true);
+            assertEquals("Picked wrong transition", TRANSIT_KEYGUARD_UNOCCLUDE,
+                    mAmWmState.getWmState().getDefaultDisplayLastTransition());
+        }
     }
 
     @Test
-    public void testNewActivityDuringOccluded() {
-        final LockScreenSession lockScreenSession = createManagedLockScreenSession();
-        launchActivity(SHOW_WHEN_LOCKED_ACTIVITY);
-        lockScreenSession.gotoKeyguard(SHOW_WHEN_LOCKED_ACTIVITY);
-        launchActivity(SHOW_WHEN_LOCKED_WITH_DIALOG_ACTIVITY);
-        mAmWmState.computeState(SHOW_WHEN_LOCKED_WITH_DIALOG_ACTIVITY);
-        assertEquals("Picked wrong transition", TRANSIT_ACTIVITY_OPEN,
-                mAmWmState.getWmState().getDefaultDisplayLastTransition());
+    public void testNewActivityDuringOccluded() throws Exception {
+        try (final LockScreenSession lockScreenSession = new LockScreenSession()) {
+            launchActivity(SHOW_WHEN_LOCKED_ACTIVITY);
+            lockScreenSession.gotoKeyguard(SHOW_WHEN_LOCKED_ACTIVITY);
+            launchActivity(SHOW_WHEN_LOCKED_WITH_DIALOG_ACTIVITY);
+            mAmWmState.computeState(SHOW_WHEN_LOCKED_WITH_DIALOG_ACTIVITY);
+            assertEquals("Picked wrong transition", TRANSIT_ACTIVITY_OPEN,
+                    mAmWmState.getWmState().getDefaultDisplayLastTransition());
+        }
     }
 
     @Test
-    public void testOccludeManifestAttr() {
-        final LockScreenSession lockScreenSession = createManagedLockScreenSession();
-        lockScreenSession.gotoKeyguard();
-        separateTestJournal();
-        launchActivity(SHOW_WHEN_LOCKED_ATTR_ACTIVITY);
-        mAmWmState.computeState(SHOW_WHEN_LOCKED_ATTR_ACTIVITY);
-        assertEquals("Picked wrong transition", TRANSIT_KEYGUARD_OCCLUDE,
-                mAmWmState.getWmState().getDefaultDisplayLastTransition());
-        assertSingleLaunch(SHOW_WHEN_LOCKED_ATTR_ACTIVITY);
+    public void testOccludeManifestAttr() throws Exception {
+        try (final LockScreenSession lockScreenSession = new LockScreenSession()) {
+            lockScreenSession.gotoKeyguard();
+            separateTestJournal();
+            launchActivity(SHOW_WHEN_LOCKED_ATTR_ACTIVITY);
+            mAmWmState.computeState(SHOW_WHEN_LOCKED_ATTR_ACTIVITY);
+            assertEquals("Picked wrong transition", TRANSIT_KEYGUARD_OCCLUDE,
+                    mAmWmState.getWmState().getDefaultDisplayLastTransition());
+            assertSingleLaunch(SHOW_WHEN_LOCKED_ATTR_ACTIVITY);
+        }
     }
 
     @Test
-    public void testOccludeAttrRemove() {
-        final LockScreenSession lockScreenSession = createManagedLockScreenSession();
-        lockScreenSession.gotoKeyguard();
-        separateTestJournal();
-        launchActivity(SHOW_WHEN_LOCKED_ATTR_REMOVE_ATTR_ACTIVITY);
-        mAmWmState.computeState(SHOW_WHEN_LOCKED_ATTR_REMOVE_ATTR_ACTIVITY);
-        assertEquals("Picked wrong transition", TRANSIT_KEYGUARD_OCCLUDE,
-                mAmWmState.getWmState().getDefaultDisplayLastTransition());
-        assertSingleLaunch(SHOW_WHEN_LOCKED_ATTR_REMOVE_ATTR_ACTIVITY);
+    public void testOccludeAttrRemove() throws Exception {
+        try (final LockScreenSession lockScreenSession = new LockScreenSession()) {
+            lockScreenSession.gotoKeyguard();
+            separateTestJournal();
+            launchActivity(SHOW_WHEN_LOCKED_ATTR_REMOVE_ATTR_ACTIVITY);
+            mAmWmState.computeState(SHOW_WHEN_LOCKED_ATTR_REMOVE_ATTR_ACTIVITY);
+            assertEquals("Picked wrong transition", TRANSIT_KEYGUARD_OCCLUDE,
+                    mAmWmState.getWmState().getDefaultDisplayLastTransition());
+            assertSingleLaunch(SHOW_WHEN_LOCKED_ATTR_REMOVE_ATTR_ACTIVITY);
 
-        // Waiting for the standard keyguard since
-        // {@link SHOW_WHEN_LOCKED_ATTR_REMOVE_ATTR_ACTIVITY} called
-        // {@link Activity#showWhenLocked(boolean)} and removed the attribute.
-        lockScreenSession.gotoKeyguard();
-        separateTestJournal();
-        // Waiting for {@link SHOW_WHEN_LOCKED_ATTR_REMOVE_ATTR_ACTIVITY} stopped since it
-        // already lost show-when-locked attribute.
-        launchActivityNoWait(SHOW_WHEN_LOCKED_ATTR_REMOVE_ATTR_ACTIVITY);
-        mAmWmState.waitForActivityState(SHOW_WHEN_LOCKED_ATTR_REMOVE_ATTR_ACTIVITY, STATE_STOPPED);
-        assertSingleStartAndStop(SHOW_WHEN_LOCKED_ATTR_REMOVE_ATTR_ACTIVITY);
+            // Waiting for the standard keyguard since
+            // {@link SHOW_WHEN_LOCKED_ATTR_REMOVE_ATTR_ACTIVITY} called
+            // {@link Activity#showWhenLocked(boolean)} and removed the attribute.
+            lockScreenSession.gotoKeyguard();
+            separateTestJournal();
+            // Waiting for {@link SHOW_WHEN_LOCKED_ATTR_REMOVE_ATTR_ACTIVITY} stopped since it
+            // already lost show-when-locked attribute.
+            launchActivityNoWait(SHOW_WHEN_LOCKED_ATTR_REMOVE_ATTR_ACTIVITY);
+            mAmWmState.waitForActivityState(
+                    SHOW_WHEN_LOCKED_ATTR_REMOVE_ATTR_ACTIVITY, STATE_STOPPED);
+            assertSingleStartAndStop(SHOW_WHEN_LOCKED_ATTR_REMOVE_ATTR_ACTIVITY);
+        }
     }
 
     @Test
-    public void testNewActivityDuringOccludedWithAttr() {
-        final LockScreenSession lockScreenSession = createManagedLockScreenSession();
-        launchActivity(SHOW_WHEN_LOCKED_ATTR_ACTIVITY);
-        lockScreenSession.gotoKeyguard(SHOW_WHEN_LOCKED_ATTR_ACTIVITY);
-        launchActivity(SHOW_WHEN_LOCKED_WITH_DIALOG_ACTIVITY);
-        mAmWmState.computeState(SHOW_WHEN_LOCKED_WITH_DIALOG_ACTIVITY);
-        assertEquals("Picked wrong transition", TRANSIT_ACTIVITY_OPEN,
-                mAmWmState.getWmState().getDefaultDisplayLastTransition());
+    public void testNewActivityDuringOccludedWithAttr() throws Exception {
+        try (final LockScreenSession lockScreenSession = new LockScreenSession()) {
+            launchActivity(SHOW_WHEN_LOCKED_ATTR_ACTIVITY);
+            lockScreenSession.gotoKeyguard(SHOW_WHEN_LOCKED_ATTR_ACTIVITY);
+            launchActivity(SHOW_WHEN_LOCKED_WITH_DIALOG_ACTIVITY);
+            mAmWmState.computeState(SHOW_WHEN_LOCKED_WITH_DIALOG_ACTIVITY);
+            assertEquals("Picked wrong transition", TRANSIT_ACTIVITY_OPEN,
+                    mAmWmState.getWmState().getDefaultDisplayLastTransition());
+        }
     }
 }
diff --git a/tests/framework/base/windowmanager/src/android/server/wm/MockImeHelper.java b/tests/framework/base/windowmanager/src/android/server/wm/MockImeHelper.java
deleted file mode 100644
index 2191164..0000000
--- a/tests/framework/base/windowmanager/src/android/server/wm/MockImeHelper.java
+++ /dev/null
@@ -1,39 +0,0 @@
-/*
- * Copyright (C) 2019 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 android.server.wm;
-
-import com.android.cts.mockime.MockImeSession;
-
-/**
- * Centralizes the creation of {@link MockImeSession}. This class ins't placed in utility group
- * because only this package uses it.
- */
-public class MockImeHelper {
-
-    /**
-     * Leverage MockImeSession to ensure at least an IME exists as default.
-     *
-     * @see ObjectTracker#manage(AutoCloseable)
-     */
-    public static MockImeSession createManagedMockImeSession(ActivityManagerTestBase base) {
-        try {
-            return base.mObjectTracker.manage(MockImeSession.create(base.mContext));
-        } catch (Exception e) {
-            throw new RuntimeException("Failed to create MockImeSession", e);
-        }
-    }
-}
diff --git a/tests/framework/base/windowmanager/src/android/server/wm/MultiDisplayActivityLaunchTests.java b/tests/framework/base/windowmanager/src/android/server/wm/MultiDisplayActivityLaunchTests.java
index 4ba740c..8f5906b 100644
--- a/tests/framework/base/windowmanager/src/android/server/wm/MultiDisplayActivityLaunchTests.java
+++ b/tests/framework/base/windowmanager/src/android/server/wm/MultiDisplayActivityLaunchTests.java
@@ -68,6 +68,7 @@
 import android.server.wm.ActivityManagerState.ActivityStack;
 import android.server.wm.CommandSession.ActivitySession;
 import android.server.wm.CommandSession.SizeInfo;
+import android.util.SparseArray;
 
 import com.android.compatibility.common.util.SystemUtil;
 
@@ -111,30 +112,32 @@
      * Tests launching an assistant activity on virtual display.
      */
     @Test
-    public void testLaunchAssistantActivityOnSecondaryDisplay() {
+    public void testLaunchAssistantActivityOnSecondaryDisplay() throws Exception {
         validateActivityLaunchOnNewDisplay(ACTIVITY_TYPE_ASSISTANT);
     }
 
-    private void validateActivityLaunchOnNewDisplay(int activityType) {
-        // Create new virtual display.
-        final ActivityDisplay newDisplay = createManagedVirtualDisplaySession().createDisplay();
+    private void validateActivityLaunchOnNewDisplay(int activityType) throws Exception {
+        try (final VirtualDisplaySession virtualDisplaySession = new VirtualDisplaySession()) {
+            // Create new virtual display.
+            final ActivityDisplay newDisplay = virtualDisplaySession.createDisplay();
 
-        // Launch activity on new secondary display.
-        separateTestJournal();
-        getLaunchActivityBuilder().setUseInstrumentation().setWithShellPermission(true)
-                .setTargetActivity(TEST_ACTIVITY).setNewTask(true)
-                .setMultipleTask(true).setActivityType(activityType)
-                .setDisplayId(newDisplay.mId).execute();
-        waitAndAssertTopResumedActivity(TEST_ACTIVITY, newDisplay.mId,
-                "Activity launched on secondary display must be focused and on top");
+            // Launch activity on new secondary display.
+            separateTestJournal();
+            getLaunchActivityBuilder().setUseInstrumentation().setWithShellPermission(true)
+                    .setTargetActivity(TEST_ACTIVITY).setNewTask(true)
+                    .setMultipleTask(true).setActivityType(activityType)
+                    .setDisplayId(newDisplay.mId).execute();
+            waitAndAssertTopResumedActivity(TEST_ACTIVITY, newDisplay.mId,
+                    "Activity launched on secondary display must be focused and on top");
 
-        // Check that activity config corresponds to display config.
-        final SizeInfo reportedSizes = getLastReportedSizesForActivity(TEST_ACTIVITY);
-        assertEquals("Activity launched on secondary display must have proper configuration",
-                CUSTOM_DENSITY_DPI, reportedSizes.densityDpi);
+            // Check that activity config corresponds to display config.
+            final SizeInfo reportedSizes = getLastReportedSizesForActivity(TEST_ACTIVITY);
+            assertEquals("Activity launched on secondary display must have proper configuration",
+                    CUSTOM_DENSITY_DPI, reportedSizes.densityDpi);
 
-        assertEquals("Top activity must have correct activity type", activityType,
-                mAmWmState.getAmState().getFrontStackActivityType(newDisplay.mId));
+            assertEquals("Top activity must have correct activity type", activityType,
+                    mAmWmState.getAmState().getFrontStackActivityType(newDisplay.mId));
+        }
     }
 
     /**
@@ -161,31 +164,32 @@
      * Tests launching an existing activity from an activity that resided on secondary display.
      */
     @Test
-    public void testLaunchActivityFromSecondaryDisplay() {
+    public void testLaunchActivityFromSecondaryDisplay() throws Exception {
         getLaunchActivityBuilder().setUseInstrumentation()
                 .setTargetActivity(TEST_ACTIVITY).setNewTask(true)
                 .setDisplayId(DEFAULT_DISPLAY).execute();
 
-        final ActivityDisplay newDisplay = createManagedVirtualDisplaySession()
-                .setSimulateDisplay(true)
-                .createDisplay();
-        final int newDisplayId = newDisplay.mId;
+        try (final VirtualDisplaySession virtualDisplaySession = new VirtualDisplaySession()) {
+            final ActivityDisplay newDisplay =
+                    virtualDisplaySession.setSimulateDisplay(true).createDisplay();
+            final int newDisplayId = newDisplay.mId;
 
-        getLaunchActivityBuilder().setUseInstrumentation()
-                .setTargetActivity(BROADCAST_RECEIVER_ACTIVITY).setNewTask(true)
-                .setDisplayId(newDisplayId).execute();
-        waitAndAssertTopResumedActivity(BROADCAST_RECEIVER_ACTIVITY, newDisplay.mId,
-                "Activity should be resumed on secondary display");
+            getLaunchActivityBuilder().setUseInstrumentation()
+                    .setTargetActivity(BROADCAST_RECEIVER_ACTIVITY).setNewTask(true)
+                    .setDisplayId(newDisplayId).execute();
+            waitAndAssertTopResumedActivity(BROADCAST_RECEIVER_ACTIVITY, newDisplay.mId,
+                    "Activity should be resumed on secondary display");
 
-        mBroadcastActionTrigger.launchActivityNewTask(getActivityName(TEST_ACTIVITY));
-        waitAndAssertTopResumedActivity(TEST_ACTIVITY, DEFAULT_DISPLAY,
-                "Activity should be the top resumed on default display");
+            mBroadcastActionTrigger.launchActivityNewTask(getActivityName(TEST_ACTIVITY));
+            waitAndAssertTopResumedActivity(TEST_ACTIVITY, DEFAULT_DISPLAY,
+                    "Activity should be the top resumed on default display");
 
-        getLaunchActivityBuilder().setUseInstrumentation()
-                .setTargetActivity(TEST_ACTIVITY).setNewTask(true)
-                .setDisplayId(newDisplayId).execute();
-        waitAndAssertTopResumedActivity(TEST_ACTIVITY, newDisplay.mId,
-                "Activity should be resumed on secondary display");
+            getLaunchActivityBuilder().setUseInstrumentation()
+                    .setTargetActivity(TEST_ACTIVITY).setNewTask(true)
+                    .setDisplayId(newDisplayId).execute();
+            waitAndAssertTopResumedActivity(TEST_ACTIVITY, newDisplay.mId,
+                    "Activity should be resumed on secondary display");
+        }
     }
 
     /**
@@ -193,33 +197,34 @@
      * display is off.
      */
     @Test
-    public void testLaunchExternalDisplayActivityWhilePrimaryOff() {
+    public void testLaunchExternalDisplayActivityWhilePrimaryOff() throws Exception {
         // Launch something on the primary display so we know there is a resumed activity there
         launchActivity(RESIZEABLE_ACTIVITY);
         waitAndAssertTopResumedActivity(RESIZEABLE_ACTIVITY, DEFAULT_DISPLAY,
                 "Activity launched on primary display must be resumed");
 
-        final PrimaryDisplayStateSession displayStateSession =
-                mObjectTracker.manage(new PrimaryDisplayStateSession());
-        final ExternalDisplaySession externalDisplaySession = createManagedExternalDisplaySession();
-        displayStateSession.turnScreenOff();
+        try (final PrimaryDisplayStateSession displayStateSession =
+                     new PrimaryDisplayStateSession();
+             final ExternalDisplaySession externalDisplaySession = new ExternalDisplaySession()) {
+            displayStateSession.turnScreenOff();
 
-        // Make sure there is no resumed activity when the primary display is off
-        waitAndAssertActivityState(RESIZEABLE_ACTIVITY, STATE_STOPPED,
-                "Activity launched on primary display must be stopped after turning off");
-        assertEquals("Unexpected resumed activity",
-                0, mAmWmState.getAmState().getResumedActivitiesCount());
+            // Make sure there is no resumed activity when the primary display is off
+            waitAndAssertActivityState(RESIZEABLE_ACTIVITY, STATE_STOPPED,
+                    "Activity launched on primary display must be stopped after turning off");
+            assertEquals("Unexpected resumed activity",
+                    0, mAmWmState.getAmState().getResumedActivitiesCount());
 
-        final ActivityDisplay newDisplay = externalDisplaySession
-                .setCanShowWithInsecureKeyguard(true).createVirtualDisplay();
+            final ActivityDisplay newDisplay = externalDisplaySession
+                    .setCanShowWithInsecureKeyguard(true).createVirtualDisplay();
 
-        launchActivityOnDisplay(TEST_ACTIVITY, newDisplay.mId);
+            launchActivityOnDisplay(TEST_ACTIVITY, newDisplay.mId);
 
-        // Check that the test activity is resumed on the external display
-        waitAndAssertTopResumedActivity(TEST_ACTIVITY, newDisplay.mId,
-                "Activity launched on external display must be resumed");
-        mAmWmState.assertFocusedAppOnDisplay("App on default display must still be focused",
-                RESIZEABLE_ACTIVITY, DEFAULT_DISPLAY);
+            // Check that the test activity is resumed on the external display
+            waitAndAssertTopResumedActivity(TEST_ACTIVITY, newDisplay.mId,
+                    "Activity launched on external display must be resumed");
+            mAmWmState.assertFocusedAppOnDisplay("App on default display must still be focused",
+                    RESIZEABLE_ACTIVITY, DEFAULT_DISPLAY);
+        }
     }
 
     /**
@@ -227,64 +232,66 @@
      * virtual display with correct configuration.
      */
     @Test
-    public void testLaunchNonResizeableActivityOnSecondaryDisplay() {
-        // Create new virtual display.
-        final ActivityDisplay newDisplay = createManagedVirtualDisplaySession().createDisplay();
+    public void testLaunchNonResizeableActivityOnSecondaryDisplay() throws Exception {
+        try (final VirtualDisplaySession virtualDisplaySession = new VirtualDisplaySession()) {
+            // Create new virtual display.
+            final ActivityDisplay newDisplay = virtualDisplaySession.createDisplay();
 
-        // Launch activity on new secondary display.
-        launchActivityOnDisplay(NON_RESIZEABLE_ACTIVITY, newDisplay.mId);
+            // Launch activity on new secondary display.
+            launchActivityOnDisplay(NON_RESIZEABLE_ACTIVITY, newDisplay.mId);
 
-        waitAndAssertTopResumedActivity(NON_RESIZEABLE_ACTIVITY, newDisplay.mId,
-                "Activity requested to launch on secondary display must be focused");
+            waitAndAssertTopResumedActivity(NON_RESIZEABLE_ACTIVITY, newDisplay.mId,
+                    "Activity requested to launch on secondary display must be focused");
 
-        final Configuration taskConfig = mAmWmState.getAmState()
-                .getTaskByActivity(NON_RESIZEABLE_ACTIVITY).mFullConfiguration;
-        final Configuration displayConfig = mAmWmState.getWmState()
-                .getDisplay(newDisplay.mId).mFullConfiguration;
+            final Configuration taskConfig = mAmWmState.getAmState()
+                    .getTaskByActivity(NON_RESIZEABLE_ACTIVITY).mFullConfiguration;
+            final Configuration displayConfig = mAmWmState.getWmState()
+                    .getDisplay(newDisplay.mId).mFullConfiguration;
 
-        // Check that activity config corresponds to display config.
-        assertEquals("Activity launched on secondary display must have proper configuration",
-                taskConfig.densityDpi, displayConfig.densityDpi);
+            // Check that activity config corresponds to display config.
+            assertEquals("Activity launched on secondary display must have proper configuration",
+                    taskConfig.densityDpi, displayConfig.densityDpi);
 
-        assertEquals("Activity launched on secondary display must have proper configuration",
-                taskConfig.windowConfiguration.getBounds(),
-                displayConfig.windowConfiguration.getBounds());
+            assertEquals("Activity launched on secondary display must have proper configuration",
+                    taskConfig.windowConfiguration.getBounds(),
+                    displayConfig.windowConfiguration.getBounds());
+        }
     }
 
     /**
      * Tests successfully moving a non-resizeable activity to a virtual display.
      */
     @Test
-    public void testMoveNonResizeableActivityToSecondaryDisplay() {
-        final VirtualDisplayLauncher virtualLauncher =
-                mObjectTracker.manage(new VirtualDisplayLauncher());
-        // Create new virtual display.
-        final ActivityDisplay newDisplay = virtualLauncher.createDisplay();
-        // Launch a non-resizeable activity on a primary display.
-        final ActivitySession nonResizeableSession = virtualLauncher.launchActivity(
-                builder -> builder.setTargetActivity(NON_RESIZEABLE_ACTIVITY).setNewTask(true));
+    public void testMoveNonResizeableActivityToSecondaryDisplay() throws Exception {
+        try (final VirtualDisplayLauncher virtualLauncher = new VirtualDisplayLauncher()) {
+            // Create new virtual display.
+            final ActivityDisplay newDisplay = virtualLauncher.createDisplay();
+            // Launch a non-resizeable activity on a primary display.
+            final ActivitySession nonResizeableSession = virtualLauncher.launchActivity(
+                    builder -> builder.setTargetActivity(NON_RESIZEABLE_ACTIVITY).setNewTask(true));
 
-        // Launch a resizeable activity on new secondary display to create a new stack there.
-        virtualLauncher.launchActivityOnDisplay(RESIZEABLE_ACTIVITY, newDisplay);
-        final int externalFrontStackId = mAmWmState.getAmState()
-                .getFrontStackId(newDisplay.mId);
+            // Launch a resizeable activity on new secondary display to create a new stack there.
+            virtualLauncher.launchActivityOnDisplay(RESIZEABLE_ACTIVITY, newDisplay);
+            final int externalFrontStackId = mAmWmState.getAmState()
+                    .getFrontStackId(newDisplay.mId);
 
-        // Clear lifecycle callback history before moving the activity so the later verification
-        // can get the callbacks which are related to the reparenting.
-        nonResizeableSession.takeCallbackHistory();
+            // Clear lifecycle callback history before moving the activity so the later verification
+            // can get the callbacks which are related to the reparenting.
+            nonResizeableSession.takeCallbackHistory();
 
-        // Try to move the non-resizeable activity to the top of stack on secondary display.
-        moveActivityToStack(NON_RESIZEABLE_ACTIVITY, externalFrontStackId);
-        // Wait for a while to check that it will move.
-        assertTrue("Non-resizeable activity should be moved",
-                mAmWmState.waitForWithAmState(
-                        state -> newDisplay.mId == state
-                                .getDisplayByActivity(NON_RESIZEABLE_ACTIVITY),
-                        "seeing if activity won't be moved"));
+            // Try to move the non-resizeable activity to the top of stack on secondary display.
+            moveActivityToStack(NON_RESIZEABLE_ACTIVITY, externalFrontStackId);
+            // Wait for a while to check that it will move.
+            assertTrue("Non-resizeable activity should be moved",
+                    mAmWmState.waitForWithAmState(
+                            state -> newDisplay.mId == state
+                                    .getDisplayByActivity(NON_RESIZEABLE_ACTIVITY),
+                            "seeing if activity won't be moved"));
 
-        waitAndAssertTopResumedActivity(NON_RESIZEABLE_ACTIVITY, newDisplay.mId,
-                "The moved non-resizeable activity must be focused");
-        assertActivityLifecycle(nonResizeableSession, true /* relaunched */);
+            waitAndAssertTopResumedActivity(NON_RESIZEABLE_ACTIVITY, newDisplay.mId,
+                    "The moved non-resizeable activity must be focused");
+            assertActivityLifecycle(nonResizeableSession, true /* relaunched */);
+        }
     }
 
     /**
@@ -292,21 +299,22 @@
      * land on the secondary display based on the resizeability of the root activity of the task.
      */
     @Test
-    public void testLaunchNonResizeableActivityFromSecondaryDisplaySameTask() {
-        // Create new simulated display.
-        final ActivityDisplay newDisplay = createManagedVirtualDisplaySession()
-                .setSimulateDisplay(true)
-                .createDisplay();
+    public void testLaunchNonResizeableActivityFromSecondaryDisplaySameTask() throws Exception {
+        try (final VirtualDisplaySession virtualDisplaySession = new VirtualDisplaySession()) {
+            // Create new simulated display.
+            final ActivityDisplay newDisplay = virtualDisplaySession.setSimulateDisplay(true)
+                    .createDisplay();
 
-        // Launch activity on new secondary display.
-        launchActivityOnDisplay(BROADCAST_RECEIVER_ACTIVITY, newDisplay.mId);
-        waitAndAssertTopResumedActivity(BROADCAST_RECEIVER_ACTIVITY, newDisplay.mId,
-                "Activity launched on secondary display must be focused");
+            // Launch activity on new secondary display.
+            launchActivityOnDisplay(BROADCAST_RECEIVER_ACTIVITY, newDisplay.mId);
+            waitAndAssertTopResumedActivity(BROADCAST_RECEIVER_ACTIVITY, newDisplay.mId,
+                    "Activity launched on secondary display must be focused");
 
-        // Launch non-resizeable activity from secondary display.
-        mBroadcastActionTrigger.launchActivityNewTask(getActivityName(NON_RESIZEABLE_ACTIVITY));
-        waitAndAssertTopResumedActivity(NON_RESIZEABLE_ACTIVITY, newDisplay.mId,
-                "Launched activity must be on the secondary display and resumed");
+            // Launch non-resizeable activity from secondary display.
+            mBroadcastActionTrigger.launchActivityNewTask(getActivityName(NON_RESIZEABLE_ACTIVITY));
+            waitAndAssertTopResumedActivity(NON_RESIZEABLE_ACTIVITY, newDisplay.mId,
+                    "Launched activity must be on the secondary display and resumed");
+        }
     }
 
     /**
@@ -314,34 +322,41 @@
      * there. It must land on the display as its caller.
      */
     @Test
-    public void testLaunchNonResizeableActivityFromSecondaryDisplayNewTask() {
-        // Create new virtual display.
-        final ActivityDisplay newDisplay = createManagedVirtualDisplaySession().createDisplay();
+    public void testLaunchNonResizeableActivityFromSecondaryDisplayNewTask() throws Exception {
+        try (final VirtualDisplaySession virtualDisplaySession = new VirtualDisplaySession()) {
+            // Create new virtual display.
+            final ActivityDisplay newDisplay = virtualDisplaySession.createDisplay();
 
-        // Launch activity on new secondary display.
-        launchActivityOnDisplay(LAUNCHING_ACTIVITY, newDisplay.mId);
-        waitAndAssertTopResumedActivity(LAUNCHING_ACTIVITY, newDisplay.mId,
-                "Activity launched on secondary display must be focused");
+            // Launch activity on new secondary display.
+            launchActivityOnDisplay(LAUNCHING_ACTIVITY, newDisplay.mId);
+            waitAndAssertTopResumedActivity(LAUNCHING_ACTIVITY, newDisplay.mId,
+                    "Activity launched on secondary display must be focused");
 
-        // Launch non-resizeable activity from secondary display in a new task.
-        getLaunchActivityBuilder().setTargetActivity(NON_RESIZEABLE_ACTIVITY)
-                .setNewTask(true).setMultipleTask(true).execute();
+            // Launch non-resizeable activity from secondary display in a new task.
+            getLaunchActivityBuilder().setTargetActivity(NON_RESIZEABLE_ACTIVITY)
+                    .setNewTask(true).setMultipleTask(true).execute();
 
-        mAmWmState.waitForActivityState(NON_RESIZEABLE_ACTIVITY, STATE_RESUMED);
+            mAmWmState.waitForActivityState(NON_RESIZEABLE_ACTIVITY, STATE_RESUMED);
 
-        // Check that non-resizeable activity is on the same display.
-        final int newFrontStackId = mAmWmState.getAmState().getFocusedStackId();
-        final ActivityStack newFrontStack = mAmWmState.getAmState().getStackById(newFrontStackId);
-        assertTrue("Launched activity must be on the same display",
-                newDisplay.mId == newFrontStack.mDisplayId);
-        assertEquals("Launched activity must be resumed",
-                getActivityName(NON_RESIZEABLE_ACTIVITY),
-                newFrontStack.mResumedActivity);
-        mAmWmState.assertFocusedStack(
-                "Top stack must be the one with just launched activity",
-                newFrontStackId);
-        assertBothDisplaysHaveResumedActivities(pair(newDisplay.mId, LAUNCHING_ACTIVITY),
-                pair(newFrontStack.mDisplayId, NON_RESIZEABLE_ACTIVITY));
+            // Check that non-resizeable activity is on the same display.
+            final int newFrontStackId = mAmWmState.getAmState().getFocusedStackId();
+            final ActivityStack newFrontStack =
+                    mAmWmState.getAmState().getStackById(newFrontStackId);
+            assertTrue("Launched activity must be on the same display",
+                    newDisplay.mId == newFrontStack.mDisplayId);
+            assertEquals("Launched activity must be resumed",
+                    getActivityName(NON_RESIZEABLE_ACTIVITY),
+                    newFrontStack.mResumedActivity);
+            mAmWmState.assertFocusedStack(
+                    "Top stack must be the one with just launched activity",
+                    newFrontStackId);
+            mAmWmState.assertResumedActivities("Both displays must have resumed activities",
+                    new SparseArray<ComponentName>(){{
+                        put(newDisplay.mId, LAUNCHING_ACTIVITY);
+                        put(newFrontStack.mDisplayId, NON_RESIZEABLE_ACTIVITY);
+                    }}
+            );
+        }
     }
 
     /**
@@ -350,24 +365,30 @@
      * primary display.
      */
     @Test
-    public void testConsequentLaunchActivity() {
-        // Create new virtual display.
-        final ActivityDisplay newDisplay = createManagedVirtualDisplaySession().createDisplay();
+    public void testConsequentLaunchActivity() throws Exception {
+        try (final VirtualDisplaySession virtualDisplaySession = new VirtualDisplaySession()) {
+            // Create new virtual display.
+            final ActivityDisplay newDisplay = virtualDisplaySession.createDisplay();
 
-        // Launch activity on new secondary display.
-        launchActivityOnDisplay(TEST_ACTIVITY, newDisplay.mId);
+            // Launch activity on new secondary display.
+            launchActivityOnDisplay(TEST_ACTIVITY, newDisplay.mId);
 
-        waitAndAssertTopResumedActivity(TEST_ACTIVITY, newDisplay.mId,
-                "Activity launched on secondary display must be on top");
+            waitAndAssertTopResumedActivity(TEST_ACTIVITY, newDisplay.mId,
+                    "Activity launched on secondary display must be on top");
 
-        // Launch second activity without specifying display.
-        launchActivity(LAUNCHING_ACTIVITY);
+            // Launch second activity without specifying display.
+            launchActivity(LAUNCHING_ACTIVITY);
 
-        // Check that activity is launched in focused stack on primary display.
-        waitAndAssertTopResumedActivity(LAUNCHING_ACTIVITY, DEFAULT_DISPLAY,
-                "Launched activity must be focused");
-        assertBothDisplaysHaveResumedActivities(pair(newDisplay.mId, TEST_ACTIVITY),
-                pair(DEFAULT_DISPLAY, LAUNCHING_ACTIVITY));
+            // Check that activity is launched in focused stack on primary display.
+            waitAndAssertTopResumedActivity(LAUNCHING_ACTIVITY, DEFAULT_DISPLAY,
+                    "Launched activity must be focused");
+            mAmWmState.assertResumedActivities("Both displays must have resumed activities",
+                    new SparseArray<ComponentName>(){{
+                        put(newDisplay.mId, TEST_ACTIVITY);
+                        put(DEFAULT_DISPLAY, LAUNCHING_ACTIVITY);
+                    }}
+            );
+        }
     }
 
     /**
@@ -375,24 +396,25 @@
      * first one - it must appear on the secondary display, because it was launched from there.
      */
     @Test
-    public void testConsequentLaunchActivityFromSecondaryDisplay() {
-        // Create new simulated display.
-        final ActivityDisplay newDisplay = createManagedVirtualDisplaySession()
-                .setSimulateDisplay(true)
-                .createDisplay();
+    public void testConsequentLaunchActivityFromSecondaryDisplay() throws Exception {
+        try (final VirtualDisplaySession virtualDisplaySession = new VirtualDisplaySession()) {
+            // Create new simulated display.
+            final ActivityDisplay newDisplay = virtualDisplaySession.setSimulateDisplay(true)
+                    .createDisplay();
 
-        // Launch activity on new secondary display.
-        launchActivityOnDisplay(LAUNCHING_ACTIVITY, newDisplay.mId);
+            // Launch activity on new secondary display.
+            launchActivityOnDisplay(LAUNCHING_ACTIVITY, newDisplay.mId);
 
-        waitAndAssertTopResumedActivity(LAUNCHING_ACTIVITY, newDisplay.mId,
-                "Activity launched on secondary display must be on top");
+            waitAndAssertTopResumedActivity(LAUNCHING_ACTIVITY, newDisplay.mId,
+                    "Activity launched on secondary display must be on top");
 
-        // Launch second activity from app on secondary display without specifying display id.
-        getLaunchActivityBuilder().setTargetActivity(TEST_ACTIVITY).execute();
+            // Launch second activity from app on secondary display without specifying display id.
+            getLaunchActivityBuilder().setTargetActivity(TEST_ACTIVITY).execute();
 
-        // Check that activity is launched in focused stack on external display.
-        waitAndAssertTopResumedActivity(TEST_ACTIVITY, newDisplay.mId,
-                "Launched activity must be on top");
+            // Check that activity is launched in focused stack on external display.
+            waitAndAssertTopResumedActivity(TEST_ACTIVITY, newDisplay.mId,
+                    "Launched activity must be on top");
+        }
     }
 
     /**
@@ -400,23 +422,25 @@
      * first one - it must appear on the secondary display, because it was launched from there.
      */
     @Test
-    public void testConsequentLaunchActivityFromVirtualDisplay() {
-        // Create new virtual display.
-        final ActivityDisplay newDisplay = createManagedVirtualDisplaySession().createDisplay();
+    public void testConsequentLaunchActivityFromVirtualDisplay() throws Exception {
+        try (final VirtualDisplaySession virtualDisplaySession = new VirtualDisplaySession()) {
+            // Create new virtual display.
+            final ActivityDisplay newDisplay = virtualDisplaySession.createDisplay();
 
-        // Launch activity on new secondary display.
-        launchActivityOnDisplay(LAUNCHING_ACTIVITY, newDisplay.mId);
+            // Launch activity on new secondary display.
+            launchActivityOnDisplay(LAUNCHING_ACTIVITY, newDisplay.mId);
 
-        waitAndAssertTopResumedActivity(LAUNCHING_ACTIVITY, newDisplay.mId,
-                "Activity launched on secondary display must be on top");
+            waitAndAssertTopResumedActivity(LAUNCHING_ACTIVITY, newDisplay.mId,
+                    "Activity launched on secondary display must be on top");
 
-        // Launch second activity from app on secondary display without specifying display id.
-        getLaunchActivityBuilder().setTargetActivity(TEST_ACTIVITY).execute();
-        mAmWmState.computeState(TEST_ACTIVITY);
+            // Launch second activity from app on secondary display without specifying display id.
+            getLaunchActivityBuilder().setTargetActivity(TEST_ACTIVITY).execute();
+            mAmWmState.computeState(TEST_ACTIVITY);
 
-        // Check that activity is launched in focused stack on external display.
-        waitAndAssertTopResumedActivity(TEST_ACTIVITY, newDisplay.mId,
-                "Launched activity must be on top");
+            // Check that activity is launched in focused stack on external display.
+            waitAndAssertTopResumedActivity(TEST_ACTIVITY, newDisplay.mId,
+                    "Launched activity must be on top");
+        }
     }
 
     /**
@@ -424,100 +448,108 @@
      * first one with specifying the target display - it must appear on the secondary display.
      */
     @Test
-    public void testConsequentLaunchActivityFromVirtualDisplayToTargetDisplay() {
-        // Create new virtual display.
-        final ActivityDisplay newDisplay = createManagedVirtualDisplaySession().createDisplay();
+    public void testConsequentLaunchActivityFromVirtualDisplayToTargetDisplay() throws Exception {
+        try (final VirtualDisplaySession virtualDisplaySession = new VirtualDisplaySession()) {
+            // Create new virtual display.
+            final ActivityDisplay newDisplay = virtualDisplaySession.createDisplay();
 
-        // Launch activity on new secondary display.
-        launchActivityOnDisplay(LAUNCHING_ACTIVITY, newDisplay.mId);
+            // Launch activity on new secondary display.
+            launchActivityOnDisplay(LAUNCHING_ACTIVITY, newDisplay.mId);
 
-        waitAndAssertTopResumedActivity(LAUNCHING_ACTIVITY, newDisplay.mId,
-                "Activity launched on secondary display must be on top");
+            waitAndAssertTopResumedActivity(LAUNCHING_ACTIVITY, newDisplay.mId,
+                    "Activity launched on secondary display must be on top");
 
-        // Launch second activity from app on secondary display specifying same display id.
-        getLaunchActivityBuilder()
-                .setTargetActivity(SECOND_ACTIVITY)
-                .setDisplayId(newDisplay.mId)
-                .execute();
+            // Launch second activity from app on secondary display specifying same display id.
+            getLaunchActivityBuilder()
+                    .setTargetActivity(SECOND_ACTIVITY)
+                    .setDisplayId(newDisplay.mId)
+                    .execute();
 
-        // Check that activity is launched in focused stack on external display.
-        waitAndAssertTopResumedActivity(SECOND_ACTIVITY, newDisplay.mId,
-                "Launched activity must be on top");
+            // Check that activity is launched in focused stack on external display.
+            waitAndAssertTopResumedActivity(SECOND_ACTIVITY, newDisplay.mId,
+                    "Launched activity must be on top");
 
-        // Launch other activity with different uid and check if it has launched successfully.
-        getLaunchActivityBuilder()
-                .setUseBroadcastReceiver(SECOND_LAUNCH_BROADCAST_RECEIVER,
-                        SECOND_LAUNCH_BROADCAST_ACTION)
-                .setDisplayId(newDisplay.mId)
-                .setTargetActivity(THIRD_ACTIVITY)
-                .execute();
+            // Launch other activity with different uid and check if it has launched successfully.
+            getLaunchActivityBuilder()
+                    .setUseBroadcastReceiver(SECOND_LAUNCH_BROADCAST_RECEIVER,
+                            SECOND_LAUNCH_BROADCAST_ACTION)
+                    .setDisplayId(newDisplay.mId)
+                    .setTargetActivity(THIRD_ACTIVITY)
+                    .execute();
 
-        // Check that activity is launched in focused stack on external display.
-        waitAndAssertTopResumedActivity(THIRD_ACTIVITY, newDisplay.mId,
-                "Launched activity must be on top");
+            // Check that activity is launched in focused stack on external display.
+            waitAndAssertTopResumedActivity(THIRD_ACTIVITY, newDisplay.mId,
+                    "Launched activity must be on top");
+        }
     }
 
     /**
      * Tests launching an activity to secondary display from activity on primary display.
      */
     @Test
-    public void testLaunchActivityFromAppToSecondaryDisplay() {
+    public void testLaunchActivityFromAppToSecondaryDisplay() throws Exception {
         // Start launching activity.
         launchActivity(LAUNCHING_ACTIVITY);
 
-        // Create new simulated display.
-        final ActivityDisplay newDisplay = createManagedVirtualDisplaySession()
-                .setSimulateDisplay(true)
-                .createDisplay();
+        try (final VirtualDisplaySession virtualDisplaySession = new VirtualDisplaySession()) {
+            // Create new simulated display.
+            final ActivityDisplay newDisplay = virtualDisplaySession.setSimulateDisplay(true)
+                    .createDisplay();
 
-        // Launch activity on secondary display from the app on primary display.
-        getLaunchActivityBuilder().setTargetActivity(TEST_ACTIVITY)
-                .setDisplayId(newDisplay.mId).execute();
+            // Launch activity on secondary display from the app on primary display.
+            getLaunchActivityBuilder().setTargetActivity(TEST_ACTIVITY)
+                    .setDisplayId(newDisplay.mId).execute();
 
-        // Check that activity is launched on external display.
-        waitAndAssertTopResumedActivity(TEST_ACTIVITY, newDisplay.mId,
-                "Activity launched on secondary display must be focused");
-        assertBothDisplaysHaveResumedActivities(pair(DEFAULT_DISPLAY, LAUNCHING_ACTIVITY),
-                pair(newDisplay.mId, TEST_ACTIVITY));
+            // Check that activity is launched on external display.
+            waitAndAssertTopResumedActivity(TEST_ACTIVITY, newDisplay.mId,
+                    "Activity launched on secondary display must be focused");
+            mAmWmState.assertResumedActivities("Both displays must have resumed activities",
+                    new SparseArray<ComponentName>(){{
+                        put(DEFAULT_DISPLAY, LAUNCHING_ACTIVITY);
+                        put(newDisplay.mId, TEST_ACTIVITY);
+                    }}
+            );
+        }
     }
 
     /** Tests that launching app from pending activity queue on external display is allowed. */
     @Test
-    public void testLaunchPendingActivityOnSecondaryDisplay() {
-        pressHomeButton();
-        // Create new simulated display.
-        final ActivityDisplay newDisplay = createManagedVirtualDisplaySession()
-                .setSimulateDisplay(true)
-                .createDisplay();
-        final Bundle bundle = ActivityOptions.makeBasic().
-                setLaunchDisplayId(newDisplay.mId).toBundle();
-        final Intent intent = new Intent(Intent.ACTION_VIEW)
-                .setComponent(SECOND_ACTIVITY)
-                .setFlags(FLAG_ACTIVITY_NEW_TASK | FLAG_ACTIVITY_MULTIPLE_TASK)
-                .putExtra(KEY_LAUNCH_ACTIVITY, true)
-                .putExtra(KEY_NEW_TASK, true);
-        mContext.startActivity(intent, bundle);
+    public void testLaunchPendingActivityOnSecondaryDisplay() throws Exception {
+        try (final VirtualDisplaySession virtualDisplaySession = new VirtualDisplaySession()) {
+            pressHomeButton();
+            // Create new simulated display.
+            final ActivityDisplay newDisplay = virtualDisplaySession.setSimulateDisplay(true)
+                    .createDisplay();
+            final Bundle bundle = ActivityOptions.makeBasic().
+                    setLaunchDisplayId(newDisplay.mId).toBundle();
+            final Intent intent = new Intent(Intent.ACTION_VIEW)
+                    .setComponent(SECOND_ACTIVITY)
+                    .setFlags(FLAG_ACTIVITY_NEW_TASK | FLAG_ACTIVITY_MULTIPLE_TASK)
+                    .putExtra(KEY_LAUNCH_ACTIVITY, true)
+                    .putExtra(KEY_NEW_TASK, true);
+            mContext.startActivity(intent, bundle);
 
-        // If home key was pressed, stopAppSwitches will be called.
-        // Since this test case is not start activity from shell, it won't grant
-        // STOP_APP_SWITCHES and this activity should be put into pending activity queue
-        // and this activity should been launched after
-        // ActivityTaskManagerService.APP_SWITCH_DELAY_TIME
-        mAmWmState.waitForPendingActivityContain(SECOND_ACTIVITY);
-        // If the activity is not pending, skip this test.
-        mAmWmState.assumePendingActivityContain(SECOND_ACTIVITY);
-        // In order to speed up test case without waiting for APP_SWITCH_DELAY_TIME, we launch
-        // another activity with LaunchActivityBuilder, in this way the activity can be start
-        // directly and also trigger pending activity to be launched.
-        getLaunchActivityBuilder()
-                .setTargetActivity(THIRD_ACTIVITY)
-                .execute();
-        mAmWmState.waitForValidState(SECOND_ACTIVITY);
-        waitAndAssertTopResumedActivity(THIRD_ACTIVITY, DEFAULT_DISPLAY,
-                "Top activity must be the newly launched one");
-        mAmWmState.assertVisibility(SECOND_ACTIVITY, true);
-        assertEquals("Activity launched by app on secondary display must be on that display",
-                newDisplay.mId, mAmWmState.getAmState().getDisplayByActivity(SECOND_ACTIVITY));
+            // If home key was pressed, stopAppSwitches will be called.
+            // Since this test case is not start activity from shell, it won't grant
+            // STOP_APP_SWITCHES and this activity should be put into pending activity queue
+            // and this activity should been launched after
+            // ActivityTaskManagerService.APP_SWITCH_DELAY_TIME
+            mAmWmState.waitForPendingActivityContain(SECOND_ACTIVITY);
+            // If the activity is not pending, skip this test.
+            mAmWmState.assumePendingActivityContain(SECOND_ACTIVITY);
+            // In order to speed up test case without waiting for APP_SWITCH_DELAY_TIME, we launch
+            // another activity with LaunchActivityBuilder, in this way the activity can be start
+            // directly and also trigger pending activity to be launched.
+            getLaunchActivityBuilder()
+                    .setTargetActivity(THIRD_ACTIVITY)
+                    .execute();
+            mAmWmState.waitForValidState(SECOND_ACTIVITY);
+            waitAndAssertTopResumedActivity(THIRD_ACTIVITY, DEFAULT_DISPLAY,
+                    "Top activity must be the newly launched one");
+            mAmWmState.assertVisibility(SECOND_ACTIVITY, true);
+            assertEquals("Activity launched by app on secondary display must be on that display",
+                    newDisplay.mId, mAmWmState.getAmState().getDisplayByActivity(SECOND_ACTIVITY));
+        }
     }
 
     /**
@@ -525,46 +557,49 @@
      * matching task on some other display - that task will moved to the target display.
      */
     @Test
-    public void testMoveToDisplayOnLaunch() {
+    public void testMoveToDisplayOnLaunch() throws Exception {
         // Launch activity with unique affinity, so it will the only one in its task.
         launchActivity(LAUNCHING_ACTIVITY);
 
-        // Create new virtual display.
-        final ActivityDisplay newDisplay = createManagedVirtualDisplaySession().createDisplay();
-        mAmWmState.assertVisibility(VIRTUAL_DISPLAY_ACTIVITY, true /* visible */);
-        // Launch something to that display so that a new stack is created. We need this to be
-        // able to compare task numbers in stacks later.
-        launchActivityOnDisplay(RESIZEABLE_ACTIVITY, newDisplay.mId);
-        mAmWmState.assertVisibility(RESIZEABLE_ACTIVITY, true /* visible */);
+        try (final VirtualDisplaySession virtualDisplaySession = new VirtualDisplaySession()) {
+            // Create new virtual display.
+            final ActivityDisplay newDisplay = virtualDisplaySession.createDisplay();
+            mAmWmState.assertVisibility(VIRTUAL_DISPLAY_ACTIVITY, true /* visible */);
+            // Launch something to that display so that a new stack is created. We need this to be
+            // able to compare task numbers in stacks later.
+            launchActivityOnDisplay(RESIZEABLE_ACTIVITY, newDisplay.mId);
+            mAmWmState.assertVisibility(RESIZEABLE_ACTIVITY, true /* visible */);
 
-        final int stackNum = mAmWmState.getAmState().getDisplay(DEFAULT_DISPLAY).mStacks.size();
-        final int stackNumOnSecondary = mAmWmState.getAmState()
-                .getDisplay(newDisplay.mId).mStacks.size();
+            final int stackNum = mAmWmState.getAmState().getDisplay(DEFAULT_DISPLAY)
+                    .mStacks.size();
+            final int stackNumOnSecondary = mAmWmState.getAmState()
+                    .getDisplay(newDisplay.mId).mStacks.size();
 
-        // Launch activity on new secondary display.
-        // Using custom command here, because normally we add flags
-        // {@link Intent#FLAG_ACTIVITY_NEW_TASK} and {@link Intent#FLAG_ACTIVITY_MULTIPLE_TASK}
-        // when launching on some specific display. We don't do it here as we want an existing
-        // task to be used.
-        final String launchCommand = "am start -n " + getActivityName(LAUNCHING_ACTIVITY)
-                + " --display " + newDisplay.mId;
-        executeShellCommand(launchCommand);
+            // Launch activity on new secondary display.
+            // Using custom command here, because normally we add flags
+            // {@link Intent#FLAG_ACTIVITY_NEW_TASK} and {@link Intent#FLAG_ACTIVITY_MULTIPLE_TASK}
+            // when launching on some specific display. We don't do it here as we want an existing
+            // task to be used.
+            final String launchCommand = "am start -n " + getActivityName(LAUNCHING_ACTIVITY)
+                    + " --display " + newDisplay.mId;
+            executeShellCommand(launchCommand);
 
-        // Check that activity is brought to front.
-        waitAndAssertTopResumedActivity(LAUNCHING_ACTIVITY, newDisplay.mId,
-                "Existing task must be brought to front");
+            // Check that activity is brought to front.
+            waitAndAssertTopResumedActivity(LAUNCHING_ACTIVITY, newDisplay.mId,
+                    "Existing task must be brought to front");
 
-        // Check that task has moved from primary display to secondary.
-        // Since it is 1-to-1 relationship between task and stack for standard type &
-        // fullscreen activity, we check the number of stacks here
-        final int stackNumFinal = mAmWmState.getAmState().getDisplay(DEFAULT_DISPLAY)
-                .mStacks.size();
-        assertEquals("Stack number in default stack must be decremented.", stackNum - 1,
-                stackNumFinal);
-        final int stackNumFinalOnSecondary = mAmWmState.getAmState()
-                .getDisplay(newDisplay.mId).mStacks.size();
-        assertEquals("Stack number on external display must be incremented.",
-                stackNumOnSecondary + 1, stackNumFinalOnSecondary);
+            // Check that task has moved from primary display to secondary.
+            // Since it is 1-to-1 relationship between task and stack for standard type &
+            // fullscreen activity, we check the number of stacks here
+            final int stackNumFinal = mAmWmState.getAmState().getDisplay(DEFAULT_DISPLAY)
+                    .mStacks.size();
+            assertEquals("Stack number in default stack must be decremented.", stackNum - 1,
+                    stackNumFinal);
+            final int stackNumFinalOnSecondary = mAmWmState.getAmState()
+                    .getDisplay(newDisplay.mId).mStacks.size();
+            assertEquals("Stack number on external display must be incremented.",
+                    stackNumOnSecondary + 1, stackNumFinalOnSecondary);
+        }
     }
 
     /**
@@ -572,36 +607,38 @@
      * matching task on some other display - that task will moved to the target display.
      */
     @Test
-    public void testMoveToEmptyDisplayOnLaunch() {
+    public void testMoveToEmptyDisplayOnLaunch() throws Exception {
         // Launch activity with unique affinity, so it will the only one in its task. And choose
         // resizeable activity to prevent the test activity be relaunched when launch it to another
         // display, which may affect on this test case.
         launchActivity(RESIZEABLE_ACTIVITY);
 
-        // Create new virtual display.
-        final ActivityDisplay newDisplay = createManagedVirtualDisplaySession().createDisplay();
-        mAmWmState.assertVisibility(VIRTUAL_DISPLAY_ACTIVITY, true /* visible */);
+        try (final VirtualDisplaySession virtualDisplaySession = new VirtualDisplaySession()) {
+            // Create new virtual display.
+            final ActivityDisplay newDisplay = virtualDisplaySession.createDisplay();
+            mAmWmState.assertVisibility(VIRTUAL_DISPLAY_ACTIVITY, true /* visible */);
 
-        final int stackNum = mAmWmState.getAmState().getDisplay(DEFAULT_DISPLAY).mStacks.size();
+            final int stackNum = mAmWmState.getAmState().getDisplay(DEFAULT_DISPLAY).mStacks.size();
 
-        // Launch activity on new secondary display.
-        // Using custom command here, because normally we add flags
-        // {@link Intent#FLAG_ACTIVITY_NEW_TASK} and {@link Intent#FLAG_ACTIVITY_MULTIPLE_TASK}
-        // when launching on some specific display. We don't do it here as we want an existing
-        // task to be used.
-        final String launchCommand = "am start -n " + getActivityName(RESIZEABLE_ACTIVITY)
-                + " --display " + newDisplay.mId;
-        executeShellCommand(launchCommand);
+            // Launch activity on new secondary display.
+            // Using custom command here, because normally we add flags
+            // {@link Intent#FLAG_ACTIVITY_NEW_TASK} and {@link Intent#FLAG_ACTIVITY_MULTIPLE_TASK}
+            // when launching on some specific display. We don't do it here as we want an existing
+            // task to be used.
+            final String launchCommand = "am start -n " + getActivityName(RESIZEABLE_ACTIVITY)
+                    + " --display " + newDisplay.mId;
+            executeShellCommand(launchCommand);
 
-        // Check that activity is brought to front.
-        waitAndAssertTopResumedActivity(RESIZEABLE_ACTIVITY, newDisplay.mId,
-                "Existing task must be brought to front");
+            // Check that activity is brought to front.
+            waitAndAssertTopResumedActivity(RESIZEABLE_ACTIVITY, newDisplay.mId,
+                    "Existing task must be brought to front");
 
-        // Check that task has moved from primary display to secondary.
-        final int stackNumFinal = mAmWmState.getAmState().getDisplay(DEFAULT_DISPLAY)
-                .mStacks.size();
-        assertEquals("Stack number in default stack must be decremented.", stackNum - 1,
-                stackNumFinal);
+            // Check that task has moved from primary display to secondary.
+            final int stackNumFinal = mAmWmState.getAmState().getDisplay(DEFAULT_DISPLAY)
+                    .mStacks.size();
+            assertEquals("Stack number in default stack must be decremented.", stackNum - 1,
+                    stackNumFinal);
+        }
     }
 
     /**
@@ -609,48 +646,52 @@
      * matching the task component root does.
      */
     @Test
-    public void testTaskMatchAcrossDisplays() {
-        final ActivityDisplay newDisplay = createManagedVirtualDisplaySession().createDisplay();
+    public void testTaskMatchAcrossDisplays() throws Exception {
+        try (final VirtualDisplaySession virtualDisplaySession = new VirtualDisplaySession()) {
+            final ActivityDisplay newDisplay = virtualDisplaySession.createDisplay();
 
-        launchActivityOnDisplay(LAUNCHING_ACTIVITY, newDisplay.mId);
-        mAmWmState.computeState(LAUNCHING_ACTIVITY);
+            launchActivityOnDisplay(LAUNCHING_ACTIVITY, newDisplay.mId);
+            mAmWmState.computeState(LAUNCHING_ACTIVITY);
 
-        // Check that activity is on the secondary display.
-        final int frontStackId = mAmWmState.getAmState().getFrontStackId(newDisplay.mId);
-        final ActivityStack firstFrontStack = mAmWmState.getAmState().getStackById(frontStackId);
-        assertEquals("Activity launched on secondary display must be resumed",
-                getActivityName(LAUNCHING_ACTIVITY), firstFrontStack.mResumedActivity);
-        mAmWmState.assertFocusedStack("Top stack must be on secondary display", frontStackId);
+            // Check that activity is on the secondary display.
+            final int frontStackId = mAmWmState.getAmState().getFrontStackId(newDisplay.mId);
+            final ActivityStack firstFrontStack =
+                    mAmWmState.getAmState().getStackById(frontStackId);
+            assertEquals("Activity launched on secondary display must be resumed",
+                    getActivityName(LAUNCHING_ACTIVITY), firstFrontStack.mResumedActivity);
+            mAmWmState.assertFocusedStack("Top stack must be on secondary display",
+                    frontStackId);
 
-        executeShellCommand("am start -n " + getActivityName(ALT_LAUNCHING_ACTIVITY));
-        mAmWmState.waitForValidState(ALT_LAUNCHING_ACTIVITY);
+            executeShellCommand("am start -n " + getActivityName(ALT_LAUNCHING_ACTIVITY));
+            mAmWmState.waitForValidState(ALT_LAUNCHING_ACTIVITY);
 
-        // Check that second activity gets launched on the default display despite
-        // the affinity match on the secondary display.
-        final int defaultDisplayFrontStackId = mAmWmState.getAmState().getFrontStackId(
-                DEFAULT_DISPLAY);
-        final ActivityStack defaultDisplayFrontStack =
-                mAmWmState.getAmState().getStackById(defaultDisplayFrontStackId);
-        assertEquals("Activity launched on default display must be resumed",
-                getActivityName(ALT_LAUNCHING_ACTIVITY),
-                defaultDisplayFrontStack.mResumedActivity);
-        mAmWmState.assertFocusedStack("Top stack must be on primary display",
-                defaultDisplayFrontStackId);
+            // Check that second activity gets launched on the default display despite
+            // the affinity match on the secondary display.
+            final int defaultDisplayFrontStackId = mAmWmState.getAmState().getFrontStackId(
+                    DEFAULT_DISPLAY);
+            final ActivityStack defaultDisplayFrontStack =
+                    mAmWmState.getAmState().getStackById(defaultDisplayFrontStackId);
+            assertEquals("Activity launched on default display must be resumed",
+                    getActivityName(ALT_LAUNCHING_ACTIVITY),
+                    defaultDisplayFrontStack.mResumedActivity);
+            mAmWmState.assertFocusedStack("Top stack must be on primary display",
+                    defaultDisplayFrontStackId);
 
-        executeShellCommand("am start -n " + getActivityName(LAUNCHING_ACTIVITY));
-        waitAndAssertTopResumedActivity(LAUNCHING_ACTIVITY, newDisplay.mId,
-                "Existing task must be brought to front");
+            executeShellCommand("am start -n " + getActivityName(LAUNCHING_ACTIVITY));
+            mAmWmState.waitForFocusedStack(frontStackId);
 
-        // Check that the third intent is redirected to the first task due to the root
-        // component match on the secondary display.
-        final ActivityStack secondFrontStack = mAmWmState.getAmState().getStackById(frontStackId);
-        assertEquals("Activity launched on secondary display must be resumed",
-                getActivityName(LAUNCHING_ACTIVITY), secondFrontStack.mResumedActivity);
-        mAmWmState.assertFocusedStack("Top stack must be on primary display", frontStackId);
-        assertEquals("Top stack must only contain 1 task",
-                1, secondFrontStack.getTasks().size());
-        assertEquals("Top task must only contain 1 activity",
-                1, secondFrontStack.getTasks().get(0).mActivities.size());
+            // Check that the third intent is redirected to the first task due to the root
+            // component match on the secondary display.
+            final ActivityStack secondFrontStack =
+                    mAmWmState.getAmState().getStackById(frontStackId);
+            assertEquals("Activity launched on secondary display must be resumed",
+                    getActivityName(LAUNCHING_ACTIVITY), secondFrontStack.mResumedActivity);
+            mAmWmState.assertFocusedStack("Top stack must be on primary display", frontStackId);
+            assertEquals("Top stack must only contain 1 task",
+                    1, secondFrontStack.getTasks().size());
+            assertEquals("Top task must only contain 1 activity",
+                    1, secondFrontStack.getTasks().get(0).mActivities.size());
+        }
     }
 
     /**
@@ -658,7 +699,7 @@
      * both displays have matching tasks.
      */
     @Test
-    public void testTaskMatchOrderAcrossDisplays() {
+    public void testTaskMatchOrderAcrossDisplays() throws Exception {
         getLaunchActivityBuilder().setUseInstrumentation()
                 .setTargetActivity(TEST_ACTIVITY).setNewTask(true)
                 .setDisplayId(DEFAULT_DISPLAY).execute();
@@ -668,54 +709,62 @@
                 .setTargetActivity(BROADCAST_RECEIVER_ACTIVITY).setNewTask(true)
                 .setDisplayId(DEFAULT_DISPLAY).execute();
 
-        final ActivityDisplay newDisplay = createManagedVirtualDisplaySession().createDisplay();
-        getLaunchActivityBuilder().setUseInstrumentation().setWithShellPermission(true)
-                .setTargetActivity(TEST_ACTIVITY).setNewTask(true)
-                .setDisplayId(newDisplay.mId).execute();
-        assertNotEquals("Top focus stack should not be on default display",
-                stackId, mAmWmState.getAmState().getFocusedStackId());
+        try (final VirtualDisplaySession virtualDisplaySession = new VirtualDisplaySession()) {
+            final ActivityDisplay newDisplay = virtualDisplaySession.createDisplay();
+            getLaunchActivityBuilder().setUseInstrumentation().setWithShellPermission(true)
+                    .setTargetActivity(TEST_ACTIVITY).setNewTask(true)
+                    .setDisplayId(newDisplay.mId).execute();
+            assertNotEquals("Top focus stack should not be on default display",
+                    stackId, mAmWmState.getAmState().getFocusedStackId());
 
-        mBroadcastActionTrigger.launchActivityNewTask(getActivityName(TEST_ACTIVITY));
-        waitAndAssertTopResumedActivity(TEST_ACTIVITY, DEFAULT_DISPLAY,
-                "Activity must be launched on default display");
-        mAmWmState.assertFocusedStack("Top focus stack must be on the default display", stackId);
+            mBroadcastActionTrigger.launchActivityNewTask(getActivityName(TEST_ACTIVITY));
+            waitAndAssertTopResumedActivity(TEST_ACTIVITY, DEFAULT_DISPLAY,
+                    "Activity must be launched on default display");
+            mAmWmState.assertFocusedStack("Top focus stack must be on the default display",
+                    stackId);
+        }
     }
 
     /**
      * Tests that the task affinity search respects the launch display id.
      */
     @Test
-    public void testLaunchDisplayAffinityMatch() {
-        final ActivityDisplay newDisplay = createManagedVirtualDisplaySession().createDisplay();
+    public void testLaunchDisplayAffinityMatch() throws Exception {
+        try (final VirtualDisplaySession virtualDisplaySession = new VirtualDisplaySession()) {
+            final ActivityDisplay newDisplay = virtualDisplaySession.createDisplay();
 
-        launchActivityOnDisplay(LAUNCHING_ACTIVITY, newDisplay.mId);
+            launchActivityOnDisplay(LAUNCHING_ACTIVITY, newDisplay.mId);
 
-        // Check that activity is on the secondary display.
-        final int frontStackId = mAmWmState.getAmState().getFrontStackId(newDisplay.mId);
-        final ActivityStack firstFrontStack = mAmWmState.getAmState().getStackById(frontStackId);
-        assertEquals("Activity launched on secondary display must be resumed",
-                getActivityName(LAUNCHING_ACTIVITY), firstFrontStack.mResumedActivity);
-        mAmWmState.assertFocusedStack("Focus must be on secondary display", frontStackId);
+            // Check that activity is on the secondary display.
+            final int frontStackId = mAmWmState.getAmState().getFrontStackId(newDisplay.mId);
+            final ActivityStack firstFrontStack =
+                    mAmWmState.getAmState().getStackById(frontStackId);
+            assertEquals("Activity launched on secondary display must be resumed",
+                    getActivityName(LAUNCHING_ACTIVITY), firstFrontStack.mResumedActivity);
+            mAmWmState.assertFocusedStack("Focus must be on secondary display", frontStackId);
 
-        // We don't want FLAG_ACTIVITY_MULTIPLE_TASK, so we can't use launchActivityOnDisplay
-        executeShellCommand("am start -n " + getActivityName(ALT_LAUNCHING_ACTIVITY)
-                + " -f 0x10000000" // FLAG_ACTIVITY_NEW_TASK
-                + " --display " + newDisplay.mId);
-        mAmWmState.computeState(ALT_LAUNCHING_ACTIVITY);
+            // We don't want FLAG_ACTIVITY_MULTIPLE_TASK, so we can't use launchActivityOnDisplay
+            executeShellCommand("am start -n " + getActivityName(ALT_LAUNCHING_ACTIVITY)
+                    + " -f 0x10000000" // FLAG_ACTIVITY_NEW_TASK
+                    + " --display " + newDisplay.mId);
+            mAmWmState.computeState(ALT_LAUNCHING_ACTIVITY);
 
-        // Check that second activity gets launched into the affinity matching
-        // task on the secondary display
-        final int secondFrontStackId = mAmWmState.getAmState().getFrontStackId(newDisplay.mId);
-        final ActivityStack secondFrontStack =
-                mAmWmState.getAmState().getStackById(secondFrontStackId);
-        assertEquals("Activity launched on secondary display must be resumed",
-                getActivityName(ALT_LAUNCHING_ACTIVITY),
-                secondFrontStack.mResumedActivity);
-        mAmWmState.assertFocusedStack("Top stack must be on secondary display", secondFrontStackId);
-        assertEquals("Top stack must only contain 1 task",
-                1, secondFrontStack.getTasks().size());
-        assertEquals("Top stack task must contain 2 activities",
-                2, secondFrontStack.getTasks().get(0).mActivities.size());
+            // Check that second activity gets launched into the affinity matching
+            // task on the secondary display
+            final int secondFrontStackId =
+                    mAmWmState.getAmState().getFrontStackId(newDisplay.mId);
+            final ActivityStack secondFrontStack =
+                    mAmWmState.getAmState().getStackById(secondFrontStackId);
+            assertEquals("Activity launched on secondary display must be resumed",
+                    getActivityName(ALT_LAUNCHING_ACTIVITY),
+                    secondFrontStack.mResumedActivity);
+            mAmWmState.assertFocusedStack("Top stack must be on secondary display",
+                    secondFrontStackId);
+            assertEquals("Top stack must only contain 1 task",
+                    1, secondFrontStack.getTasks().size());
+            assertEquals("Top stack task must contain 2 activities",
+                    2, secondFrontStack.getTasks().get(0).mActivities.size());
+        }
     }
 
     /**
@@ -723,163 +772,181 @@
      * even if the focused stack is not on that activity's display.
      */
     @Test
-    public void testNewTaskSameDisplay() {
-        final ActivityDisplay newDisplay = createManagedVirtualDisplaySession()
-                .setSimulateDisplay(true)
-                .createDisplay();
+    public void testNewTaskSameDisplay() throws Exception {
+        try (final VirtualDisplaySession virtualDisplaySession = new VirtualDisplaySession()) {
+            final ActivityDisplay newDisplay = virtualDisplaySession.setSimulateDisplay(true)
+                    .createDisplay();
 
-        launchActivityOnDisplay(BROADCAST_RECEIVER_ACTIVITY, newDisplay.mId);
+            launchActivityOnDisplay(BROADCAST_RECEIVER_ACTIVITY, newDisplay.mId);
 
-        // Check that the first activity is launched onto the secondary display
-        waitAndAssertTopResumedActivity(BROADCAST_RECEIVER_ACTIVITY, newDisplay.mId,
-                "Activity launched on secondary display must be resumed");
+            // Check that the first activity is launched onto the secondary display
+            waitAndAssertTopResumedActivity(BROADCAST_RECEIVER_ACTIVITY, newDisplay.mId,
+                    "Activity launched on secondary display must be resumed");
 
-        executeShellCommand("am start -n " + getActivityName(TEST_ACTIVITY));
+            executeShellCommand("am start -n " + getActivityName(TEST_ACTIVITY));
 
-        // Check that the second activity is launched on the default display
-        waitAndAssertTopResumedActivity(TEST_ACTIVITY, DEFAULT_DISPLAY,
-                "Activity launched on default display must be resumed");
-        assertBothDisplaysHaveResumedActivities(pair(DEFAULT_DISPLAY, TEST_ACTIVITY),
-                pair(newDisplay.mId, BROADCAST_RECEIVER_ACTIVITY));
+            // Check that the second activity is launched on the default display
+            waitAndAssertTopResumedActivity(TEST_ACTIVITY, DEFAULT_DISPLAY,
+                    "Activity launched on default display must be resumed");
+            mAmWmState.assertResumedActivities("Both displays should have resumed activities",
+                    new SparseArray<ComponentName>(){{
+                        put(DEFAULT_DISPLAY, TEST_ACTIVITY);
+                        put(newDisplay.mId, BROADCAST_RECEIVER_ACTIVITY);
+                    }}
+            );
 
-        mBroadcastActionTrigger.launchActivityNewTask(getActivityName(LAUNCHING_ACTIVITY));
+            mBroadcastActionTrigger.launchActivityNewTask(getActivityName(LAUNCHING_ACTIVITY));
 
-        // Check that the third activity ends up in a new stack in the same display where the
-        // first activity lands
-        waitAndAssertTopResumedActivity(LAUNCHING_ACTIVITY, newDisplay.mId,
-                "Activity must be launched on secondary display");
-        assertEquals("Secondary display must contain 2 stacks", 2,
-                mAmWmState.getAmState().getDisplay(newDisplay.mId).mStacks.size());
-        assertBothDisplaysHaveResumedActivities(pair(DEFAULT_DISPLAY, TEST_ACTIVITY),
-                pair(newDisplay.mId, LAUNCHING_ACTIVITY));
+            // Check that the third activity ends up in a new stack in the same display where the
+            // first activity lands
+            waitAndAssertTopResumedActivity(LAUNCHING_ACTIVITY, newDisplay.mId,
+                    "Activity must be launched on secondary display");
+            assertEquals("Secondary display must contain 2 stacks", 2,
+                    mAmWmState.getAmState().getDisplay(newDisplay.mId).mStacks.size());
+            mAmWmState.assertResumedActivities("Both displays should have resumed activities",
+                    new SparseArray<ComponentName>(){{
+                        put(DEFAULT_DISPLAY, TEST_ACTIVITY);
+                        put(newDisplay.mId, LAUNCHING_ACTIVITY);
+                    }}
+            );
+        }
     }
 
     /**
      * Tests than an immediate launch after new display creation is handled correctly.
      */
     @Test
-    public void testImmediateLaunchOnNewDisplay() {
-        // Create new virtual display and immediately launch an activity on it.
-        final ActivityDisplay newDisplay = createManagedVirtualDisplaySession()
-                .setLaunchActivity(TEST_ACTIVITY)
-                .createDisplay();
+    public void testImmediateLaunchOnNewDisplay() throws Exception {
+        try (final VirtualDisplaySession virtualDisplaySession = new VirtualDisplaySession()) {
+            // Create new virtual display and immediately launch an activity on it.
+            final ActivityDisplay newDisplay = virtualDisplaySession
+                    .setLaunchActivity(TEST_ACTIVITY)
+                    .createDisplay();
 
-        // Check that activity is launched and placed correctly.
-        waitAndAssertTopResumedActivity(TEST_ACTIVITY, newDisplay.mId,
-                "Test activity must be on top");
-        final int frontStackId = mAmWmState.getAmState().getFrontStackId(newDisplay.mId);
-        final ActivityStack firstFrontStack = mAmWmState.getAmState().getStackById(frontStackId);
-        assertEquals("Activity launched on secondary display must be resumed",
-                getActivityName(TEST_ACTIVITY), firstFrontStack.mResumedActivity);
-        mAmWmState.assertFocusedStack("Top stack must be on secondary display",
-                frontStackId);
+            // Check that activity is launched and placed correctly.
+            waitAndAssertTopResumedActivity(TEST_ACTIVITY, newDisplay.mId,
+                    "Test activity must be on top");
+            final int frontStackId = mAmWmState.getAmState().getFrontStackId(newDisplay.mId);
+            final ActivityStack firstFrontStack =
+                    mAmWmState.getAmState().getStackById(frontStackId);
+            assertEquals("Activity launched on secondary display must be resumed",
+                    getActivityName(TEST_ACTIVITY), firstFrontStack.mResumedActivity);
+            mAmWmState.assertFocusedStack("Top stack must be on secondary display",
+                    frontStackId);
+        }
     }
 
     /** Tests launching of activities on a single task instance display. */
     @Test
-    public void testSingleTaskInstanceDisplay() {
-        ActivityDisplay display = createManagedVirtualDisplaySession()
-                .setSimulateDisplay(true)
-                .createDisplay();
-        final int displayId = display.mId;
+    public void testSingleTaskInstanceDisplay() throws Exception {
+        try (final VirtualDisplaySession virtualDisplaySession = new VirtualDisplaySession()) {
+            ActivityDisplay display =
+                    virtualDisplaySession.setSimulateDisplay(true).createDisplay();
+            final int displayId = display.mId;
 
-        SystemUtil.runWithShellPermissionIdentity(
-                () -> mAtm.setDisplayToSingleTaskInstance(displayId));
-        display = getDisplayState(displayId);
-        assertTrue("Display must be set to singleTaskInstance", display.mSingleTaskInstance);
+            SystemUtil.runWithShellPermissionIdentity(
+                    () -> mAtm.setDisplayToSingleTaskInstance(displayId));
+            display = getDisplayState(displayId);
+            assertTrue("Display must be set to singleTaskInstance", display.mSingleTaskInstance);
 
-        // SINGLE_TASK_INSTANCE_DISPLAY_ACTIVITY will launch
-        // SINGLE_TASK_INSTANCE_DISPLAY_ACTIVITY2 in the same task and
-        // SINGLE_TASK_INSTANCE_DISPLAY_ACTIVITY3 in different task.
-        launchActivityOnDisplay(SINGLE_TASK_INSTANCE_DISPLAY_ACTIVITY, displayId);
+            // SINGLE_TASK_INSTANCE_DISPLAY_ACTIVITY will launch
+            // SINGLE_TASK_INSTANCE_DISPLAY_ACTIVITY2 in the same task and
+            // SINGLE_TASK_INSTANCE_DISPLAY_ACTIVITY3 in different task.
+            launchActivityOnDisplay(SINGLE_TASK_INSTANCE_DISPLAY_ACTIVITY, displayId);
 
-        waitAndAssertTopResumedActivity(SINGLE_TASK_INSTANCE_DISPLAY_ACTIVITY3, DEFAULT_DISPLAY,
-                "Activity should be resumed on default display");
+            waitAndAssertTopResumedActivity(SINGLE_TASK_INSTANCE_DISPLAY_ACTIVITY3, DEFAULT_DISPLAY,
+                    "Activity should be resumed on default display");
 
-        display = getDisplayState(displayId);
-        // Verify that the 2 activities in the same task are on the display and the one in a
-        // different task isn't on the display, but on the default display
-        assertTrue("Display should contain SINGLE_TASK_INSTANCE_DISPLAY_ACTIVITY",
-                display.containsActivity(SINGLE_TASK_INSTANCE_DISPLAY_ACTIVITY));
-        assertTrue("Display should contain SINGLE_TASK_INSTANCE_DISPLAY_ACTIVITY2",
-                display.containsActivity(SINGLE_TASK_INSTANCE_DISPLAY_ACTIVITY2));
+            display = getDisplayState(displayId);
+            // Verify that the 2 activities in the same task are on the display and the one in a
+            // different task isn't on the display, but on the default display
+            assertTrue("Display should contain SINGLE_TASK_INSTANCE_DISPLAY_ACTIVITY",
+                    display.containsActivity(SINGLE_TASK_INSTANCE_DISPLAY_ACTIVITY));
+            assertTrue("Display should contain SINGLE_TASK_INSTANCE_DISPLAY_ACTIVITY2",
+                    display.containsActivity(SINGLE_TASK_INSTANCE_DISPLAY_ACTIVITY2));
 
-        assertFalse("Display shouldn't contain SINGLE_TASK_INSTANCE_DISPLAY_ACTIVITY3",
-                display.containsActivity(SINGLE_TASK_INSTANCE_DISPLAY_ACTIVITY3));
-        assertTrue("Display should contain SINGLE_TASK_INSTANCE_DISPLAY_ACTIVITY3",
-                getDisplayState(DEFAULT_DISPLAY).containsActivity(
-                        SINGLE_TASK_INSTANCE_DISPLAY_ACTIVITY3));
+            assertFalse("Display shouldn't contain SINGLE_TASK_INSTANCE_DISPLAY_ACTIVITY3",
+                    display.containsActivity(SINGLE_TASK_INSTANCE_DISPLAY_ACTIVITY3));
+            assertTrue("Display should contain SINGLE_TASK_INSTANCE_DISPLAY_ACTIVITY3",
+                    getDisplayState(DEFAULT_DISPLAY).containsActivity(
+                            SINGLE_TASK_INSTANCE_DISPLAY_ACTIVITY3));
+        }
     }
 
     @Test
     public void testLaunchPendingIntentActivity() throws Exception {
-        final DisplayManager displayManager = mContext.getSystemService(DisplayManager.class);
+        final DisplayManager displayManager =
+                getInstrumentation().getContext().getSystemService(DisplayManager.class);
 
-        final ActivityDisplay activityDisplay = createManagedVirtualDisplaySession()
-                .setSimulateDisplay(true)
-                .createDisplay();
+        try (final VirtualDisplaySession virtualDisplaySession = new VirtualDisplaySession()) {
+            final ActivityDisplay activityDisplay =
+                    virtualDisplaySession.setSimulateDisplay(true).createDisplay();
 
-        // Activity should be launched on primary display by default.
-        getPendingIntentActivity(TEST_ACTIVITY).send();
-        waitAndAssertTopResumedActivity(TEST_ACTIVITY, DEFAULT_DISPLAY,
-                "Activity launched on primary display and on top");
+            // Activity should be launched on primary display by default.
+            getPendingIntentActivity(TEST_ACTIVITY).send();
+            waitAndAssertTopResumedActivity(TEST_ACTIVITY, DEFAULT_DISPLAY,
+                    "Activity launched on primary display and on top");
 
-        final int resultCode = 1;
-        // Activity should be launched on target display according to the caller context.
-        final Context displayContext =
-                mContext.createDisplayContext(displayManager.getDisplay(activityDisplay.mId));
-        getPendingIntentActivity(TOP_ACTIVITY).send(displayContext, resultCode, null /* intent */);
-        waitAndAssertTopResumedActivity(TOP_ACTIVITY, activityDisplay.mId,
-                "Activity launched on secondary display and on top");
+            // Activity should be launched on target display according to the caller context.
+            final Context displayContext =
+                    mContext.createDisplayContext(displayManager.getDisplay(activityDisplay.mId));
+            getPendingIntentActivity(TOP_ACTIVITY).send(displayContext, 1 /* code */,
+                    null /* intent */);
+            waitAndAssertTopResumedActivity(TOP_ACTIVITY, activityDisplay.mId,
+                    "Activity launched on secondary display and on top");
 
-        // Activity should be brought to front on the same display if it already existed.
-        getPendingIntentActivity(TEST_ACTIVITY).send(displayContext, resultCode, null /* intent */);
-        waitAndAssertTopResumedActivity(TEST_ACTIVITY, DEFAULT_DISPLAY,
-                "Activity launched on primary display and on top");
+            // Activity should be brought to front on the same display if it already existed.
+            getPendingIntentActivity(TEST_ACTIVITY).send(displayContext, 1 /* code */,
+                    null /* intent */);
+            waitAndAssertTopResumedActivity(TEST_ACTIVITY, DEFAULT_DISPLAY,
+                    "Activity launched on primary display and on top");
 
-        // Activity should be moved to target display.
-        final ActivityOptions options = ActivityOptions.makeBasic();
-        options.setLaunchDisplayId(activityDisplay.mId);
-        getPendingIntentActivity(TEST_ACTIVITY).send(mContext, resultCode, null /* intent */,
-                null /* onFinished */, null /* handler */, null /* requiredPermission */,
-                options.toBundle());
-        waitAndAssertTopResumedActivity(TEST_ACTIVITY, activityDisplay.mId,
-                "Activity launched on secondary display and on top");
+            // Activity should be moved to target display.
+            ActivityOptions options = ActivityOptions.makeBasic();
+            options.setLaunchDisplayId(activityDisplay.mId);
+            getPendingIntentActivity(TEST_ACTIVITY).send(getInstrumentation().getContext(),
+                    1 /* code */, null /* intent */, null /* onFinished */, null /* handler */,
+                    null /* requiredPermission */, options.toBundle());
+            waitAndAssertTopResumedActivity(TEST_ACTIVITY, activityDisplay.mId,
+                    "Activity launched on secondary display and on top");
+        }
     }
 
     @Test
-    public void testLaunchActivityClearTask() {
+    public void testLaunchActivityClearTask() throws Exception {
         assertBroughtExistingTaskToAnotherDisplay(FLAG_ACTIVITY_CLEAR_TASK, LAUNCHING_ACTIVITY);
     }
 
     @Test
-    public void testLaunchActivityClearTop() {
+    public void testLaunchActivityClearTop() throws Exception {
         assertBroughtExistingTaskToAnotherDisplay(FLAG_ACTIVITY_CLEAR_TOP, LAUNCHING_ACTIVITY);
     }
 
     @Test
-    public void testLaunchActivitySingleTop() {
+    public void testLaunchActivitySingleTop() throws Exception {
         assertBroughtExistingTaskToAnotherDisplay(FLAG_ACTIVITY_SINGLE_TOP, TEST_ACTIVITY);
     }
 
-    private void assertBroughtExistingTaskToAnotherDisplay(int flags, ComponentName topActivity) {
+    private void assertBroughtExistingTaskToAnotherDisplay(int flags, ComponentName topActivity)
+            throws Exception {
         // Start TEST_ACTIVITY on top of LAUNCHING_ACTIVITY within the same task
         getLaunchActivityBuilder().setTargetActivity(TEST_ACTIVITY).execute();
 
-        final ActivityDisplay newDisplay = createManagedVirtualDisplaySession()
-                .setSimulateDisplay(true)
-                .createDisplay();
+        try (final VirtualDisplaySession virtualDisplaySession = new VirtualDisplaySession()) {
+            final ActivityDisplay newDisplay =
+                    virtualDisplaySession.setSimulateDisplay(true).createDisplay();
 
-        // Start LAUNCHING_ACTIVITY on secondary display with target flags, verify the task
-        // be reparented to secondary display
-        getLaunchActivityBuilder()
-                .setUseInstrumentation()
-                .setTargetActivity(LAUNCHING_ACTIVITY)
-                .setIntentFlags(flags)
-                .allowMultipleInstances(false)
-                .setDisplayId(newDisplay.mId).execute();
-        waitAndAssertTopResumedActivity(topActivity, newDisplay.mId,
-                "Activity launched on secondary display and on top");
+            // Start LAUNCHING_ACTIVITY on secondary display with target flags, verify the task
+            // be reparented to secondary display
+            getLaunchActivityBuilder()
+                    .setUseInstrumentation()
+                    .setTargetActivity(LAUNCHING_ACTIVITY)
+                    .setIntentFlags(flags)
+                    .allowMultipleInstances(false)
+                    .setDisplayId(newDisplay.mId).execute();
+            waitAndAssertTopResumedActivity(topActivity, newDisplay.mId,
+                    "Activity launched on secondary display and on top");
+        }
     }
 
     private PendingIntent getPendingIntentActivity(ComponentName activity) {
diff --git a/tests/framework/base/windowmanager/src/android/server/wm/MultiDisplayClientTests.java b/tests/framework/base/windowmanager/src/android/server/wm/MultiDisplayClientTests.java
index 4ab381c..272ef00 100755
--- a/tests/framework/base/windowmanager/src/android/server/wm/MultiDisplayClientTests.java
+++ b/tests/framework/base/windowmanager/src/android/server/wm/MultiDisplayClientTests.java
@@ -40,7 +40,6 @@
 import android.hardware.display.DisplayManager;
 import android.os.Bundle;
 import android.platform.test.annotations.Presubmit;
-import android.server.wm.ActivityManagerState.ActivityDisplay;
 import android.view.Display;
 import android.view.View;
 import android.view.WindowManager;
@@ -88,9 +87,9 @@
         testDisplayIdUpdateOnMove(NoRelaunchActivity.class, true /* handlesConfigChange */);
     }
 
-    private <T extends Activity> void testDisplayIdUpdateOnMove(Class<T> activityClass,
+    private void testDisplayIdUpdateOnMove(Class<? extends Activity> activityClass,
             boolean handlesConfigChange) throws Exception {
-        final ActivityTestRule<T> activityTestRule = new ActivityTestRule<>(
+        final ActivityTestRule activityTestRule = new ActivityTestRule(
                 activityClass, true /* initialTouchMode */, false /* launchActivity */);
 
         // Launch activity display.
@@ -99,44 +98,45 @@
         final ComponentName activityName = activity.getComponentName();
         waitAndAssertResume(activityName);
 
-        // Create new simulated display
-        final ActivityDisplay newDisplay = createManagedVirtualDisplaySession()
-                .setSimulateDisplay(true)
-                .createDisplay();
+        try (final VirtualDisplaySession virtualDisplaySession = new VirtualDisplaySession()) {
+            // Create new simulated display
+            final ActivityManagerState.ActivityDisplay newDisplay =
+                    virtualDisplaySession.setSimulateDisplay(true).createDisplay();
 
-        // Move the activity to the new secondary display.
-        separateTestJournal();
-        final ActivityOptions launchOptions = ActivityOptions.makeBasic();
-        final int displayId = newDisplay.mId;
-        launchOptions.setLaunchDisplayId(displayId);
-        final Intent newDisplayIntent = new Intent(mContext, activityClass);
-        newDisplayIntent.setFlags(FLAG_ACTIVITY_NEW_TASK);
-        getInstrumentation().getTargetContext().startActivity(newDisplayIntent,
-                launchOptions.toBundle());
-        waitAndAssertTopResumedActivity(activityName, displayId,
-                "Activity moved to secondary display must be focused");
+            // Move the activity to the new secondary display.
+            separateTestJournal();
+            final ActivityOptions launchOptions = ActivityOptions.makeBasic();
+            final int displayId = newDisplay.mId;
+            launchOptions.setLaunchDisplayId(displayId);
+            final Intent newDisplayIntent = new Intent(mContext, activityClass);
+            newDisplayIntent.setFlags(FLAG_ACTIVITY_NEW_TASK);
+            getInstrumentation().getTargetContext().startActivity(newDisplayIntent,
+                    launchOptions.toBundle());
+            waitAndAssertTopResumedActivity(activityName, displayId,
+                    "Activity moved to secondary display must be focused");
 
-        if (handlesConfigChange) {
-            // Wait for activity to receive the configuration change after move
-            waitAndAssertConfigurationChange(activityName);
-        } else {
-            // Activity will be re-created, wait for resumed state
-            waitAndAssertResume(activityName);
-            activity = activityTestRule.getActivity();
+            if (handlesConfigChange) {
+                // Wait for activity to receive the configuration change after move
+                waitAndAssertConfigurationChange(activityName);
+            } else {
+                // Activity will be re-created, wait for resumed state
+                waitAndAssertResume(activityName);
+                activity = activityTestRule.getActivity();
+            }
+
+            final String suffix = " must be updated.";
+            assertEquals("Activity#getDisplayId()" + suffix, displayId, activity.getDisplayId());
+            assertEquals("Activity#getDisplay" + suffix,
+                    displayId, activity.getDisplay().getDisplayId());
+
+            final WindowManager wm = activity.getWindowManager();
+            assertEquals("WM#getDefaultDisplay()" + suffix,
+                    displayId, wm.getDefaultDisplay().getDisplayId());
+
+            final View view = activity.getWindow().getDecorView();
+            assertEquals("View#getDisplay()" + suffix,
+                    displayId, view.getDisplay().getDisplayId());
         }
-
-        final String suffix = " must be updated.";
-        assertEquals("Activity#getDisplayId()" + suffix, displayId, activity.getDisplayId());
-        assertEquals("Activity#getDisplay" + suffix,
-                displayId, activity.getDisplay().getDisplayId());
-
-        final WindowManager wm = activity.getWindowManager();
-        assertEquals("WM#getDefaultDisplay()" + suffix,
-                displayId, wm.getDefaultDisplay().getDisplayId());
-
-        final View view = activity.getWindow().getDecorView();
-        assertEquals("View#getDisplay()" + suffix,
-                displayId, view.getDisplay().getDisplayId());
     }
 
     @Test
@@ -151,18 +151,19 @@
 
     private void testDisplayIdUpdateWhenImeMove(Class<? extends ImeTestActivity> activityClass)
             throws Exception {
-        final VirtualDisplaySession virtualDisplaySession = createManagedVirtualDisplaySession();
-        final MockImeSession mockImeSession = MockImeHelper.createManagedMockImeSession(this);
+        try (final VirtualDisplaySession virtualDisplaySession = new VirtualDisplaySession();
+            final MockImeSession mockImeSession = MockImeSession.create(mContext)) {
 
-        assertImeShownAndMatchesDisplayId(
-                activityClass, mockImeSession, DEFAULT_DISPLAY);
+            assertImeShownAndMatchesDisplayId(
+                    activityClass, mockImeSession, DEFAULT_DISPLAY);
 
-        final ActivityDisplay newDisplay = virtualDisplaySession
-                .setSimulateDisplay(true).setShowSystemDecorations(true).createDisplay();
+            final ActivityManagerState.ActivityDisplay newDisplay = virtualDisplaySession
+                    .setSimulateDisplay(true).setShowSystemDecorations(true).createDisplay();
 
-        // Launch activity on the secondary display and make IME show.
-        assertImeShownAndMatchesDisplayId(
-                activityClass, mockImeSession, newDisplay.mId);
+            // Launch activity on the secondary display and make IME show.
+            assertImeShownAndMatchesDisplayId(
+                    activityClass, mockImeSession, newDisplay.mId);
+        }
     }
 
     private  void assertImeShownAndMatchesDisplayId(Class<? extends ImeTestActivity> activityClass,
@@ -188,76 +189,83 @@
     }
 
     @Test
-    public void testInputMethodManagerDisplayId() {
-        // Create a simulated display.
-        final ActivityDisplay newDisplay = createManagedVirtualDisplaySession()
-                .setSimulateDisplay(true)
-                .createDisplay();
+    public void testInputMethodManagerDisplayId() throws Exception {
+        try (final VirtualDisplaySession virtualDisplaySession = new VirtualDisplaySession()) {
+            // Create a simulated display.
+            final ActivityManagerState.ActivityDisplay newDisplay = virtualDisplaySession
+                    .setSimulateDisplay(true).createDisplay();
 
-        final Display display = mContext.getSystemService(DisplayManager.class)
-                .getDisplay(newDisplay.mId);
-        final Context newDisplayContext = mContext.createDisplayContext(display);
-        final InputMethodManager imm = newDisplayContext.getSystemService(InputMethodManager.class);
+            final Display display = mContext.getSystemService(DisplayManager.class)
+                    .getDisplay(newDisplay.mId);
+            final Context newDisplayContext = mContext.createDisplayContext(display);
+            final InputMethodManager imm =
+                    newDisplayContext.getSystemService(InputMethodManager.class);
 
-        assertEquals("IMM#getDisplayId() must match.", newDisplay.mId, imm.getDisplayId());
+            assertEquals("IMM#getDisplayId() must match.", newDisplay.mId, imm.getDisplayId());
+        }
     }
 
     @Test
-    public void testViewGetDisplayOnPrimaryDisplay() {
+    public void testViewGetDisplayOnPrimaryDisplay() throws Exception {
         testViewGetDisplay(true /* isPrimary */);
     }
 
     @Test
-    public void testViewGetDisplayOnSecondaryDisplay() {
+    public void testViewGetDisplayOnSecondaryDisplay() throws Exception {
         testViewGetDisplay(false /* isPrimary */);
     }
 
-    private void testViewGetDisplay(boolean isPrimary) {
-        final TestActivitySession<ClientTestActivity> activitySession =
-                createManagedTestActivitySession();
-        final ActivityDisplay newDisplay = createManagedVirtualDisplaySession()
-                .setSimulateDisplay(true)
-                .createDisplay();
-        final int displayId = isPrimary ? DEFAULT_DISPLAY : newDisplay.mId;
+    private void testViewGetDisplay(boolean isPrimary) throws Exception {
+        try (final VirtualDisplaySession virtualDisplaySession = new VirtualDisplaySession();
+             final TestActivitySession<ClientTestActivity> activitySession =
+                     new TestActivitySession<>()) {
 
-        separateTestJournal();
-        activitySession.launchTestActivityOnDisplaySync(ClientTestActivity.class, displayId);
+            final ActivityManagerState.ActivityDisplay newDisplay = virtualDisplaySession
+                    .setSimulateDisplay(true).createDisplay();
+            final int displayId = (isPrimary) ? DEFAULT_DISPLAY : newDisplay.mId;
 
-        final Activity activity = activitySession.getActivity();
-        final ComponentName activityName = activity.getComponentName();
+            separateTestJournal();
+            activitySession.launchTestActivityOnDisplaySync(
+                    ClientTestActivity.class, displayId);
 
-        waitAndAssertTopResumedActivity(activityName, displayId,
-                "Activity launched on display:" + displayId + " must be focused");
+            final Activity activity = activitySession.getActivity();
+            final ComponentName activityName = activity.getComponentName();
 
-        // Test View#getdisplay() from activity
-        final View view = activity.getWindow().getDecorView();
-        assertEquals("View#getDisplay() must match.", displayId, view.getDisplay().getDisplayId());
+            waitAndAssertTopResumedActivity(activityName, displayId,
+                    "Activity launched on display:" + displayId + " must be focused");
 
-        final int[] resultDisplayId = { INVALID_DISPLAY };
-        activitySession.runOnMainAndAssertWithTimeout(
-                () -> {
-                    // Test View#getdisplay() from WM#addView()
-                    final WindowManager wm = activity.getWindowManager();
-                    final View addedView = new View(activity);
-                    wm.addView(addedView, new WindowManager.LayoutParams());
+            // Test View#getdisplay() from activity
+            final View view = activity.getWindow().getDecorView();
+            assertEquals("View#getDisplay() must match.",
+                    displayId, view.getDisplay().getDisplayId());
 
-                    // Get display ID from callback in case the added view has not be attached.
-                    addedView.addOnAttachStateChangeListener(
-                            new View.OnAttachStateChangeListener() {
-                                @Override
-                                public void onViewAttachedToWindow(View view) {
-                                    resultDisplayId[0] = view.getDisplay().getDisplayId();
-                                }
+            final int[] resultDisplayId = new int[] { INVALID_DISPLAY };
+            activitySession.runOnMainAndAssertWithTimeout(
+                    () -> {
+                        // Test View#getdisplay() from WM#addView()
+                        final WindowManager wm = activity.getWindowManager();
+                        final View addedView = new View(activity);
+                        wm.addView(addedView, new WindowManager.LayoutParams());
 
-                                @Override
-                                public void onViewDetachedFromWindow(View view) {}
-                            });
+                        // Get display ID from callback in case the added view has not be attached.
+                        addedView.addOnAttachStateChangeListener(
+                                new View.OnAttachStateChangeListener() {
+                                    @Override
+                                    public void onViewAttachedToWindow(View view) {
+                                        resultDisplayId[0] = view.getDisplay().getDisplayId();
+                                    }
 
-                    return displayId == resultDisplayId[0];
-                }, TIMEOUT, "Display from added view must match. "
-                        + "Should be display:" + displayId
-                        + ", but was display:" + resultDisplayId[0]
-        );
+                                    @Override
+                                    public void onViewDetachedFromWindow(View view) {}
+                                });
+
+                        return displayId == resultDisplayId[0];
+                    }, TIMEOUT, "Display from added view must match. "
+                            + "Should be display:" + displayId
+                            + ", but was display:" + resultDisplayId[0]
+            );
+
+        }
     }
 
     private void waitAndAssertConfigurationChange(ComponentName activityName) {
diff --git a/tests/framework/base/windowmanager/src/android/server/wm/MultiDisplayKeyguardTests.java b/tests/framework/base/windowmanager/src/android/server/wm/MultiDisplayKeyguardTests.java
index 780712a..7007120 100644
--- a/tests/framework/base/windowmanager/src/android/server/wm/MultiDisplayKeyguardTests.java
+++ b/tests/framework/base/windowmanager/src/android/server/wm/MultiDisplayKeyguardTests.java
@@ -56,16 +56,18 @@
      * insecure keyguard).
      */
     @Test
-    public void testDismissKeyguardActivity_secondaryDisplay() {
-        final LockScreenSession lockScreenSession = createManagedLockScreenSession();
-        final ActivityDisplay newDisplay = createManagedVirtualDisplaySession().createDisplay();
+    public void testDismissKeyguardActivity_secondaryDisplay() throws Exception {
+        try (final LockScreenSession lockScreenSession = new LockScreenSession();
+             final VirtualDisplaySession virtualDisplaySession = new VirtualDisplaySession()) {
+            final ActivityDisplay newDisplay = virtualDisplaySession.createDisplay();
 
-        lockScreenSession.gotoKeyguard();
-        mAmWmState.assertKeyguardShowingAndNotOccluded();
-        launchActivityOnDisplay(DISMISS_KEYGUARD_ACTIVITY, newDisplay.mId);
-        mAmWmState.waitForKeyguardShowingAndNotOccluded();
-        mAmWmState.assertKeyguardShowingAndNotOccluded();
-        mAmWmState.assertVisibility(DISMISS_KEYGUARD_ACTIVITY, true);
+            lockScreenSession.gotoKeyguard();
+            mAmWmState.assertKeyguardShowingAndNotOccluded();
+            launchActivityOnDisplay(DISMISS_KEYGUARD_ACTIVITY, newDisplay.mId);
+            mAmWmState.waitForKeyguardShowingAndNotOccluded();
+            mAmWmState.assertKeyguardShowingAndNotOccluded();
+            mAmWmState.assertVisibility(DISMISS_KEYGUARD_ACTIVITY, true);
+        }
     }
 
     /**
@@ -73,23 +75,23 @@
      * @throws Exception
      */
     @Test
-    public void testShowKeyguardDialogOnSecondaryDisplay() {
-        final LockScreenSession lockScreenSession = createManagedLockScreenSession();
-        final ActivityDisplay publicDisplay = createManagedVirtualDisplaySession()
-                .setPublicDisplay(true)
-                .createDisplay();
+    public void testShowKeyguardDialogOnSecondaryDisplay() throws Exception {
+        try (final LockScreenSession lockScreenSession = new LockScreenSession();
+             final VirtualDisplaySession virtualDisplaySession = new VirtualDisplaySession()) {
+            final ActivityDisplay publicDisplay = virtualDisplaySession.setPublicDisplay(true)
+                    .createDisplay();
+            lockScreenSession.gotoKeyguard();
+            assertTrue("KeyguardDialog must show on external public display",
+                    mAmWmState.waitForWithWmState(
+                            state -> isKeyguardOnDisplay(state, publicDisplay.mId),
+                            "keyguard window to show"));
 
-        lockScreenSession.gotoKeyguard();
-        assertTrue("KeyguardDialog must show on external public display",
-                mAmWmState.waitForWithWmState(
-                        state -> isKeyguardOnDisplay(state, publicDisplay.mId),
-                        "keyguard window to show"));
-
-        // Keyguard dialog mustn't be removed when press back key
-        pressBackButton();
-        mAmWmState.computeState(true);
-        assertTrue("KeyguardDialog must not be removed when press back key",
-                isKeyguardOnDisplay(mAmWmState.getWmState(), publicDisplay.mId));
+            // Keyguard dialog mustn't be removed when press back key
+            pressBackButton();
+            mAmWmState.computeState(true);
+            assertTrue("KeyguardDialog must not be removed when press back key",
+                    isKeyguardOnDisplay(mAmWmState.getWmState(), publicDisplay.mId));
+        }
     }
 
     /**
@@ -97,23 +99,23 @@
      * @throws Exception
      */
     @Test
-    public void testNoKeyguardDialogOnPrivateDisplay() {
-        final LockScreenSession lockScreenSession = createManagedLockScreenSession();
-        final VirtualDisplaySession virtualDisplaySession = createManagedVirtualDisplaySession();
+    public void testNoKeyguardDialogOnPrivateDisplay() throws Exception {
+        try (final LockScreenSession lockScreenSession = new LockScreenSession();
+             final VirtualDisplaySession virtualDisplaySession = new VirtualDisplaySession()) {
+            final ActivityDisplay privateDisplay = virtualDisplaySession.setPublicDisplay(false)
+                    .createDisplay();
+            final ActivityDisplay publicDisplay = virtualDisplaySession.setPublicDisplay(true)
+                    .createDisplay();
 
-        final ActivityDisplay privateDisplay =
-                virtualDisplaySession.setPublicDisplay(false).createDisplay();
-        final ActivityDisplay publicDisplay =
-                virtualDisplaySession.setPublicDisplay(true).createDisplay();
+            lockScreenSession.gotoKeyguard();
+            assertTrue("KeyguardDialog must show on external public display",
+                    mAmWmState.waitForWithWmState(
+                            state -> isKeyguardOnDisplay(state, publicDisplay.mId),
+                            "keyguard window to show"));
 
-        lockScreenSession.gotoKeyguard();
-        assertTrue("KeyguardDialog must show on external public display",
-                mAmWmState.waitForWithWmState(
-                        state -> isKeyguardOnDisplay(state, publicDisplay.mId),
-                        "keyguard window to show"));
-
-        assertFalse("KeyguardDialog must not show on external private display",
-                isKeyguardOnDisplay(mAmWmState.getWmState(), privateDisplay.mId));
+            assertFalse("KeyguardDialog must not show on external private display",
+                    isKeyguardOnDisplay(mAmWmState.getWmState(), privateDisplay.mId));
+        }
     }
 
     private boolean isKeyguardOnDisplay(WindowManagerState windowManagerState, int displayId) {
diff --git a/tests/framework/base/windowmanager/src/android/server/wm/MultiDisplayLockedKeyguardTests.java b/tests/framework/base/windowmanager/src/android/server/wm/MultiDisplayLockedKeyguardTests.java
index 590ff27..353558db 100644
--- a/tests/framework/base/windowmanager/src/android/server/wm/MultiDisplayLockedKeyguardTests.java
+++ b/tests/framework/base/windowmanager/src/android/server/wm/MultiDisplayLockedKeyguardTests.java
@@ -57,111 +57,113 @@
      * Test that virtual display content is hidden when device is locked.
      */
     @Test
-    public void testVirtualDisplayHidesContentWhenLocked() {
-        final LockScreenSession lockScreenSession = createManagedLockScreenSession();
-        lockScreenSession.setLockCredential();
+    public void testVirtualDisplayHidesContentWhenLocked() throws Exception {
+        try (final LockScreenSession lockScreenSession = new LockScreenSession();
+             final VirtualDisplaySession virtualDisplaySession = new VirtualDisplaySession()) {
+            lockScreenSession.setLockCredential();
 
-        // Create new usual virtual display.
-        final ActivityDisplay newDisplay = createManagedVirtualDisplaySession()
-                .setPublicDisplay(true)
-                .createDisplay();
-        mAmWmState.assertVisibility(VIRTUAL_DISPLAY_ACTIVITY, true /* visible */);
+            // Create new usual virtual display.
+            final ActivityDisplay newDisplay = virtualDisplaySession.setPublicDisplay(true)
+                    .createDisplay();
+            mAmWmState.assertVisibility(VIRTUAL_DISPLAY_ACTIVITY, true /* visible */);
 
-        // Launch activity on new secondary display.
-        launchActivityOnDisplay(TEST_ACTIVITY, newDisplay.mId);
-        mAmWmState.assertVisibility(TEST_ACTIVITY, true /* visible */);
+            // Launch activity on new secondary display.
+            launchActivityOnDisplay(TEST_ACTIVITY, newDisplay.mId);
+            mAmWmState.assertVisibility(TEST_ACTIVITY, true /* visible */);
 
-        // Lock the device.
-        lockScreenSession.gotoKeyguard();
-        waitAndAssertActivityState(TEST_ACTIVITY, STATE_STOPPED,
-                "Expected stopped activity on secondary display ");
-        mAmWmState.assertVisibility(TEST_ACTIVITY, false /* visible */);
+            // Lock the device.
+            lockScreenSession.gotoKeyguard();
+            waitAndAssertActivityState(TEST_ACTIVITY, STATE_STOPPED,
+                    "Expected stopped activity on secondary display ");
+            mAmWmState.assertVisibility(TEST_ACTIVITY, false /* visible */);
 
-        // Unlock and check if visibility is back.
-        lockScreenSession.unlockDevice();
+            // Unlock and check if visibility is back.
+            lockScreenSession.unlockDevice();
 
-        lockScreenSession.enterAndConfirmLockCredential();
-        mAmWmState.waitForKeyguardGone();
-        mAmWmState.assertKeyguardGone();
-        waitAndAssertActivityState(TEST_ACTIVITY, STATE_RESUMED,
-                "Expected resumed activity on secondary display");
-        mAmWmState.assertVisibility(TEST_ACTIVITY, true /* visible */);
+            lockScreenSession.enterAndConfirmLockCredential();
+            mAmWmState.waitForKeyguardGone();
+            mAmWmState.assertKeyguardGone();
+            waitAndAssertActivityState(TEST_ACTIVITY, STATE_RESUMED,
+                    "Expected resumed activity on secondary display");
+            mAmWmState.assertVisibility(TEST_ACTIVITY, true /* visible */);
+        }
     }
 
     /**
      * Tests that private display cannot show content while device locked.
      */
     @Test
-    public void testPrivateDisplayHideContentWhenLocked() {
-        final LockScreenSession lockScreenSession = createManagedLockScreenSession();
-        lockScreenSession.setLockCredential();
+    public void testPrivateDisplayHideContentWhenLocked() throws Exception {
+        try (final LockScreenSession lockScreenSession = new LockScreenSession();
+             final VirtualDisplaySession virtualDisplaySession = new VirtualDisplaySession()) {
+            lockScreenSession.setLockCredential();
 
-        final ActivityDisplay newDisplay = createManagedVirtualDisplaySession()
-                .setPublicDisplay(false)
-                .createDisplay();
-        launchActivityOnDisplay(TEST_ACTIVITY, newDisplay.mId);
+            final ActivityDisplay newDisplay =
+                    virtualDisplaySession.setPublicDisplay(false).createDisplay();
+            launchActivityOnDisplay(TEST_ACTIVITY, newDisplay.mId);
 
-        lockScreenSession.gotoKeyguard();
+            lockScreenSession.gotoKeyguard();
 
-        waitAndAssertActivityState(TEST_ACTIVITY, STATE_STOPPED,
-                "Expected stopped activity on private display");
-        mAmWmState.assertVisibility(TEST_ACTIVITY, false /* visible */);
+            waitAndAssertActivityState(TEST_ACTIVITY, STATE_STOPPED,
+                    "Expected stopped activity on private display");
+            mAmWmState.assertVisibility(TEST_ACTIVITY, false /* visible */);
+        }
     }
 
     /**
      * Tests whether a FLAG_DISMISS_KEYGUARD activity on a secondary display dismisses the keyguard.
      */
     @Test
-    public void testDismissKeyguard_secondaryDisplay() {
-        final LockScreenSession lockScreenSession =
-                mObjectTracker.manage(new LockScreenSession(FLAG_REMOVE_ACTIVITIES_ON_CLOSE));
-        lockScreenSession.setLockCredential();
+    public void testDismissKeyguard_secondaryDisplay() throws Exception {
+        try (final LockScreenSession lockScreenSession =
+                     new LockScreenSession(FLAG_REMOVE_ACTIVITIES_ON_CLOSE);
+             final VirtualDisplaySession virtualDisplaySession = new VirtualDisplaySession()) {
+            lockScreenSession.setLockCredential();
+            final ActivityDisplay newDisplay = virtualDisplaySession.setPublicDisplay(true).
+                    createDisplay();
 
-        final ActivityDisplay newDisplay = createManagedVirtualDisplaySession()
-                .setPublicDisplay(true)
-                .createDisplay();
-
-        lockScreenSession.gotoKeyguard();
-        mAmWmState.assertKeyguardShowingAndNotOccluded();
-        getLaunchActivityBuilder().setUseInstrumentation()
-                .setTargetActivity(DISMISS_KEYGUARD_ACTIVITY).setNewTask(true)
-                .setMultipleTask(true).setDisplayId(newDisplay.mId)
-                .setWaitForLaunched(false).execute();
-        waitAndAssertActivityState(DISMISS_KEYGUARD_ACTIVITY, STATE_STOPPED,
-                "Expected stopped activity on secondary display");
-        lockScreenSession.enterAndConfirmLockCredential();
-        mAmWmState.waitForKeyguardGone();
-        mAmWmState.assertKeyguardGone();
-        mAmWmState.assertVisibility(DISMISS_KEYGUARD_ACTIVITY, true);
+            lockScreenSession.gotoKeyguard();
+            mAmWmState.assertKeyguardShowingAndNotOccluded();
+            getLaunchActivityBuilder().setUseInstrumentation()
+                    .setTargetActivity(DISMISS_KEYGUARD_ACTIVITY).setNewTask(true)
+                    .setMultipleTask(true).setDisplayId(newDisplay.mId)
+                    .setWaitForLaunched(false).execute();
+            waitAndAssertActivityState(DISMISS_KEYGUARD_ACTIVITY, STATE_STOPPED,
+                    "Expected stopped activity on secondary display");
+            lockScreenSession.enterAndConfirmLockCredential();
+            mAmWmState.waitForKeyguardGone();
+            mAmWmState.assertKeyguardGone();
+            mAmWmState.assertVisibility(DISMISS_KEYGUARD_ACTIVITY, true);
+        }
     }
 
     @FlakyTest(bugId = 141674516)
     @Test
-    public void testDismissKeyguard_whileOccluded_secondaryDisplay() {
-        final LockScreenSession lockScreenSession =
-                mObjectTracker.manage(new LockScreenSession(FLAG_REMOVE_ACTIVITIES_ON_CLOSE));
-        lockScreenSession.setLockCredential();
+    public void testDismissKeyguard_whileOccluded_secondaryDisplay() throws Exception {
+        try (final LockScreenSession lockScreenSession =
+                     new LockScreenSession(FLAG_REMOVE_ACTIVITIES_ON_CLOSE);
+             final VirtualDisplaySession virtualDisplaySession = new VirtualDisplaySession()) {
+            lockScreenSession.setLockCredential();
+            final ActivityDisplay newDisplay = virtualDisplaySession.setPublicDisplay(true).
+                    createDisplay();
 
-        final ActivityDisplay newDisplay = createManagedVirtualDisplaySession()
-                .setPublicDisplay(true)
-                .createDisplay();
-
-        lockScreenSession.gotoKeyguard();
-        mAmWmState.assertKeyguardShowingAndNotOccluded();
-        launchActivity(SHOW_WHEN_LOCKED_ACTIVITY);
-        mAmWmState.computeState(SHOW_WHEN_LOCKED_ACTIVITY);
-        mAmWmState.assertVisibility(SHOW_WHEN_LOCKED_ACTIVITY, true);
-        getLaunchActivityBuilder().setUseInstrumentation()
-                .setTargetActivity(DISMISS_KEYGUARD_ACTIVITY).setNewTask(true)
-                .setMultipleTask(true).setDisplayId(newDisplay.mId)
-                .setWaitForLaunched(false).execute();
-        waitAndAssertActivityState(DISMISS_KEYGUARD_ACTIVITY, STATE_STOPPED,
-                "Expected stopped activity on secondary display");
-        lockScreenSession.enterAndConfirmLockCredential();
-        mAmWmState.waitForKeyguardGone();
-        mAmWmState.assertKeyguardGone();
-        mAmWmState.computeState(DISMISS_KEYGUARD_ACTIVITY);
-        mAmWmState.assertVisibility(DISMISS_KEYGUARD_ACTIVITY, true);
-        mAmWmState.assertVisibility(SHOW_WHEN_LOCKED_ACTIVITY, true);
+            lockScreenSession.gotoKeyguard();
+            mAmWmState.assertKeyguardShowingAndNotOccluded();
+            launchActivity(SHOW_WHEN_LOCKED_ACTIVITY);
+            mAmWmState.computeState(SHOW_WHEN_LOCKED_ACTIVITY);
+            mAmWmState.assertVisibility(SHOW_WHEN_LOCKED_ACTIVITY, true);
+            getLaunchActivityBuilder().setUseInstrumentation()
+                    .setTargetActivity(DISMISS_KEYGUARD_ACTIVITY).setNewTask(true)
+                    .setMultipleTask(true).setDisplayId(newDisplay.mId)
+                    .setWaitForLaunched(false).execute();
+            waitAndAssertActivityState(DISMISS_KEYGUARD_ACTIVITY, STATE_STOPPED,
+                    "Expected stopped activity on secondary display");
+            lockScreenSession.enterAndConfirmLockCredential();
+            mAmWmState.waitForKeyguardGone();
+            mAmWmState.assertKeyguardGone();
+            mAmWmState.computeState(DISMISS_KEYGUARD_ACTIVITY);
+            mAmWmState.assertVisibility(DISMISS_KEYGUARD_ACTIVITY, true);
+            mAmWmState.assertVisibility(SHOW_WHEN_LOCKED_ACTIVITY, true);
+        }
     }
 }
diff --git a/tests/framework/base/windowmanager/src/android/server/wm/MultiDisplayPolicyTests.java b/tests/framework/base/windowmanager/src/android/server/wm/MultiDisplayPolicyTests.java
index 4336cf1..76086ed 100644
--- a/tests/framework/base/windowmanager/src/android/server/wm/MultiDisplayPolicyTests.java
+++ b/tests/framework/base/windowmanager/src/android/server/wm/MultiDisplayPolicyTests.java
@@ -47,12 +47,14 @@
 import static org.junit.Assume.assumeFalse;
 import static org.junit.Assume.assumeTrue;
 
+import android.content.ComponentName;
 import android.platform.test.annotations.Presubmit;
 import android.server.wm.ActivityManagerState.ActivityDisplay;
 import android.server.wm.ActivityManagerState.ActivityStack;
 import android.server.wm.CommandSession.ActivityCallback;
 import android.server.wm.CommandSession.ActivitySession;
 import android.server.wm.CommandSession.SizeInfo;
+import android.util.SparseArray;
 
 import org.junit.Before;
 import org.junit.Test;
@@ -77,7 +79,7 @@
      * Tests that all activities that were on the private display are destroyed on display removal.
      */
     @Test
-    public void testContentDestroyOnDisplayRemoved() {
+    public void testContentDestroyOnDisplayRemoved() throws Exception {
         try (final VirtualDisplaySession virtualDisplaySession = new VirtualDisplaySession()) {
             // Create new private virtual display.
             final ActivityDisplay newDisplay = virtualDisplaySession.createDisplay();
@@ -118,7 +120,7 @@
      * Tests that newly launched activity will be landing on default display on display removal.
      */
     @Test
-    public void testActivityLaunchOnContentDestroyDisplayRemoved() {
+    public void testActivityLaunchOnContentDestroyDisplayRemoved() throws Exception {
         try (final VirtualDisplaySession virtualDisplaySession = new VirtualDisplaySession()) {
             // Create new private virtual display.
             final ActivityDisplay newDisplay = virtualDisplaySession.createDisplay();
@@ -141,93 +143,102 @@
      * Tests that the update of display metrics updates all its content.
      */
     @Test
-    public void testDisplayResize() {
-        final VirtualDisplaySession virtualDisplaySession = createManagedVirtualDisplaySession();
-        // Create new virtual display.
-        final ActivityDisplay newDisplay = virtualDisplaySession.createDisplay();
-        mAmWmState.assertVisibility(VIRTUAL_DISPLAY_ACTIVITY, true /* visible */);
+    public void testDisplayResize() throws Exception {
+        try (final VirtualDisplaySession virtualDisplaySession = new VirtualDisplaySession()) {
+            // Create new virtual display.
+            final ActivityDisplay newDisplay = virtualDisplaySession.createDisplay();
+            mAmWmState.assertVisibility(VIRTUAL_DISPLAY_ACTIVITY, true /* visible */);
 
-        // Launch a resizeable activity on new secondary display.
-        separateTestJournal();
-        launchActivityOnDisplay(RESIZEABLE_ACTIVITY, newDisplay.mId);
-        waitAndAssertTopResumedActivity(RESIZEABLE_ACTIVITY, newDisplay.mId,
-                "Launched activity must be on top");
+            // Launch a resizeable activity on new secondary display.
+            separateTestJournal();
+            launchActivityOnDisplay(RESIZEABLE_ACTIVITY, newDisplay.mId);
+            waitAndAssertTopResumedActivity(RESIZEABLE_ACTIVITY, newDisplay.mId,
+                    "Launched activity must be on top");
 
-        // Grab reported sizes and compute new with slight size change.
-        final SizeInfo initialSize = getLastReportedSizesForActivity(RESIZEABLE_ACTIVITY);
+            // Grab reported sizes and compute new with slight size change.
+            final SizeInfo initialSize = getLastReportedSizesForActivity(RESIZEABLE_ACTIVITY);
 
-        // Resize the display
-        separateTestJournal();
-        virtualDisplaySession.resizeDisplay();
+            // Resize the display
+            separateTestJournal();
+            virtualDisplaySession.resizeDisplay();
 
-        mAmWmState.waitForWithAmState(amState -> {
-            try {
-                return amState.hasActivityState(RESIZEABLE_ACTIVITY, STATE_RESUMED)
-                        && new ActivityLifecycleCounts(RESIZEABLE_ACTIVITY)
-                                .getCount(ActivityCallback.ON_CONFIGURATION_CHANGED) == 1;
-            } catch (Exception e) {
-                logE("Error waiting for valid state: " + e.getMessage());
-                return false;
-            }
-        }, "the configuration change to happen and activity to be resumed");
+            mAmWmState.waitForWithAmState(amState -> {
+                try {
+                    return readConfigChangeNumber(RESIZEABLE_ACTIVITY) == 1
+                            && amState.hasActivityState(RESIZEABLE_ACTIVITY, STATE_RESUMED);
+                } catch (Exception e) {
+                    logE("Error waiting for valid state: " + e.getMessage());
+                    return false;
+                }
+            }, "the configuration change to happen and activity to be resumed");
 
-        mAmWmState.computeState(false /* compareTaskAndStackBounds */,
-                new WaitForValidActivityState(RESIZEABLE_ACTIVITY),
-                new WaitForValidActivityState(VIRTUAL_DISPLAY_ACTIVITY));
-        mAmWmState.assertVisibility(VIRTUAL_DISPLAY_ACTIVITY, true);
-        mAmWmState.assertVisibility(RESIZEABLE_ACTIVITY, true);
+            mAmWmState.computeState(false /* compareTaskAndStackBounds */,
+                    new WaitForValidActivityState(RESIZEABLE_ACTIVITY),
+                    new WaitForValidActivityState(VIRTUAL_DISPLAY_ACTIVITY));
+            mAmWmState.assertVisibility(VIRTUAL_DISPLAY_ACTIVITY, true);
+            mAmWmState.assertVisibility(RESIZEABLE_ACTIVITY, true);
 
-        // Check if activity in virtual display was resized properly.
-        assertRelaunchOrConfigChanged(RESIZEABLE_ACTIVITY, 0 /* numRelaunch */,
-                1 /* numConfigChange */);
+            // Check if activity in virtual display was resized properly.
+            assertRelaunchOrConfigChanged(RESIZEABLE_ACTIVITY, 0 /* numRelaunch */,
+                    1 /* numConfigChange */);
 
-        final SizeInfo updatedSize = getLastReportedSizesForActivity(RESIZEABLE_ACTIVITY);
-        assertTrue(updatedSize.widthDp <= initialSize.widthDp);
-        assertTrue(updatedSize.heightDp <= initialSize.heightDp);
-        assertTrue(updatedSize.displayWidth == initialSize.displayWidth / 2);
-        assertTrue(updatedSize.displayHeight == initialSize.displayHeight / 2);
+            final SizeInfo updatedSize = getLastReportedSizesForActivity(RESIZEABLE_ACTIVITY);
+            assertTrue(updatedSize.widthDp <= initialSize.widthDp);
+            assertTrue(updatedSize.heightDp <= initialSize.heightDp);
+            assertTrue(updatedSize.displayWidth == initialSize.displayWidth / 2);
+            assertTrue(updatedSize.displayHeight == initialSize.displayHeight / 2);
+        }
+    }
+
+    /** Read the number of configuration changes sent to activity from logs. */
+    private int readConfigChangeNumber(ComponentName activityName) throws Exception {
+        return (new ActivityLifecycleCounts(activityName))
+                .getCount(ActivityCallback.ON_CONFIGURATION_CHANGED);
     }
 
     /**
      * Tests that when primary display is rotated secondary displays are not affected.
      */
     @Test
-    public void testRotationNotAffectingSecondaryScreen() {
-        final VirtualDisplayLauncher virtualLauncher =
-                mObjectTracker.manage(new VirtualDisplayLauncher());
-        // Create new virtual display.
-        final ActivityDisplay newDisplay = virtualLauncher.setResizeDisplay(false).createDisplay();
-        mAmWmState.assertVisibility(VIRTUAL_DISPLAY_ACTIVITY, true /* visible */);
+    public void testRotationNotAffectingSecondaryScreen() throws Exception {
+        try (final VirtualDisplayLauncher virtualLauncher = new VirtualDisplayLauncher()) {
+            // Create new virtual display.
+            final ActivityDisplay newDisplay = virtualLauncher.setResizeDisplay(false)
+                    .createDisplay();
+            mAmWmState.assertVisibility(VIRTUAL_DISPLAY_ACTIVITY, true /* visible */);
 
-        // Launch activity on new secondary display.
-        final ActivitySession resizeableActivitySession =
-                virtualLauncher.launchActivityOnDisplay(RESIZEABLE_ACTIVITY, newDisplay);
-        waitAndAssertTopResumedActivity(RESIZEABLE_ACTIVITY, newDisplay.mId,
-                "Top activity must be on secondary display");
-        final SizeInfo initialSize = resizeableActivitySession.getConfigInfo().sizeInfo;
+            // Launch activity on new secondary display.
+            final ActivitySession resizeableActivitySession =
+                    virtualLauncher.launchActivityOnDisplay(RESIZEABLE_ACTIVITY, newDisplay);
+            waitAndAssertTopResumedActivity(RESIZEABLE_ACTIVITY, newDisplay.mId,
+                    "Top activity must be on secondary display");
+            final SizeInfo initialSize = resizeableActivitySession.getConfigInfo().sizeInfo;
 
-        assertNotNull("Test activity must have reported initial size on launch", initialSize);
+            assertNotNull("Test activity must have reported initial size on launch", initialSize);
 
-        final RotationSession rotationSession = createManagedRotationSession();
-        // Rotate primary display and check that activity on secondary display is not affected.
-        rotateAndCheckSameSizes(rotationSession, resizeableActivitySession, initialSize);
+            try (final RotationSession rotationSession = new RotationSession()) {
+                // Rotate primary display and check that activity on secondary display is not
+                // affected.
+                rotateAndCheckSameSizes(rotationSession, resizeableActivitySession, initialSize);
 
-        // Launch activity to secondary display when primary one is rotated.
-        final int initialRotation = mAmWmState.getWmState().getRotation();
-        rotationSession.set((initialRotation + 1) % 4);
+                // Launch activity to secondary display when primary one is rotated.
+                final int initialRotation = mAmWmState.getWmState().getRotation();
+                rotationSession.set((initialRotation + 1) % 4);
 
-        final ActivitySession testActivitySession =
-                virtualLauncher.launchActivityOnDisplay(TEST_ACTIVITY, newDisplay);
-        waitAndAssertTopResumedActivity(TEST_ACTIVITY, newDisplay.mId,
-                "Top activity must be on secondary display");
-        final SizeInfo testActivitySize = testActivitySession.getConfigInfo().sizeInfo;
+                final ActivitySession testActivitySession =
+                        virtualLauncher.launchActivityOnDisplay(TEST_ACTIVITY, newDisplay);
+                waitAndAssertTopResumedActivity(TEST_ACTIVITY, newDisplay.mId,
+                        "Top activity must be on secondary display");
+                final SizeInfo testActivitySize = testActivitySession.getConfigInfo().sizeInfo;
 
-        assertEquals("Sizes of secondary display must not change after rotation of primary"
-                + " display", initialSize, testActivitySize);
+                assertEquals("Sizes of secondary display must not change after rotation of primary"
+                        + " display", initialSize, testActivitySize);
+            }
+        }
     }
 
     private void rotateAndCheckSameSizes(RotationSession rotationSession,
-            ActivitySession activitySession, SizeInfo initialSize) {
+            ActivitySession activitySession, SizeInfo initialSize) throws Exception {
         for (int rotation = 3; rotation >= 0; --rotation) {
             rotationSession.set(rotation);
             final SizeInfo rotatedSize = activitySession.getConfigInfo().sizeInfo;
@@ -241,39 +252,43 @@
      * on an external secondary display.
      */
     @Test
-    public void testExternalDisplayActivityTurnPrimaryOff() {
+    public void testExternalDisplayActivityTurnPrimaryOff() throws Exception {
         // Launch something on the primary display so we know there is a resumed activity there
         launchActivity(RESIZEABLE_ACTIVITY);
         waitAndAssertTopResumedActivity(RESIZEABLE_ACTIVITY, DEFAULT_DISPLAY,
                 "Activity launched on primary display must be resumed");
 
-        final ActivityDisplay newDisplay = createManagedExternalDisplaySession()
-                .setCanShowWithInsecureKeyguard(true).createVirtualDisplay();
+        try (final ExternalDisplaySession externalDisplaySession = new ExternalDisplaySession();
+             final PrimaryDisplayStateSession displayStateSession =
+                     new PrimaryDisplayStateSession()) {
+            final ActivityDisplay newDisplay = externalDisplaySession
+                    .setCanShowWithInsecureKeyguard(true).createVirtualDisplay();
 
-        launchActivityOnDisplay(TEST_ACTIVITY, newDisplay.mId);
+            launchActivityOnDisplay(TEST_ACTIVITY, newDisplay.mId);
 
-        // Check that the activity is launched onto the external display
-        waitAndAssertTopResumedActivity(TEST_ACTIVITY, newDisplay.mId,
-                "Activity launched on external display must be resumed");
-        mAmWmState.assertFocusedAppOnDisplay("App on default display must still be focused",
-                RESIZEABLE_ACTIVITY, DEFAULT_DISPLAY);
+            // Check that the activity is launched onto the external display
+            waitAndAssertTopResumedActivity(TEST_ACTIVITY, newDisplay.mId,
+                    "Activity launched on external display must be resumed");
+            mAmWmState.assertFocusedAppOnDisplay("App on default display must still be focused",
+                    RESIZEABLE_ACTIVITY, DEFAULT_DISPLAY);
 
-        separateTestJournal();
-        mObjectTracker.manage(new PrimaryDisplayStateSession()).turnScreenOff();
+            separateTestJournal();
+            displayStateSession.turnScreenOff();
 
-        // Wait for the fullscreen stack to start sleeping, and then make sure the
-        // test activity is still resumed.
-        final ActivityLifecycleCounts counts = new ActivityLifecycleCounts(RESIZEABLE_ACTIVITY);
-        if (!Condition.waitFor(counts.countWithRetry(RESIZEABLE_ACTIVITY + " to be stopped",
-                countSpec(ActivityCallback.ON_STOP, CountSpec.EQUALS, 1)))) {
-            fail(RESIZEABLE_ACTIVITY + " has received "
-                    + counts.getCount(ActivityCallback.ON_STOP)
-                    + " onStop() calls, expecting 1");
+            // Wait for the fullscreen stack to start sleeping, and then make sure the
+            // test activity is still resumed.
+            final ActivityLifecycleCounts counts = new ActivityLifecycleCounts(RESIZEABLE_ACTIVITY);
+            if (!Condition.waitFor(counts.countWithRetry(RESIZEABLE_ACTIVITY + " to be stopped",
+                    countSpec(ActivityCallback.ON_STOP, CountSpec.EQUALS, 1)))) {
+                fail(RESIZEABLE_ACTIVITY + " has received "
+                        + counts.getCount(ActivityCallback.ON_STOP)
+                        + " onStop() calls, expecting 1");
+            }
+            // For this test we create this virtual display with flag showContentWhenLocked, so it
+            // cannot be effected when default display screen off.
+            waitAndAssertTopResumedActivity(TEST_ACTIVITY, newDisplay.mId,
+                    "Activity launched on external display must be resumed");
         }
-        // For this test we create this virtual display with flag showContentWhenLocked, so it
-        // cannot be effected when default display screen off.
-        waitAndAssertTopResumedActivity(TEST_ACTIVITY, newDisplay.mId,
-                "Activity launched on external display must be resumed");
     }
 
     /**
@@ -281,29 +296,30 @@
      * on that display.
      */
     @Test
-    public void testExternalDisplayToggleState() {
-        final ExternalDisplaySession externalDisplaySession = createManagedExternalDisplaySession();
-        final ActivityDisplay newDisplay = externalDisplaySession.createVirtualDisplay();
+    public void testExternalDisplayToggleState() throws Exception {
+        try (final ExternalDisplaySession externalDisplaySession = new ExternalDisplaySession()) {
+            final ActivityDisplay newDisplay = externalDisplaySession.createVirtualDisplay();
 
-        launchActivityOnDisplay(TEST_ACTIVITY, newDisplay.mId);
+            launchActivityOnDisplay(TEST_ACTIVITY, newDisplay.mId);
 
-        // Check that the test activity is resumed on the external display
-        waitAndAssertTopResumedActivity(TEST_ACTIVITY, newDisplay.mId,
-                "Activity launched on external display must be resumed");
+            // Check that the test activity is resumed on the external display
+            waitAndAssertTopResumedActivity(TEST_ACTIVITY, newDisplay.mId,
+                    "Activity launched on external display must be resumed");
 
-        externalDisplaySession.turnDisplayOff();
+            externalDisplaySession.turnDisplayOff();
 
-        // Check that turning off the external display stops the activity, and makes it
-        // invisible.
-        waitAndAssertActivityState(TEST_ACTIVITY, STATE_STOPPED,
-                "Activity launched on external display must be stopped after turning off");
-        mAmWmState.assertVisibility(TEST_ACTIVITY, false /* visible */);
+            // Check that turning off the external display stops the activity, and makes it
+            // invisible.
+            waitAndAssertActivityState(TEST_ACTIVITY, STATE_STOPPED,
+                    "Activity launched on external display must be stopped after turning off");
+            mAmWmState.assertVisibility(TEST_ACTIVITY, false /* visible */);
 
-        externalDisplaySession.turnDisplayOn();
+            externalDisplaySession.turnDisplayOn();
 
-        // Check that turning on the external display resumes the activity
-        waitAndAssertTopResumedActivity(TEST_ACTIVITY, newDisplay.mId,
-                "Activity launched on external display must be resumed");
+            // Check that turning on the external display resumes the activity
+            waitAndAssertTopResumedActivity(TEST_ACTIVITY, newDisplay.mId,
+                    "Activity launched on external display must be resumed");
+        }
     }
 
     /**
@@ -340,57 +356,69 @@
      * visibility is not affected.
      */
     @Test
-    public void testLaunchActivitiesAffectsVisibility() {
+    public void testLaunchActivitiesAffectsVisibility() throws Exception {
         // Start launching activity.
         launchActivity(LAUNCHING_ACTIVITY);
 
-        // Create new virtual display.
-        final ActivityDisplay newDisplay = createManagedVirtualDisplaySession().createDisplay();
-        mAmWmState.assertVisibility(VIRTUAL_DISPLAY_ACTIVITY, true /* visible */);
+        try (final VirtualDisplaySession virtualDisplaySession = new VirtualDisplaySession()) {
+            // Create new virtual display.
+            final ActivityDisplay newDisplay = virtualDisplaySession.createDisplay();
+            mAmWmState.assertVisibility(VIRTUAL_DISPLAY_ACTIVITY, true /* visible */);
 
-        // Launch activity on new secondary display.
-        launchActivityOnDisplay(TEST_ACTIVITY, newDisplay.mId);
-        mAmWmState.assertVisibility(TEST_ACTIVITY, true /* visible */);
-        mAmWmState.assertVisibility(VIRTUAL_DISPLAY_ACTIVITY, true /* visible */);
+            // Launch activity on new secondary display.
+            launchActivityOnDisplay(TEST_ACTIVITY, newDisplay.mId);
+            mAmWmState.assertVisibility(TEST_ACTIVITY, true /* visible */);
+            mAmWmState.assertVisibility(VIRTUAL_DISPLAY_ACTIVITY, true /* visible */);
 
-        // Launch activity on primary display and check if it doesn't affect activity on
-        // secondary display.
-        getLaunchActivityBuilder().setTargetActivity(RESIZEABLE_ACTIVITY).execute();
-        mAmWmState.waitForValidState(RESIZEABLE_ACTIVITY);
-        mAmWmState.assertVisibility(TEST_ACTIVITY, true /* visible */);
-        mAmWmState.assertVisibility(RESIZEABLE_ACTIVITY, true /* visible */);
-        assertBothDisplaysHaveResumedActivities(pair(DEFAULT_DISPLAY, RESIZEABLE_ACTIVITY),
-                pair(newDisplay.mId, TEST_ACTIVITY));
+            // Launch activity on primary display and check if it doesn't affect activity on
+            // secondary display.
+            getLaunchActivityBuilder().setTargetActivity(RESIZEABLE_ACTIVITY).execute();
+            mAmWmState.waitForValidState(RESIZEABLE_ACTIVITY);
+            mAmWmState.assertVisibility(TEST_ACTIVITY, true /* visible */);
+            mAmWmState.assertVisibility(RESIZEABLE_ACTIVITY, true /* visible */);
+            mAmWmState.assertResumedActivities("Both displays must have resumed activities",
+                    new SparseArray<ComponentName>(){{
+                        put(DEFAULT_DISPLAY, RESIZEABLE_ACTIVITY);
+                        put(newDisplay.mId, TEST_ACTIVITY);
+                    }}
+            );
+        }
     }
 
     /**
      * Test that move-task works when moving between displays.
      */
     @Test
-    public void testMoveTaskBetweenDisplays() {
-        // Create new virtual display.
-        final ActivityDisplay newDisplay = createManagedVirtualDisplaySession().createDisplay();
-        mAmWmState.assertVisibility(VIRTUAL_DISPLAY_ACTIVITY, true /* visible */);
-        mAmWmState.assertFocusedActivity("Virtual display activity must be on top",
-                VIRTUAL_DISPLAY_ACTIVITY);
-        final int defaultDisplayStackId = mAmWmState.getAmState().getFocusedStackId();
-        ActivityStack frontStack = mAmWmState.getAmState().getStackById(
-                defaultDisplayStackId);
-        assertEquals("Top stack must remain on primary display",
-                DEFAULT_DISPLAY, frontStack.mDisplayId);
+    public void testMoveTaskBetweenDisplays() throws Exception {
+        try (final VirtualDisplaySession virtualDisplaySession = new VirtualDisplaySession()) {
+            // Create new virtual display.
+            final ActivityDisplay newDisplay = virtualDisplaySession.createDisplay();
+            mAmWmState.assertVisibility(VIRTUAL_DISPLAY_ACTIVITY, true /* visible */);
+            mAmWmState.assertFocusedActivity("Virtual display activity must be on top",
+                    VIRTUAL_DISPLAY_ACTIVITY);
+            final int defaultDisplayStackId = mAmWmState.getAmState().getFocusedStackId();
+            ActivityStack frontStack = mAmWmState.getAmState().getStackById(
+                    defaultDisplayStackId);
+            assertEquals("Top stack must remain on primary display",
+                    DEFAULT_DISPLAY, frontStack.mDisplayId);
 
-        // Launch activity on new secondary display.
-        launchActivityOnDisplay(TEST_ACTIVITY, newDisplay.mId);
+            // Launch activity on new secondary display.
+            launchActivityOnDisplay(TEST_ACTIVITY, newDisplay.mId);
 
-        waitAndAssertTopResumedActivity(TEST_ACTIVITY, newDisplay.mId,
-                "Top activity must be on secondary display");
-        assertBothDisplaysHaveResumedActivities(pair(DEFAULT_DISPLAY, VIRTUAL_DISPLAY_ACTIVITY),
-                pair(newDisplay.mId, TEST_ACTIVITY));
+            waitAndAssertTopResumedActivity(TEST_ACTIVITY, newDisplay.mId,
+                    "Top activity must be on secondary display");
+            mAmWmState.assertResumedActivities("Both displays must have resumed activities",
+                    new SparseArray<ComponentName>(){{
+                        put(DEFAULT_DISPLAY, VIRTUAL_DISPLAY_ACTIVITY);
+                        put(newDisplay.mId, TEST_ACTIVITY);
+                    }}
+            );
 
-        // Move activity from secondary display to primary.
-        moveActivityToStack(TEST_ACTIVITY, defaultDisplayStackId);
-        waitAndAssertTopResumedActivity(TEST_ACTIVITY, DEFAULT_DISPLAY,
-                "Moved activity must be on top");
+            // Move activity from secondary display to primary.
+            moveActivityToStack(TEST_ACTIVITY, defaultDisplayStackId);
+            waitAndAssertTopResumedActivity(TEST_ACTIVITY, DEFAULT_DISPLAY,
+                    "Moved activity must be on top");
+        }
     }
 
     /**
@@ -399,7 +427,7 @@
      * This version launches virtual display creator to fullscreen stack in split-screen.
      */
     @Test
-    public void testStackFocusSwitchOnDisplayRemoved() {
+    public void testStackFocusSwitchOnDisplayRemoved() throws Exception {
         assumeTrue(supportsSplitScreenMultiWindow());
 
         // Start launching activity into docked stack.
@@ -418,7 +446,7 @@
      * This version launches virtual display creator to docked stack in split-screen.
      */
     @Test
-    public void testStackFocusSwitchOnDisplayRemoved2() {
+    public void testStackFocusSwitchOnDisplayRemoved2() throws Exception {
         assumeTrue(supportsSplitScreenMultiWindow());
 
         // Setup split-screen.
@@ -437,7 +465,7 @@
      * This version works without split-screen.
      */
     @Test
-    public void testStackFocusSwitchOnDisplayRemoved3() {
+    public void testStackFocusSwitchOnDisplayRemoved3() throws Exception {
         // Start an activity on default display to determine default stack.
         launchActivity(BROADCAST_RECEIVER_ACTIVITY);
         final int focusedStackWindowingMode = mAmWmState.getAmState().getFrontStackWindowingMode(
@@ -453,7 +481,8 @@
      * Create a virtual display, launch a test activity there, destroy the display and check if test
      * activity is moved to a stack on the default display.
      */
-    private void tryCreatingAndRemovingDisplayWithActivity(boolean splitScreen, int windowingMode) {
+    private void tryCreatingAndRemovingDisplayWithActivity(boolean splitScreen, int windowingMode)
+            throws Exception {
         try (final VirtualDisplaySession virtualDisplaySession = new VirtualDisplaySession()) {
             // Create new virtual display.
             final ActivityDisplay newDisplay = virtualDisplaySession
@@ -500,9 +529,11 @@
      * is moved correctly.
      */
     @Test
-    public void testStackFocusSwitchOnStackEmptiedInSleeping() {
-        validateStackFocusSwitchOnStackEmptied(createManagedVirtualDisplaySession(),
-                createManagedLockScreenSession());
+    public void testStackFocusSwitchOnStackEmptiedInSleeping() throws Exception {
+        try (final VirtualDisplaySession virtualDisplaySession = new VirtualDisplaySession();
+             final LockScreenSession lockScreenSession = new LockScreenSession()) {
+            validateStackFocusSwitchOnStackEmptied(virtualDisplaySession, lockScreenSession);
+        }
     }
 
     /**
@@ -510,16 +541,19 @@
      * is moved correctly.
      */
     @Test
-    public void testStackFocusSwitchOnStackEmptied() {
-        validateStackFocusSwitchOnStackEmptied(createManagedVirtualDisplaySession(),
-                null /* lockScreenSession */);
+    public void testStackFocusSwitchOnStackEmptied() throws Exception {
+        try (final VirtualDisplaySession virtualDisplaySession = new VirtualDisplaySession()) {
+            validateStackFocusSwitchOnStackEmptied(virtualDisplaySession,
+                    null /* lockScreenSession */);
+        }
     }
 
     private void validateStackFocusSwitchOnStackEmptied(VirtualDisplaySession virtualDisplaySession,
-            LockScreenSession lockScreenSession) {
+            LockScreenSession lockScreenSession) throws Exception {
         // Create new virtual display.
         final ActivityDisplay newDisplay = virtualDisplaySession.createDisplay();
         mAmWmState.assertVisibility(VIRTUAL_DISPLAY_ACTIVITY, true /* visible */);
+        final int focusedStackId = mAmWmState.getAmState().getFrontStackId(DEFAULT_DISPLAY);
 
         // Launch activity on new secondary display.
         launchActivityOnDisplay(BROADCAST_RECEIVER_ACTIVITY, newDisplay.mId);
@@ -547,31 +581,37 @@
      * Tests that input events on the primary display take focus from the virtual display.
      */
     @Test
-    public void testStackFocusSwitchOnTouchEvent() {
+    public void testStackFocusSwitchOnTouchEvent() throws Exception {
         // If config_perDisplayFocusEnabled, the focus will not move even if touching on
         // the Activity in the different display.
         assumeFalse(perDisplayFocusEnabled());
 
-        // Create new virtual display.
-        final ActivityDisplay newDisplay = createManagedVirtualDisplaySession().createDisplay();
+        try (final VirtualDisplaySession virtualDisplaySession = new VirtualDisplaySession()) {
+            // Create new virtual display.
+            final ActivityDisplay newDisplay = virtualDisplaySession.createDisplay();
 
-        mAmWmState.computeState(VIRTUAL_DISPLAY_ACTIVITY);
-        mAmWmState.assertFocusedActivity("Top activity must be the latest launched one",
-                VIRTUAL_DISPLAY_ACTIVITY);
+            mAmWmState.computeState(VIRTUAL_DISPLAY_ACTIVITY);
+            mAmWmState.assertFocusedActivity("Top activity must be the latest launched one",
+                    VIRTUAL_DISPLAY_ACTIVITY);
 
-        launchActivityOnDisplay(TEST_ACTIVITY, newDisplay.mId);
+            launchActivityOnDisplay(TEST_ACTIVITY, newDisplay.mId);
 
-        waitAndAssertTopResumedActivity(TEST_ACTIVITY, newDisplay.mId,
-                "Activity launched on secondary display must be on top");
+            waitAndAssertTopResumedActivity(TEST_ACTIVITY, newDisplay.mId,
+                    "Activity launched on secondary display must be on top");
 
-        tapOnDisplayCenter(DEFAULT_DISPLAY);
+            tapOnDisplayCenter(DEFAULT_DISPLAY);
 
-        waitAndAssertTopResumedActivity(VIRTUAL_DISPLAY_ACTIVITY, DEFAULT_DISPLAY,
-                "Top activity must be on the primary display");
-        assertBothDisplaysHaveResumedActivities(pair(DEFAULT_DISPLAY, VIRTUAL_DISPLAY_ACTIVITY),
-                pair(newDisplay.mId, TEST_ACTIVITY));
-        mAmWmState.assertFocusedAppOnDisplay("App on secondary display must still be focused",
-                TEST_ACTIVITY, newDisplay.mId);
+            waitAndAssertTopResumedActivity(VIRTUAL_DISPLAY_ACTIVITY, DEFAULT_DISPLAY,
+                    "Top activity must be on the primary display");
+            mAmWmState.assertResumedActivities("Both displays must have resumed activities",
+                    new SparseArray<ComponentName>(){{
+                        put(DEFAULT_DISPLAY, VIRTUAL_DISPLAY_ACTIVITY);
+                        put(newDisplay.mId, TEST_ACTIVITY);
+                    }}
+            );
+            mAmWmState.assertFocusedAppOnDisplay("App on secondary display must still be focused",
+                    TEST_ACTIVITY, newDisplay.mId);
+        }
     }
 
 
@@ -580,7 +620,7 @@
      * activity on the primary display.
      */
     @Test
-    public void testStackFocusSwitchOnTouchEventAfterKeyguard() {
+    public void testStackFocusSwitchOnTouchEventAfterKeyguard() throws Exception {
         assumeFalse(perDisplayFocusEnabled());
 
         // Launch something on the primary display so we know there is a resumed activity there
@@ -588,120 +628,133 @@
         waitAndAssertTopResumedActivity(RESIZEABLE_ACTIVITY, DEFAULT_DISPLAY,
                 "Activity launched on primary display must be resumed");
 
-        final LockScreenSession lockScreenSession = createManagedLockScreenSession();
-        lockScreenSession.sleepDevice();
+        try (final LockScreenSession lockScreenSession = new LockScreenSession();
+             final ExternalDisplaySession externalDisplaySession = new ExternalDisplaySession()) {
+            lockScreenSession.sleepDevice();
 
-        // Make sure there is no resumed activity when the primary display is off
-        waitAndAssertActivityState(RESIZEABLE_ACTIVITY, STATE_STOPPED,
-                "Activity launched on primary display must be stopped after turning off");
-        assertEquals("Unexpected resumed activity",
-                0, mAmWmState.getAmState().getResumedActivitiesCount());
+            // Make sure there is no resumed activity when the primary display is off
+            waitAndAssertActivityState(RESIZEABLE_ACTIVITY, STATE_STOPPED,
+                    "Activity launched on primary display must be stopped after turning off");
+            assertEquals("Unexpected resumed activity",
+                    0, mAmWmState.getAmState().getResumedActivitiesCount());
 
-        final ActivityDisplay newDisplay = createManagedExternalDisplaySession()
-                .setCanShowWithInsecureKeyguard(true).createVirtualDisplay();
+            final ActivityDisplay newDisplay = externalDisplaySession
+                    .setCanShowWithInsecureKeyguard(true).createVirtualDisplay();
 
-        launchActivityOnDisplay(TEST_ACTIVITY, newDisplay.mId);
+            launchActivityOnDisplay(TEST_ACTIVITY, newDisplay.mId);
 
-        // Unlock the device and tap on the middle of the primary display
-        lockScreenSession.wakeUpDevice();
-        executeShellCommand("wm dismiss-keyguard");
-        mAmWmState.waitForKeyguardGone();
-        mAmWmState.waitForValidState(RESIZEABLE_ACTIVITY, TEST_ACTIVITY);
+            // Unlock the device and tap on the middle of the primary display
+            lockScreenSession.wakeUpDevice();
+            executeShellCommand("wm dismiss-keyguard");
+            mAmWmState.waitForKeyguardGone();
+            mAmWmState.waitForValidState(RESIZEABLE_ACTIVITY, TEST_ACTIVITY);
 
-        // Check that the test activity is resumed on the external display and is on top
-        waitAndAssertTopResumedActivity(TEST_ACTIVITY, newDisplay.mId,
-                "Activity on external display must be resumed and on top");
-        assertBothDisplaysHaveResumedActivities(pair(DEFAULT_DISPLAY, RESIZEABLE_ACTIVITY),
-                pair(newDisplay.mId, TEST_ACTIVITY));
+            // Check that the test activity is resumed on the external display and is on top
+            waitAndAssertTopResumedActivity(TEST_ACTIVITY, newDisplay.mId,
+                    "Activity on external display must be resumed and on top");
+            mAmWmState.assertResumedActivities("Both displays should have resumed activities",
+                    new SparseArray<ComponentName>(){{
+                        put(DEFAULT_DISPLAY, RESIZEABLE_ACTIVITY);
+                        put(newDisplay.mId, TEST_ACTIVITY);
+                    }}
+            );
 
-        tapOnDisplayCenter(DEFAULT_DISPLAY);
+            tapOnDisplayCenter(DEFAULT_DISPLAY);
 
-        // Check that the activity on the primary display is the topmost resumed
-        waitAndAssertTopResumedActivity(RESIZEABLE_ACTIVITY, DEFAULT_DISPLAY,
-                "Activity on primary display must be resumed and on top");
-        assertBothDisplaysHaveResumedActivities(pair(DEFAULT_DISPLAY, RESIZEABLE_ACTIVITY),
-                pair(newDisplay.mId, TEST_ACTIVITY));
-        mAmWmState.assertFocusedAppOnDisplay("App on external display must still be focused",
-                TEST_ACTIVITY, newDisplay.mId);
+            // Check that the activity on the primary display is the topmost resumed
+            waitAndAssertTopResumedActivity(RESIZEABLE_ACTIVITY, DEFAULT_DISPLAY,
+                    "Activity on primary display must be resumed and on top");
+            mAmWmState.assertResumedActivities("Both displays should have resumed activities",
+                    new SparseArray<ComponentName>(){{
+                        put(DEFAULT_DISPLAY, RESIZEABLE_ACTIVITY);
+                        put(newDisplay.mId, TEST_ACTIVITY);
+                    }}
+            );
+            mAmWmState.assertFocusedAppOnDisplay("App on external display must still be focused",
+                    TEST_ACTIVITY, newDisplay.mId);
+        }
     }
 
     /**
      * Tests that showWhenLocked works on a secondary display.
      */
     @Test
-    public void testSecondaryDisplayShowWhenLocked() {
+    public void testSecondaryDisplayShowWhenLocked() throws Exception {
         assumeTrue(supportsSecureLock());
 
-        final LockScreenSession lockScreenSession = createManagedLockScreenSession();
-        lockScreenSession.setLockCredential();
+        try (final ExternalDisplaySession externalDisplaySession = new ExternalDisplaySession();
+             final LockScreenSession lockScreenSession = new LockScreenSession()) {
+            lockScreenSession.setLockCredential();
 
-        launchActivity(TEST_ACTIVITY);
+            launchActivity(TEST_ACTIVITY);
 
-        final ActivityDisplay newDisplay = createManagedExternalDisplaySession()
-                .createVirtualDisplay();
-        launchActivityOnDisplay(SHOW_WHEN_LOCKED_ATTR_ACTIVITY, newDisplay.mId);
+            final ActivityDisplay newDisplay = externalDisplaySession.createVirtualDisplay();
+            launchActivityOnDisplay(SHOW_WHEN_LOCKED_ATTR_ACTIVITY, newDisplay.mId);
 
-        lockScreenSession.gotoKeyguard();
+            lockScreenSession.gotoKeyguard();
 
-        waitAndAssertActivityState(TEST_ACTIVITY, STATE_STOPPED,
-                "Expected stopped activity on default display");
-        waitAndAssertTopResumedActivity(SHOW_WHEN_LOCKED_ATTR_ACTIVITY, newDisplay.mId,
-                "Expected resumed activity on secondary display");
+            waitAndAssertActivityState(TEST_ACTIVITY, STATE_STOPPED,
+                    "Expected stopped activity on default display");
+            waitAndAssertTopResumedActivity(SHOW_WHEN_LOCKED_ATTR_ACTIVITY, newDisplay.mId,
+                    "Expected resumed activity on secondary display");
+        }
     }
 
     /**
      * Tests tap and set focus between displays.
      */
     @Test
-    public void testSecondaryDisplayFocus() {
+    public void testSecondaryDisplayFocus() throws Exception {
         assumeFalse(perDisplayFocusEnabled());
 
-        launchActivity(TEST_ACTIVITY);
-        mAmWmState.waitForActivityState(TEST_ACTIVITY, STATE_RESUMED);
+        try (final ExternalDisplaySession externalDisplaySession = new ExternalDisplaySession()) {
+            launchActivity(TEST_ACTIVITY);
+            mAmWmState.waitForActivityState(TEST_ACTIVITY, STATE_RESUMED);
 
-        final ActivityDisplay newDisplay = createManagedExternalDisplaySession()
-                .createVirtualDisplay();
-        launchActivityOnDisplay(VIRTUAL_DISPLAY_ACTIVITY, newDisplay.mId);
-        waitAndAssertTopResumedActivity(VIRTUAL_DISPLAY_ACTIVITY, newDisplay.mId,
-                "Virtual activity should be Top Resumed Activity.");
-        mAmWmState.assertFocusedAppOnDisplay("Activity on second display must be focused.",
-                VIRTUAL_DISPLAY_ACTIVITY, newDisplay.mId);
+            final ActivityDisplay newDisplay = externalDisplaySession.createVirtualDisplay();
+            launchActivityOnDisplay(VIRTUAL_DISPLAY_ACTIVITY, newDisplay.mId);
+            waitAndAssertTopResumedActivity(VIRTUAL_DISPLAY_ACTIVITY, newDisplay.mId,
+                    "Virtual activity should be Top Resumed Activity.");
+            mAmWmState.assertFocusedAppOnDisplay("Activity on second display must be focused.",
+                    VIRTUAL_DISPLAY_ACTIVITY, newDisplay.mId);
 
-        tapOnDisplayCenter(DEFAULT_DISPLAY);
+            tapOnDisplayCenter(DEFAULT_DISPLAY);
 
-        waitAndAssertTopResumedActivity(TEST_ACTIVITY, DEFAULT_DISPLAY,
-                "Activity should be top resumed when tapped.");
-        mAmWmState.assertFocusedActivity("Activity on default display must be top focused.",
-                TEST_ACTIVITY);
+            waitAndAssertTopResumedActivity(TEST_ACTIVITY, DEFAULT_DISPLAY,
+                    "Activity should be top resumed when tapped.");
+            mAmWmState.assertFocusedActivity("Activity on default display must be top focused.",
+                    TEST_ACTIVITY);
 
-        tapOnDisplayCenter(newDisplay.mId);
+            tapOnDisplayCenter(newDisplay.mId);
 
-        waitAndAssertTopResumedActivity(VIRTUAL_DISPLAY_ACTIVITY, newDisplay.mId,
-                "Virtual display activity should be top resumed when tapped.");
-        mAmWmState.assertFocusedActivity("Activity on second display must be top focused.",
-                VIRTUAL_DISPLAY_ACTIVITY);
-        mAmWmState.assertFocusedAppOnDisplay(
-                "Activity on default display must be still focused.",
-                TEST_ACTIVITY, DEFAULT_DISPLAY);
+            waitAndAssertTopResumedActivity(VIRTUAL_DISPLAY_ACTIVITY, newDisplay.mId,
+                    "Virtual display activity should be top resumed when tapped.");
+            mAmWmState.assertFocusedActivity("Activity on second display must be top focused.",
+                    VIRTUAL_DISPLAY_ACTIVITY);
+            mAmWmState.assertFocusedAppOnDisplay(
+                    "Activity on default display must be still focused.",
+                    TEST_ACTIVITY, DEFAULT_DISPLAY);
+        }
     }
 
     /**
      * Tests that toast works on a secondary display.
      */
     @Test
-    public void testSecondaryDisplayShowToast() {
-        final ActivityDisplay newDisplay = createManagedVirtualDisplaySession()
-                .setPublicDisplay(true)
-                .createDisplay();
-        final String TOAST_NAME = "Toast";
-        launchActivityOnDisplay(TOAST_ACTIVITY, newDisplay.mId);
-        waitAndAssertTopResumedActivity(TOAST_ACTIVITY, newDisplay.mId,
-                "Activity launched on external display must be resumed");
+    public void testSecondaryDisplayShowToast() throws Exception {
+        try (final VirtualDisplaySession virtualDisplaySession = new VirtualDisplaySession()) {
+            final ActivityDisplay newDisplay =
+                    virtualDisplaySession.setPublicDisplay(true).createDisplay();
+            final String TOAST_NAME = "Toast";
+            launchActivityOnDisplay(TOAST_ACTIVITY, newDisplay.mId);
+            waitAndAssertTopResumedActivity(TOAST_ACTIVITY, newDisplay.mId,
+                    "Activity launched on external display must be resumed");
 
-        assertTrue("Toast window must be shown", mAmWmState.waitForWithWmState(
-                state -> state.containsWindow(TOAST_NAME), "toast window to show"));
-        assertTrue("Toast window must be visible",
-                mAmWmState.getWmState().isWindowVisible(TOAST_NAME));
+            assertTrue("Toast window must be shown", mAmWmState.waitForWithWmState(
+                    state -> state.containsWindow(TOAST_NAME), "toast window to show"));
+            assertTrue("Toast window must be visible",
+                    mAmWmState.getWmState().isWindowVisible(TOAST_NAME));
+        }
     }
 
     /**
@@ -709,7 +762,7 @@
      * Also check that the surface size has updated after reparenting to other display.
      */
     @Test
-    public void testTaskSurfaceSizeAfterReparentDisplay() {
+    public void testTaskSurfaceSizeAfterReparentDisplay() throws Exception {
         try (final VirtualDisplaySession virtualDisplaySession = new VirtualDisplaySession()) {
             // Create new simulated display and launch an activity on it.
             final ActivityDisplay newDisplay = virtualDisplaySession.setSimulateDisplay(true)
@@ -746,99 +799,105 @@
     }
 
     @Test
-    public void testAppTransitionForActivityOnDifferentDisplay() {
-        final TestActivitySession<StandardActivity> transitionActivitySession =
-                createManagedTestActivitySession();
-        // Create new simulated display.
-        final ActivityDisplay newDisplay = createManagedVirtualDisplaySession()
-                .setSimulateDisplay(true).createDisplay();
+    public void testAppTransitionForActivityOnDifferentDisplay() throws Exception {
+        try (final VirtualDisplaySession virtualDisplaySession = new VirtualDisplaySession();
+             final TestActivitySession<StandardActivity> transitionActivitySession = new
+                     TestActivitySession<>()) {
+            // Create new simulated display.
+            final ActivityDisplay newDisplay = virtualDisplaySession
+                    .setSimulateDisplay(true).createDisplay();
 
-        // Launch BottomActivity on top of launcher activity to prevent transition state
-        // affected by wallpaper theme.
-        launchActivityOnDisplay(BOTTOM_ACTIVITY, DEFAULT_DISPLAY);
-        waitAndAssertTopResumedActivity(BOTTOM_ACTIVITY, DEFAULT_DISPLAY,
-                "Activity must be resumed");
+            // Launch BottomActivity on top of launcher activity to prevent transition state
+            // affected by wallpaper theme.
+            launchActivityOnDisplay(BOTTOM_ACTIVITY, DEFAULT_DISPLAY);
+            waitAndAssertTopResumedActivity(BOTTOM_ACTIVITY, DEFAULT_DISPLAY,
+                    "Activity must be resumed");
 
-        // Launch StandardActivity on default display, verify last transition if is correct.
-        transitionActivitySession.launchTestActivityOnDisplaySync(StandardActivity.class,
-                DEFAULT_DISPLAY);
-        mAmWmState.waitForAppTransitionIdleOnDisplay(DEFAULT_DISPLAY);
-        mAmWmState.assertSanity();
-        assertEquals(TRANSIT_TASK_OPEN,
-                mAmWmState.getWmState().getDisplay(DEFAULT_DISPLAY).getLastTransition());
+            // Launch StandardActivity on default display, verify last transition if is correct.
+            transitionActivitySession.launchTestActivityOnDisplaySync(StandardActivity.class,
+                    DEFAULT_DISPLAY);
+            mAmWmState.waitForAppTransitionIdleOnDisplay(DEFAULT_DISPLAY);
+            mAmWmState.assertSanity();
+            assertEquals(TRANSIT_TASK_OPEN,
+                    mAmWmState.getWmState().getDisplay(DEFAULT_DISPLAY).getLastTransition());
 
-        // Finish current activity & launch another TestActivity in virtual display in parallel.
-        transitionActivitySession.finishCurrentActivityNoWait();
-        launchActivityOnDisplayNoWait(TEST_ACTIVITY, newDisplay.mId);
-        mAmWmState.waitForAppTransitionIdleOnDisplay(DEFAULT_DISPLAY);
-        mAmWmState.waitForAppTransitionIdleOnDisplay(newDisplay.mId);
-        mAmWmState.assertSanity();
+            // Finish current activity & launch another TestActivity in virtual display in parallel.
+            transitionActivitySession.finishCurrentActivityNoWait();
+            launchActivityOnDisplayNoWait(TEST_ACTIVITY, newDisplay.mId);
+            mAmWmState.waitForAppTransitionIdleOnDisplay(DEFAULT_DISPLAY);
+            mAmWmState.waitForAppTransitionIdleOnDisplay(newDisplay.mId);
+            mAmWmState.assertSanity();
 
-        // Verify each display's last transition if is correct as expected.
-        assertEquals(TRANSIT_TASK_CLOSE,
-                mAmWmState.getWmState().getDisplay(DEFAULT_DISPLAY).getLastTransition());
-        assertEquals(TRANSIT_TASK_OPEN,
-                mAmWmState.getWmState().getDisplay(newDisplay.mId).getLastTransition());
+            // Verify each display's last transition if is correct as expected.
+            assertEquals(TRANSIT_TASK_CLOSE,
+                    mAmWmState.getWmState().getDisplay(DEFAULT_DISPLAY).getLastTransition());
+            assertEquals(TRANSIT_TASK_OPEN,
+                    mAmWmState.getWmState().getDisplay(newDisplay.mId).getLastTransition());
+        }
     }
 
     @Test
     public void testNoTransitionWhenMovingActivityToDisplay() throws Exception {
-        // Create new simulated display & capture new display's transition state.
-        final ActivityDisplay newDisplay = createManagedVirtualDisplaySession()
-                .setSimulateDisplay(true).createDisplay();
+        try (final VirtualDisplaySession virtualDisplaySession = new VirtualDisplaySession()) {
+            // Create new simulated display & capture new display's transition state.
+            final ActivityDisplay newDisplay = virtualDisplaySession
+                    .setSimulateDisplay(true).createDisplay();
 
-        // Launch TestActivity in virtual display & capture its transition state.
-        launchActivityOnDisplay(TEST_ACTIVITY, newDisplay.mId);
-        mAmWmState.waitForAppTransitionIdleOnDisplay(newDisplay.mId);
-        mAmWmState.assertSanity();
-        final String lastTranstionOnVirtualDisplay = mAmWmState.getWmState()
-                .getDisplay(newDisplay.mId).getLastTransition();
+            // Launch TestActivity in virtual display & capture its transition state.
+            launchActivityOnDisplay(TEST_ACTIVITY, newDisplay.mId);
+            mAmWmState.waitForAppTransitionIdleOnDisplay(newDisplay.mId);
+            mAmWmState.assertSanity();
+            final String lastTranstionOnVirtualDisplay = mAmWmState.getWmState()
+                    .getDisplay(newDisplay.mId).getLastTransition();
 
-        // Move TestActivity from virtual display to default display.
-        getLaunchActivityBuilder().setTargetActivity(TEST_ACTIVITY)
-                .allowMultipleInstances(false).setNewTask(true)
-                .setDisplayId(DEFAULT_DISPLAY).execute();
+            // Move TestActivity from virtual display to default display.
+            getLaunchActivityBuilder().setTargetActivity(TEST_ACTIVITY)
+                    .allowMultipleInstances(false).setNewTask(true)
+                    .setDisplayId(DEFAULT_DISPLAY).execute();
 
-        // Verify TestActivity moved to virtual display.
-        waitAndAssertTopResumedActivity(TEST_ACTIVITY, DEFAULT_DISPLAY,
-                "Existing task must be brought to front");
+            // Verify TestActivity moved to virtual display.
+            waitAndAssertTopResumedActivity(TEST_ACTIVITY, DEFAULT_DISPLAY,
+                    "Existing task must be brought to front");
 
-        // Make sure last transition will not change when task move to another display.
-        assertEquals(lastTranstionOnVirtualDisplay,
-                mAmWmState.getWmState().getDisplay(newDisplay.mId).getLastTransition());
+            // Make sure last transition will not change when task move to another display.
+            assertEquals(lastTranstionOnVirtualDisplay,
+                    mAmWmState.getWmState().getDisplay(newDisplay.mId).getLastTransition());
+        }
     }
 
     @Test
-    public void testPreQTopProcessResumedActivity() {
-        final ActivityDisplay newDisplay = createManagedVirtualDisplaySession()
-                .setSimulateDisplay(true).createDisplay();
+    public void testPreQTopProcessResumedActivity() throws Exception {
+        try (final VirtualDisplaySession virtualDisplaySession = new VirtualDisplaySession()) {
+            final ActivityDisplay newDisplay =
+                    virtualDisplaySession.setSimulateDisplay(true).createDisplay();
 
-        getLaunchActivityBuilder().setUseInstrumentation()
-                .setTargetActivity(SDK_27_TEST_ACTIVITY).setNewTask(true)
-                .setDisplayId(newDisplay.mId).execute();
-        waitAndAssertTopResumedActivity(SDK_27_TEST_ACTIVITY, newDisplay.mId,
-                "Activity launched on secondary display must be resumed and focused");
+            getLaunchActivityBuilder().setUseInstrumentation()
+                    .setTargetActivity(SDK_27_TEST_ACTIVITY).setNewTask(true)
+                    .setDisplayId(newDisplay.mId).execute();
+            waitAndAssertTopResumedActivity(SDK_27_TEST_ACTIVITY, newDisplay.mId,
+                    "Activity launched on secondary display must be resumed and focused");
 
-        getLaunchActivityBuilder().setUseInstrumentation()
-                .setTargetActivity(SDK_27_LAUNCHING_ACTIVITY).setNewTask(true)
-                .setDisplayId(DEFAULT_DISPLAY).execute();
-        waitAndAssertTopResumedActivity(SDK_27_LAUNCHING_ACTIVITY, DEFAULT_DISPLAY,
-                "Activity launched on default display must be resumed and focused");
+            getLaunchActivityBuilder().setUseInstrumentation()
+                    .setTargetActivity(SDK_27_LAUNCHING_ACTIVITY).setNewTask(true)
+                    .setDisplayId(DEFAULT_DISPLAY).execute();
+            waitAndAssertTopResumedActivity(SDK_27_LAUNCHING_ACTIVITY, DEFAULT_DISPLAY,
+                    "Activity launched on default display must be resumed and focused");
 
-        assertEquals("There must be only one resumed activity in the package.", 1,
-                mAmWmState.getAmState().getResumedActivitiesCountInPackage(
-                        SDK_27_LAUNCHING_ACTIVITY.getPackageName()));
+            assertEquals("There must be only one resumed activity in the package.", 1,
+                    mAmWmState.getAmState().getResumedActivitiesCountInPackage(
+                            SDK_27_LAUNCHING_ACTIVITY.getPackageName()));
 
-        getLaunchActivityBuilder().setUseInstrumentation()
-                .setTargetActivity(SDK_27_SEPARATE_PROCESS_ACTIVITY).setNewTask(true)
-                .setDisplayId(DEFAULT_DISPLAY).execute();
-        waitAndAssertTopResumedActivity(SDK_27_SEPARATE_PROCESS_ACTIVITY, DEFAULT_DISPLAY,
-                "Activity launched on default display must be resumed and focused");
-        assertTrue("Activity that was on secondary display must be resumed",
-                mAmWmState.getAmState().hasActivityState(SDK_27_TEST_ACTIVITY, STATE_RESUMED));
-        assertEquals("There must be only two resumed activities in the package.", 2,
-                mAmWmState.getAmState().getResumedActivitiesCountInPackage(
-                        SDK_27_TEST_ACTIVITY.getPackageName()));
+            getLaunchActivityBuilder().setUseInstrumentation()
+                    .setTargetActivity(SDK_27_SEPARATE_PROCESS_ACTIVITY).setNewTask(true)
+                    .setDisplayId(DEFAULT_DISPLAY).execute();
+            waitAndAssertTopResumedActivity(SDK_27_SEPARATE_PROCESS_ACTIVITY, DEFAULT_DISPLAY,
+                    "Activity launched on default display must be resumed and focused");
+            assertTrue("Activity that was on secondary display must be resumed",
+                    mAmWmState.getAmState().hasActivityState(SDK_27_TEST_ACTIVITY, STATE_RESUMED));
+            assertEquals("There must be only two resumed activities in the package.", 2,
+                    mAmWmState.getAmState().getResumedActivitiesCountInPackage(
+                            SDK_27_TEST_ACTIVITY.getPackageName()));
+        }
     }
 
 }
diff --git a/tests/framework/base/windowmanager/src/android/server/wm/MultiDisplaySecurityTests.java b/tests/framework/base/windowmanager/src/android/server/wm/MultiDisplaySecurityTests.java
index eafe516..dd32bbb 100644
--- a/tests/framework/base/windowmanager/src/android/server/wm/MultiDisplaySecurityTests.java
+++ b/tests/framework/base/windowmanager/src/android/server/wm/MultiDisplaySecurityTests.java
@@ -59,6 +59,7 @@
 import android.server.wm.ActivityManagerState.ActivityStack;
 import android.server.wm.CommandSession.ActivitySession;
 import android.server.wm.TestJournalProvider.TestJournalContainer;
+import android.util.SparseArray;
 import android.view.Display;
 import android.view.View;
 import android.view.ViewGroup;
@@ -92,29 +93,31 @@
      * for activities with same UID.
      */
     @Test
-    public void testLaunchWithoutPermissionOnVirtualDisplayByOwner() {
-        // Create new virtual display.
-        final ActivityDisplay newDisplay = createManagedVirtualDisplaySession().createDisplay();
+    public void testLaunchWithoutPermissionOnVirtualDisplayByOwner() throws Exception {
+        try (final VirtualDisplaySession virtualDisplaySession = new VirtualDisplaySession()) {
+            // Create new virtual display.
+            final ActivityDisplay newDisplay = virtualDisplaySession.createDisplay();
 
-        // Try to launch an activity and check it security exception was triggered.
-        getLaunchActivityBuilder()
-                .setUseBroadcastReceiver(LAUNCH_BROADCAST_RECEIVER, LAUNCH_BROADCAST_ACTION)
-                .setDisplayId(newDisplay.mId)
-                .setTargetActivity(TEST_ACTIVITY)
-                .execute();
+            // Try to launch an activity and check it security exception was triggered.
+            getLaunchActivityBuilder()
+                    .setUseBroadcastReceiver(LAUNCH_BROADCAST_RECEIVER, LAUNCH_BROADCAST_ACTION)
+                    .setDisplayId(newDisplay.mId)
+                    .setTargetActivity(TEST_ACTIVITY)
+                    .execute();
 
-        mAmWmState.waitForValidState(TEST_ACTIVITY);
+            mAmWmState.waitForValidState(TEST_ACTIVITY);
 
-        final int externalFocusedStackId = mAmWmState.getAmState().getFocusedStackId();
-        final ActivityStack focusedStack =
-                mAmWmState.getAmState().getStackById(externalFocusedStackId);
-        assertEquals("Focused stack must be on secondary display", newDisplay.mId,
-                focusedStack.mDisplayId);
+            final int externalFocusedStackId = mAmWmState.getAmState().getFocusedStackId();
+            final ActivityStack focusedStack =
+                    mAmWmState.getAmState().getStackById(externalFocusedStackId);
+            assertEquals("Focused stack must be on secondary display", newDisplay.mId,
+                    focusedStack.mDisplayId);
 
-        mAmWmState.assertFocusedActivity("Focus must be on newly launched app",
-                TEST_ACTIVITY);
-        assertEquals("Activity launched by owner must be on external display",
-                externalFocusedStackId, mAmWmState.getAmState().getFocusedStackId());
+            mAmWmState.assertFocusedActivity("Focus must be on newly launched app",
+                    TEST_ACTIVITY);
+            assertEquals("Activity launched by owner must be on external display",
+                    externalFocusedStackId, mAmWmState.getAmState().getFocusedStackId());
+        }
     }
 
     /**
@@ -122,25 +125,27 @@
      * allowed.
      */
     @Test
-    public void testLaunchWithoutPermissionOnVirtualDisplay() {
-        // Create new virtual display.
-        final ActivityDisplay newDisplay = createManagedVirtualDisplaySession().createDisplay();
+    public void testLaunchWithoutPermissionOnVirtualDisplay() throws Exception {
+        try (final VirtualDisplaySession virtualDisplaySession = new VirtualDisplaySession()) {
+            // Create new virtual display.
+            final ActivityDisplay newDisplay = virtualDisplaySession.createDisplay();
 
-        separateTestJournal();
+            separateTestJournal();
 
-        // Try to launch an activity and check it security exception was triggered.
-        getLaunchActivityBuilder()
-                .setUseBroadcastReceiver(SECOND_LAUNCH_BROADCAST_RECEIVER,
-                        SECOND_LAUNCH_BROADCAST_ACTION)
-                .setDisplayId(newDisplay.mId)
-                .setTargetActivity(TEST_ACTIVITY)
-                .execute();
+            // Try to launch an activity and check it security exception was triggered.
+            getLaunchActivityBuilder()
+                    .setUseBroadcastReceiver(SECOND_LAUNCH_BROADCAST_RECEIVER,
+                            SECOND_LAUNCH_BROADCAST_ACTION)
+                    .setDisplayId(newDisplay.mId)
+                    .setTargetActivity(TEST_ACTIVITY)
+                    .execute();
 
-        assertSecurityExceptionFromActivityLauncher();
+            assertSecurityExceptionFromActivityLauncher();
 
-        mAmWmState.computeState(TEST_ACTIVITY);
-        assertFalse("Restricted activity must not be launched",
-                mAmWmState.getAmState().containsActivity(TEST_ACTIVITY));
+            mAmWmState.computeState(TEST_ACTIVITY);
+            assertFalse("Restricted activity must not be launched",
+                    mAmWmState.getAmState().containsActivity(TEST_ACTIVITY));
+        }
     }
 
     /**
@@ -148,31 +153,27 @@
      * doesn't allow embedding - it should fail with security exception.
      */
     @Test
-    public void testConsequentLaunchActivityFromVirtualDisplayNoEmbedding() {
-        // Create new virtual display.
-        final ActivityDisplay newDisplay = createManagedVirtualDisplaySession().createDisplay();
+    public void testConsequentLaunchActivityFromVirtualDisplayNoEmbedding() throws Exception {
+        try (final VirtualDisplaySession virtualDisplaySession = new VirtualDisplaySession()) {
+            // Create new virtual display.
+            final ActivityDisplay newDisplay = virtualDisplaySession.createDisplay();
 
-        // Launch activity on new secondary display.
-        launchActivityOnDisplay(LAUNCHING_ACTIVITY, newDisplay.mId);
+            // Launch activity on new secondary display.
+            launchActivityOnDisplay(LAUNCHING_ACTIVITY, newDisplay.mId);
 
-        waitAndAssertTopResumedActivity(LAUNCHING_ACTIVITY, newDisplay.mId,
-                "Activity launched on secondary display must be resumed");
+            waitAndAssertTopResumedActivity(LAUNCHING_ACTIVITY, newDisplay.mId,
+                    "Activity launched on secondary display must be resumed");
 
-        separateTestJournal();
+            separateTestJournal();
 
-        // Launch second activity from app on secondary display specifying same display id.
-        getLaunchActivityBuilder()
-                .setTargetActivity(SECOND_NO_EMBEDDING_ACTIVITY)
-                .setDisplayId(newDisplay.mId)
-                .execute();
+            // Launch second activity from app on secondary display specifying same display id.
+            getLaunchActivityBuilder()
+                    .setTargetActivity(SECOND_NO_EMBEDDING_ACTIVITY)
+                    .setDisplayId(newDisplay.mId)
+                    .execute();
 
-        assertSecurityExceptionFromActivityLauncher();
-    }
-
-    private boolean isActivityStartAllowedOnDisplay(int displayId, ComponentName activity) {
-        final Intent intent = new Intent(Intent.ACTION_VIEW).setComponent(activity);
-        return mTargetContext.getSystemService(ActivityManager.class)
-                .isActivityStartAllowedOnDisplay(mTargetContext, displayId, intent);
+            assertSecurityExceptionFromActivityLauncher();
+        }
     }
 
     /**
@@ -181,12 +182,18 @@
      * for simulated display. It is owned by system and is public, so should be accessible.
      */
     @Test
-    public void testCanAccessSystemOwnedDisplay()  {
-        final ActivityDisplay newDisplay = createManagedVirtualDisplaySession()
-                .setSimulateDisplay(true)
-                .createDisplay();
+    public void testCanAccessSystemOwnedDisplay() throws Exception {
+        try (final VirtualDisplaySession virtualDisplaySession = new VirtualDisplaySession()) {
+            final ActivityDisplay newDisplay = virtualDisplaySession.setSimulateDisplay(true)
+                    .createDisplay();
 
-        assertTrue(isActivityStartAllowedOnDisplay(newDisplay.mId, TEST_ACTIVITY));
+            final ActivityManager activityManager =
+                    (ActivityManager) mTargetContext.getSystemService(Context.ACTIVITY_SERVICE);
+            final Intent intent = new Intent(Intent.ACTION_VIEW).setComponent(TEST_ACTIVITY);
+
+            assertTrue(activityManager.isActivityStartAllowedOnDisplay(mTargetContext,
+                    newDisplay.mId, intent));
+        }
     }
 
     /**
@@ -195,15 +202,20 @@
      * for a public virtual display and an activity that doesn't support embedding from shell.
      */
     @Test
-    public void testCanAccessPublicVirtualDisplayWithInternalPermission() {
-        final ActivityDisplay newDisplay = createManagedVirtualDisplaySession()
-                .setPublicDisplay(true)
-                .createDisplay();
+    public void testCanAccessPublicVirtualDisplayWithInternalPermission() throws Exception {
+        try (final VirtualDisplaySession virtualDisplaySession = new VirtualDisplaySession()) {
+            final ActivityDisplay newDisplay = virtualDisplaySession.setPublicDisplay(true)
+                    .createDisplay();
 
-        SystemUtil.runWithShellPermissionIdentity(
-                () -> assertTrue(isActivityStartAllowedOnDisplay(
-                        newDisplay.mId, SECOND_NO_EMBEDDING_ACTIVITY)),
-                "android.permission.INTERNAL_SYSTEM_WINDOW");
+            final ActivityManager activityManager =
+                    (ActivityManager) mTargetContext.getSystemService(Context.ACTIVITY_SERVICE);
+            final Intent intent = new Intent(Intent.ACTION_VIEW)
+                    .setComponent(SECOND_NO_EMBEDDING_ACTIVITY);
+
+            SystemUtil.runWithShellPermissionIdentity(() ->
+                    assertTrue(activityManager.isActivityStartAllowedOnDisplay(mTargetContext,
+                            newDisplay.mId, intent)), "android.permission.INTERNAL_SYSTEM_WINDOW");
+        }
     }
 
     /**
@@ -212,15 +224,20 @@
      * for a private virtual display and an activity that doesn't support embedding from shell.
      */
     @Test
-    public void testCanAccessPrivateVirtualDisplayWithInternalPermission() {
-        final ActivityDisplay newDisplay = createManagedVirtualDisplaySession()
-                .setPublicDisplay(false)
-                .createDisplay();
+    public void testCanAccessPrivateVirtualDisplayWithInternalPermission() throws Exception {
+        try (final VirtualDisplaySession virtualDisplaySession = new VirtualDisplaySession()) {
+            final ActivityDisplay newDisplay = virtualDisplaySession.setPublicDisplay(false)
+                    .createDisplay();
 
-        SystemUtil.runWithShellPermissionIdentity(
-                () -> assertTrue(isActivityStartAllowedOnDisplay(
-                        newDisplay.mId, SECOND_NO_EMBEDDING_ACTIVITY)),
-                "android.permission.INTERNAL_SYSTEM_WINDOW");
+            final ActivityManager activityManager =
+                    (ActivityManager) mTargetContext.getSystemService(Context.ACTIVITY_SERVICE);
+            final Intent intent = new Intent(Intent.ACTION_VIEW)
+                    .setComponent(SECOND_NO_EMBEDDING_ACTIVITY);
+
+            SystemUtil.runWithShellPermissionIdentity(() ->
+                    assertTrue(activityManager.isActivityStartAllowedOnDisplay(mTargetContext,
+                            newDisplay.mId, intent)), "android.permission.INTERNAL_SYSTEM_WINDOW");
+        }
     }
 
     /**
@@ -230,12 +247,18 @@
      * does not have required permission to embed an activity from other app.
      */
     @Test
-    public void testCantAccessPublicVirtualDisplayNoEmbeddingPermission() {
-        final ActivityDisplay newDisplay = createManagedVirtualDisplaySession()
-                .setPublicDisplay(true)
-                .createDisplay();
+    public void testCantAccessPublicVirtualDisplayNoEmbeddingPermission() throws Exception {
+        try (final VirtualDisplaySession virtualDisplaySession = new VirtualDisplaySession()) {
+            final ActivityDisplay newDisplay = virtualDisplaySession.setPublicDisplay(true)
+                    .createDisplay();
 
-        assertFalse(isActivityStartAllowedOnDisplay(newDisplay.mId, SECOND_ACTIVITY));
+            final ActivityManager activityManager =
+                    (ActivityManager) mTargetContext.getSystemService(Context.ACTIVITY_SERVICE);
+            final Intent intent = new Intent(Intent.ACTION_VIEW).setComponent(SECOND_ACTIVITY);
+
+            assertFalse(activityManager.isActivityStartAllowedOnDisplay(mTargetContext,
+                    newDisplay.mId, intent));
+        }
     }
 
     /**
@@ -244,15 +267,20 @@
      * for a public virtual display and an activity that does not support embedding.
      */
     @Test
-    public void testCantAccessPublicVirtualDisplayActivityEmbeddingNotAllowed() {
-        final ActivityDisplay newDisplay = createManagedVirtualDisplaySession()
-                .setPublicDisplay(true)
-                .createDisplay();
+    public void testCantAccessPublicVirtualDisplayActivityEmbeddingNotAllowed() throws Exception {
+        try (final VirtualDisplaySession virtualDisplaySession = new VirtualDisplaySession()) {
+            final ActivityDisplay newDisplay = virtualDisplaySession.setPublicDisplay(true)
+                    .createDisplay();
 
-        SystemUtil.runWithShellPermissionIdentity(
-                () -> assertFalse(isActivityStartAllowedOnDisplay(
-                        newDisplay.mId, SECOND_NO_EMBEDDING_ACTIVITY)),
-                "android.permission.ACTIVITY_EMBEDDING");
+            final ActivityManager activityManager =
+                    (ActivityManager) mTargetContext.getSystemService(Context.ACTIVITY_SERVICE);
+            final Intent intent = new Intent(Intent.ACTION_VIEW)
+                    .setComponent(SECOND_NO_EMBEDDING_ACTIVITY);
+
+            SystemUtil.runWithShellPermissionIdentity(() ->
+                    assertFalse(activityManager.isActivityStartAllowedOnDisplay(mTargetContext,
+                            newDisplay.mId, intent)), "android.permission.ACTIVITY_EMBEDDING");
+        }
     }
 
     /**
@@ -261,15 +289,20 @@
      * for a public virtual display and an activity that supports embedding.
      */
     @Test
-    public void testCanAccessPublicVirtualDisplayActivityEmbeddingAllowed() {
-        final ActivityDisplay newDisplay = createManagedVirtualDisplaySession()
-                .setPublicDisplay(true)
-                .createDisplay();
+    public void testCanAccessPublicVirtualDisplayActivityEmbeddingAllowed() throws Exception {
+        try (final VirtualDisplaySession virtualDisplaySession = new VirtualDisplaySession()) {
+            final ActivityDisplay newDisplay = virtualDisplaySession.setPublicDisplay(true)
+                    .createDisplay();
 
-        SystemUtil.runWithShellPermissionIdentity(
-                () -> assertTrue(isActivityStartAllowedOnDisplay(
-                        newDisplay.mId, SECOND_ACTIVITY)),
-                "android.permission.ACTIVITY_EMBEDDING");
+            final ActivityManager activityManager =
+                    (ActivityManager) mTargetContext.getSystemService(Context.ACTIVITY_SERVICE);
+            final Intent intent = new Intent(Intent.ACTION_VIEW)
+                    .setComponent(SECOND_ACTIVITY);
+
+            SystemUtil.runWithShellPermissionIdentity(() ->
+                    assertTrue(activityManager.isActivityStartAllowedOnDisplay(mTargetContext,
+                            newDisplay.mId, intent)), "android.permission.ACTIVITY_EMBEDDING");
+        }
     }
 
     /**
@@ -278,12 +311,18 @@
      * for a private virtual display.
      */
     @Test
-    public void testCantAccessPrivateVirtualDisplay() {
-        final ActivityDisplay newDisplay = createManagedVirtualDisplaySession()
-                .setPublicDisplay(false)
-                .createDisplay();
+    public void testCantAccessPrivateVirtualDisplay() throws Exception {
+        try (final VirtualDisplaySession virtualDisplaySession = new VirtualDisplaySession()) {
+            final ActivityDisplay newDisplay = virtualDisplaySession.setPublicDisplay(false)
+                    .createDisplay();
 
-        assertFalse(isActivityStartAllowedOnDisplay(newDisplay.mId, SECOND_ACTIVITY));
+            final ActivityManager activityManager =
+                    (ActivityManager) mTargetContext.getSystemService(Context.ACTIVITY_SERVICE);
+            final Intent intent = new Intent(Intent.ACTION_VIEW).setComponent(SECOND_ACTIVITY);
+
+            assertFalse(activityManager.isActivityStartAllowedOnDisplay(mTargetContext,
+                    newDisplay.mId, intent));
+        }
     }
 
     /**
@@ -292,20 +331,21 @@
      * for a private virtual display to check the start of its own activity.
      */
     @Test
-    public void testCanAccessPrivateVirtualDisplayByOwner() {
-        final ActivityDisplay newDisplay = createManagedVirtualDisplaySession()
-                .setPublicDisplay(false)
-                .createDisplay();
+    public void testCanAccessPrivateVirtualDisplayByOwner() throws Exception {
+        try (final VirtualDisplaySession virtualDisplaySession = new VirtualDisplaySession()) {
+            final ActivityDisplay newDisplay = virtualDisplaySession.setPublicDisplay(false)
+                    .createDisplay();
 
-        // Check the embedding call.
-        separateTestJournal();
-        mContext.sendBroadcast(new Intent(ACTION_TEST_ACTIVITY_START)
-                .setPackage(LAUNCH_BROADCAST_RECEIVER.getPackageName())
-                .setFlags(Intent.FLAG_RECEIVER_FOREGROUND)
-                .putExtra(EXTRA_COMPONENT_NAME, TEST_ACTIVITY)
-                .putExtra(EXTRA_TARGET_DISPLAY, newDisplay.mId));
+            // Check the embedding call
+            separateTestJournal();
+            mContext.sendBroadcast(new Intent(ACTION_TEST_ACTIVITY_START)
+                    .setPackage(LAUNCH_BROADCAST_RECEIVER.getPackageName())
+                    .setFlags(Intent.FLAG_RECEIVER_FOREGROUND)
+                    .putExtra(EXTRA_COMPONENT_NAME, TEST_ACTIVITY)
+                    .putExtra(EXTRA_TARGET_DISPLAY, newDisplay.mId));
 
-        assertActivityStartCheckResult(true);
+            assertActivityStartCheckResult(true);
+        }
     }
 
     /**
@@ -315,21 +355,23 @@
      * embedding.
      */
     @Test
-    public void testCanAccessPrivateVirtualDisplayByUidPresentOnDisplayActivityEmbeddingAllowed() {
-        final ActivityDisplay newDisplay = createManagedVirtualDisplaySession()
-                .setPublicDisplay(false)
-                .createDisplay();
-        // Launch a test activity into the target display.
-        launchActivityOnDisplay(EMBEDDING_ACTIVITY, newDisplay.mId);
+    public void testCanAccessPrivateVirtualDisplayByUidPresentOnDisplayActivityEmbeddingAllowed()
+            throws Exception {
+        try (final VirtualDisplaySession virtualDisplaySession = new VirtualDisplaySession()) {
+            final ActivityDisplay newDisplay = virtualDisplaySession.setPublicDisplay(false)
+                    .createDisplay();
+            // Launch a test activity into the target display
+            launchActivityOnDisplay(EMBEDDING_ACTIVITY, newDisplay.mId);
 
-        // Check the embedding call.
-        separateTestJournal();
-        mContext.sendBroadcast(new Intent(ACTION_EMBEDDING_TEST_ACTIVITY_START)
-                .setFlags(Intent.FLAG_RECEIVER_FOREGROUND)
-                .putExtra(EXTRA_EMBEDDING_COMPONENT_NAME, SECOND_ACTIVITY)
-                .putExtra(EXTRA_EMBEDDING_TARGET_DISPLAY, newDisplay.mId));
+            // Check the embedding call
+            separateTestJournal();
+            mContext.sendBroadcast(new Intent(ACTION_EMBEDDING_TEST_ACTIVITY_START)
+                    .setFlags(Intent.FLAG_RECEIVER_FOREGROUND)
+                    .putExtra(EXTRA_EMBEDDING_COMPONENT_NAME, SECOND_ACTIVITY)
+                    .putExtra(EXTRA_EMBEDDING_TARGET_DISPLAY, newDisplay.mId));
 
-        assertActivityStartCheckResult(true);
+            assertActivityStartCheckResult(true);
+        }
     }
 
     /**
@@ -341,20 +383,21 @@
     @Test
     public void testCanAccessPrivateVirtualDisplayByUidPresentOnDisplayActivityEmbeddingNotAllowed()
             throws Exception {
-        final ActivityDisplay newDisplay = createManagedVirtualDisplaySession()
-                .setPublicDisplay(false)
-                .createDisplay();
-        // Launch a test activity into the target display.
-        launchActivityOnDisplay(EMBEDDING_ACTIVITY, newDisplay.mId);
+        try (final VirtualDisplaySession virtualDisplaySession = new VirtualDisplaySession()) {
+            final ActivityDisplay newDisplay = virtualDisplaySession.setPublicDisplay(false)
+                    .createDisplay();
+            // Launch a test activity into the target display
+            launchActivityOnDisplay(EMBEDDING_ACTIVITY, newDisplay.mId);
 
-        // Check the embedding call.
-        separateTestJournal();
-        mContext.sendBroadcast(new Intent(ACTION_EMBEDDING_TEST_ACTIVITY_START)
-                .setFlags(Intent.FLAG_RECEIVER_FOREGROUND)
-                .putExtra(EXTRA_EMBEDDING_COMPONENT_NAME, SECOND_NO_EMBEDDING_ACTIVITY)
-                .putExtra(EXTRA_EMBEDDING_TARGET_DISPLAY, newDisplay.mId));
+            // Check the embedding call
+            separateTestJournal();
+            mContext.sendBroadcast(new Intent(ACTION_EMBEDDING_TEST_ACTIVITY_START)
+                    .setFlags(Intent.FLAG_RECEIVER_FOREGROUND)
+                    .putExtra(EXTRA_EMBEDDING_COMPONENT_NAME, SECOND_NO_EMBEDDING_ACTIVITY)
+                    .putExtra(EXTRA_EMBEDDING_TARGET_DISPLAY, newDisplay.mId));
 
-        assertActivityStartCheckResult(false);
+            assertActivityStartCheckResult(false);
+        }
     }
 
     private void assertActivityStartCheckResult(boolean expected) {
@@ -374,78 +417,81 @@
     }
 
     @Test
-    public void testDisplayHasAccess_UIDCanPresentOnPrivateDisplay() {
-        final VirtualDisplayLauncher virtualDisplayLauncher =
-                mObjectTracker.manage(new VirtualDisplayLauncher());
-        // Create a virtual private display.
-        final ActivityDisplay newDisplay = virtualDisplayLauncher
-                .setPublicDisplay(false)
-                .createDisplay();
-        // Launch an embeddable activity into the private display.
-        // Assert that the UID can present on display.
-        final ActivitySession session1 = virtualDisplayLauncher.launchActivityOnDisplay(
-                DISPLAY_ACCESS_CHECK_EMBEDDING_ACTIVITY, newDisplay);
-        assertEquals("Activity which the UID should accessible on private display",
-                isUidAccesibleOnDisplay(session1), true);
+    public void testDisplayHasAccess_UIDCanPresentOnPrivateDisplay() throws Exception {
+        try (final VirtualDisplayLauncher virtualDisplayLauncher = new VirtualDisplayLauncher()) {
+            // Create a virtual private display.
+            final ActivityDisplay newDisplay = virtualDisplayLauncher
+                    .setPublicDisplay(false)
+                    .createDisplay();
+            // Launch an embeddable activity into the private display.
+            // Assert that the UID can present on display.
+            final ActivitySession session1 = virtualDisplayLauncher.launchActivityOnDisplay(
+                    DISPLAY_ACCESS_CHECK_EMBEDDING_ACTIVITY, newDisplay);
+            assertEquals("Activity which the UID should accessible on private display",
+                    isUidAccesibleOnDisplay(session1), true);
 
-        // Launch another embeddable activity with a different UID, verify that it will be
-        // able to access the display where it was put.
-        // Note that set withShellPermission as true in launchActivityOnDisplay is to
-        // make sure ACTIVITY_EMBEDDING can be granted by shell.
-        final ActivitySession session2 = virtualDisplayLauncher.launchActivityOnDisplay(
-                SECOND_ACTIVITY, newDisplay,
-                bundle -> bundle.putBoolean(EXTRA_DISPLAY_ACCESS_CHECK, true),
-                true /* withShellPermission */, true /* waitForLaunch */);
+            // Launch another embeddable activity with a different UID, verify that it will be
+            // able to access the display where it was put.
+            // Note that set withShellPermission as true in launchActivityOnDisplay is to
+            // make sure ACTIVITY_EMBEDDING can be granted by shell.
+            final ActivitySession session2 = virtualDisplayLauncher.launchActivityOnDisplay(
+                    SECOND_ACTIVITY, newDisplay,
+                    bundle -> bundle.putBoolean(EXTRA_DISPLAY_ACCESS_CHECK, true),
+                    true /* withShellPermission */, true /* waitForLaunch */);
 
-        // Verify SECOND_ACTIVITY's UID has access to this virtual private display.
-        assertEquals("Second activity which the UID should accessible on private display",
-                isUidAccesibleOnDisplay(session2), true);
-    }
-
-    @Test
-    public void testDisplayHasAccess_NoAccessWhenUIDNotPresentOnPrivateDisplay() {
-        final VirtualDisplayLauncher virtualDisplayLauncher =
-                mObjectTracker.manage(new VirtualDisplayLauncher());
-        // Create a virtual private display.
-        final ActivityDisplay newDisplay = virtualDisplayLauncher
-                .setPublicDisplay(false)
-                .createDisplay();
-        // Launch an embeddable activity into the private display.
-        // Assume that the UID can access on display.
-        final ActivitySession session1 = virtualDisplayLauncher.launchActivityOnDisplay(
-                DISPLAY_ACCESS_CHECK_EMBEDDING_ACTIVITY, newDisplay);
-        assertEquals("Activity which the UID should accessible on private display",
-                isUidAccesibleOnDisplay(session1), true);
-
-        // Verify SECOND_NO_EMBEDDING_ACTIVITY's UID can't access this virtual private display
-        // since there is no entity with this UID on this display.
-        // Note that set withShellPermission as false in launchActivityOnDisplay is to
-        // prevent activity can launch when INTERNAL_SYSTEM_WINDOW granted by shell case.
-        separateTestJournal();
-        final ActivitySession session2 = virtualDisplayLauncher.launchActivityOnDisplay(
-                SECOND_NO_EMBEDDING_ACTIVITY, newDisplay, null /* extrasConsumer */,
-                false /* withShellPermission */, false /* waitForLaunch */);
-        assertEquals("Second activity which the UID should not accessible on private display",
-                isUidAccesibleOnDisplay(session2), false);
-    }
-
-    @Test
-    public void testDisplayHasAccess_ExceptionWhenAddViewWithoutPresentOnPrivateDisplay() {
-        // Create a virtual private display.
-        final ActivityDisplay newDisplay = createManagedVirtualDisplaySession()
-                .setPublicDisplay(false)
-                .createDisplay();
-        try {
-            final Display display = mContext.getSystemService(DisplayManager.class).getDisplay(
-                    newDisplay.mId);
-            final Context newDisplayContext = mContext.createDisplayContext(display);
-            newDisplayContext.getSystemService(WindowManager.class).addView(new View(mContext),
-                    new ViewGroup.LayoutParams(WRAP_CONTENT, WRAP_CONTENT));
-        } catch (IllegalArgumentException e) {
-            // Exception happened when createDisplayContext with invalid display.
-            return;
+            // Verify SECOND_ACTIVITY's UID has access to this virtual private display.
+            assertEquals("Second activity which the UID should accessible on private display",
+                    isUidAccesibleOnDisplay(session2), true);
         }
-        fail("UID should not have access to private display without present entities.");
+    }
+
+    @Test
+    public void testDisplayHasAccess_NoAccessWhenUIDNotPresentOnPrivateDisplay() throws Exception {
+        try (final VirtualDisplayLauncher virtualDisplayLauncher = new VirtualDisplayLauncher()) {
+            // Create a virtual private display.
+            final ActivityDisplay newDisplay = virtualDisplayLauncher
+                    .setPublicDisplay(false)
+                    .createDisplay();
+            // Launch an embeddable activity into the private display.
+            // Assume that the UID can access on display.
+            final ActivitySession session1 = virtualDisplayLauncher.launchActivityOnDisplay(
+                    DISPLAY_ACCESS_CHECK_EMBEDDING_ACTIVITY, newDisplay);
+            assertEquals("Activity which the UID should accessible on private display",
+                    isUidAccesibleOnDisplay(session1), true);
+
+            // Verify SECOND_NO_EMBEDDING_ACTIVITY's UID can't access this virtual private display
+            // since there is no entity with this UID on this display.
+            // Note that set withShellPermission as false in launchActivityOnDisplay is to
+            // prevent activity can launch when INTERNAL_SYSTEM_WINDOW granted by shell case.
+            separateTestJournal();
+            final ActivitySession session2 = virtualDisplayLauncher.launchActivityOnDisplay(
+                    SECOND_NO_EMBEDDING_ACTIVITY, newDisplay, null /* extrasConsumer */,
+                    false /* withShellPermission */, false /* waitForLaunch */);
+            assertEquals("Second activity which the UID should not accessible on private display",
+                    isUidAccesibleOnDisplay(session2), false);
+        }
+    }
+
+    @Test
+    public void testDisplayHasAccess_ExceptionWhenAddViewWithoutPresentOnPrivateDisplay()
+            throws Exception {
+        try (final VirtualDisplaySession virtualDisplaySession = new VirtualDisplaySession()) {
+            // Create a virtual private display.
+            final ActivityDisplay newDisplay = virtualDisplaySession
+                    .setPublicDisplay(false)
+                    .createDisplay();
+            try {
+                final Display display = mContext.getSystemService(DisplayManager.class).getDisplay(
+                        newDisplay.mId);
+                final Context newDisplayContext = mContext.createDisplayContext(display);
+                newDisplayContext.getSystemService(WindowManager.class).addView(new View(mContext),
+                        new ViewGroup.LayoutParams(WRAP_CONTENT, WRAP_CONTENT));
+            } catch (IllegalArgumentException e) {
+                // Exception happened when createDisplayContext with invalid display.
+                return;
+            }
+            fail("UID should not have access to private display without present entities.");
+        }
     }
 
     private boolean isUidAccesibleOnDisplay(ActivitySession session) {
@@ -460,91 +506,106 @@
 
     /** Test that shell is allowed to launch on secondary displays. */
     @Test
-    public void testPermissionLaunchFromShell(){
-        // Create new virtual display.
-        final ActivityDisplay newDisplay = createManagedVirtualDisplaySession().createDisplay();
-        mAmWmState.assertVisibility(VIRTUAL_DISPLAY_ACTIVITY, true /* visible */);
-        mAmWmState.assertFocusedActivity("Virtual display activity must be on top",
-                VIRTUAL_DISPLAY_ACTIVITY);
-        final int defaultDisplayFocusedStackId = mAmWmState.getAmState().getFocusedStackId();
-        ActivityStack frontStack = mAmWmState.getAmState().getStackById(
-                defaultDisplayFocusedStackId);
-        assertEquals("Top stack must remain on primary display",
-                DEFAULT_DISPLAY, frontStack.mDisplayId);
+    public void testPermissionLaunchFromShell() throws Exception {
+        try (final VirtualDisplaySession virtualDisplaySession = new VirtualDisplaySession()) {
+            // Create new virtual display.
+            final ActivityDisplay newDisplay = virtualDisplaySession.createDisplay();
+            mAmWmState.assertVisibility(VIRTUAL_DISPLAY_ACTIVITY, true /* visible */);
+            mAmWmState.assertFocusedActivity("Virtual display activity must be on top",
+                    VIRTUAL_DISPLAY_ACTIVITY);
+            final int defaultDisplayFocusedStackId = mAmWmState.getAmState().getFocusedStackId();
+            ActivityStack frontStack = mAmWmState.getAmState().getStackById(
+                    defaultDisplayFocusedStackId);
+            assertEquals("Top stack must remain on primary display",
+                    DEFAULT_DISPLAY, frontStack.mDisplayId);
 
-        // Launch activity on new secondary display.
-        launchActivityOnDisplay(TEST_ACTIVITY, newDisplay.mId);
+            // Launch activity on new secondary display.
+            launchActivityOnDisplay(TEST_ACTIVITY, newDisplay.mId);
 
-        waitAndAssertTopResumedActivity(TEST_ACTIVITY, newDisplay.mId,
-                "Front activity must be on secondary display");
-        assertBothDisplaysHaveResumedActivities(pair(DEFAULT_DISPLAY, VIRTUAL_DISPLAY_ACTIVITY),
-                pair(newDisplay.mId, TEST_ACTIVITY));
+            waitAndAssertTopResumedActivity(TEST_ACTIVITY, newDisplay.mId,
+                    "Front activity must be on secondary display");
+            mAmWmState.assertResumedActivities("Both displays must have resumed activities",
+                    new SparseArray<ComponentName>(){{
+                        put(DEFAULT_DISPLAY, VIRTUAL_DISPLAY_ACTIVITY);
+                        put(newDisplay.mId, TEST_ACTIVITY);
+                    }}
+            );
 
-        // Launch other activity with different uid and check it is launched on dynamic stack on
-        // secondary display.
-        final String startCmd = "am start -n " + getActivityName(SECOND_ACTIVITY)
-                + " --display " + newDisplay.mId;
-        executeShellCommand(startCmd);
+            // Launch other activity with different uid and check it is launched on dynamic stack on
+            // secondary display.
+            final String startCmd = "am start -n " + getActivityName(SECOND_ACTIVITY)
+                    + " --display " + newDisplay.mId;
+            executeShellCommand(startCmd);
 
-        waitAndAssertTopResumedActivity(SECOND_ACTIVITY, newDisplay.mId,
-                "Focus must be on newly launched app");
-        assertBothDisplaysHaveResumedActivities(pair(DEFAULT_DISPLAY, VIRTUAL_DISPLAY_ACTIVITY),
-                pair(newDisplay.mId, SECOND_ACTIVITY));
+            waitAndAssertTopResumedActivity(SECOND_ACTIVITY, newDisplay.mId,
+                    "Focus must be on newly launched app");
+            mAmWmState.assertResumedActivities("Both displays must have resumed activities",
+                    new SparseArray<ComponentName>(){{
+                        put(DEFAULT_DISPLAY, VIRTUAL_DISPLAY_ACTIVITY);
+                        put(newDisplay.mId, SECOND_ACTIVITY);
+                    }}
+            );
+        }
     }
 
     /** Test that launching from app that is on external display is allowed. */
     @Test
-    public void testPermissionLaunchFromAppOnSecondary() {
-        // Create new simulated display.
-        final ActivityDisplay newDisplay = createManagedVirtualDisplaySession()
-                .setSimulateDisplay(true)
-                .createDisplay();
+    public void testPermissionLaunchFromAppOnSecondary() throws Exception {
+        try (final VirtualDisplaySession virtualDisplaySession = new VirtualDisplaySession()) {
+            // Create new simulated display.
+            final ActivityDisplay newDisplay = virtualDisplaySession.setSimulateDisplay(true)
+                    .createDisplay();
 
-        // Launch activity with different uid on secondary display.
-        final String startCmd = "am start -n " + getActivityName(SECOND_ACTIVITY)
-                + " --display " + newDisplay.mId;
-        executeShellCommand(startCmd);
+            // Launch activity with different uid on secondary display.
+            final String startCmd = "am start -n " + getActivityName(SECOND_ACTIVITY)
+                    + " --display " + newDisplay.mId;
+            executeShellCommand(startCmd);
 
-        waitAndAssertTopResumedActivity(SECOND_ACTIVITY, newDisplay.mId,
-                "Top activity must be the newly launched one");
+            waitAndAssertTopResumedActivity(SECOND_ACTIVITY, newDisplay.mId,
+                    "Top activity must be the newly launched one");
 
-        // Launch another activity with third different uid from app on secondary display and
-        // check it is launched on secondary display.
-        getLaunchActivityBuilder()
-                .setUseBroadcastReceiver(SECOND_LAUNCH_BROADCAST_RECEIVER,
-                        SECOND_LAUNCH_BROADCAST_ACTION)
-                .setDisplayId(newDisplay.mId)
-                .setTargetActivity(THIRD_ACTIVITY)
-                .execute();
+            // Launch another activity with third different uid from app on secondary display and
+            // check it is launched on secondary display.
+            getLaunchActivityBuilder()
+                    .setUseBroadcastReceiver(SECOND_LAUNCH_BROADCAST_RECEIVER,
+                            SECOND_LAUNCH_BROADCAST_ACTION)
+                    .setDisplayId(newDisplay.mId)
+                    .setTargetActivity(THIRD_ACTIVITY)
+                    .execute();
 
-        waitAndAssertTopResumedActivity(THIRD_ACTIVITY, newDisplay.mId,
-                "Top activity must be the newly launched one");
+            waitAndAssertTopResumedActivity(THIRD_ACTIVITY, newDisplay.mId,
+                    "Top activity must be the newly launched one");
+        }
     }
 
     /** Tests that an activity can launch an activity from a different UID into its own task. */
     @Test
-    public void testPermissionLaunchMultiUidTask() {
-        final ActivityDisplay newDisplay = createManagedVirtualDisplaySession()
-                .setSimulateDisplay(true)
-                .createDisplay();
+    public void testPermissionLaunchMultiUidTask() throws Exception {
+        try (final VirtualDisplaySession virtualDisplaySession = new VirtualDisplaySession()) {
+            final ActivityDisplay newDisplay = virtualDisplaySession.setSimulateDisplay(true)
+                    .createDisplay();
 
-        launchActivityOnDisplay(LAUNCHING_ACTIVITY, newDisplay.mId);
-        mAmWmState.computeState(LAUNCHING_ACTIVITY);
+            launchActivityOnDisplay(LAUNCHING_ACTIVITY, newDisplay.mId);
+            mAmWmState.computeState(LAUNCHING_ACTIVITY);
 
-        // Check that the first activity is launched onto the secondary display.
-        final int frontStackId = mAmWmState.getAmState().getFrontStackId(newDisplay.mId);
-        ActivityStack frontStack = mAmWmState.getAmState().getStackById(frontStackId);
-        assertEquals("Activity launched on secondary display must be resumed",
-                getActivityName(LAUNCHING_ACTIVITY), frontStack.mResumedActivity);
-        mAmWmState.assertFocusedStack("Top stack must be on secondary display", frontStackId);
+            // Check that the first activity is launched onto the secondary display
+            final int frontStackId = mAmWmState.getAmState().getFrontStackId(newDisplay.mId);
+            ActivityStack frontStack = mAmWmState.getAmState().getStackById(
+                    frontStackId);
+            assertEquals("Activity launched on secondary display must be resumed",
+                    getActivityName(LAUNCHING_ACTIVITY),
+                    frontStack.mResumedActivity);
+            mAmWmState.assertFocusedStack("Top stack must be on secondary display",
+                    frontStackId);
 
-        // Launch an activity from a different UID into the first activity's task.
-        getLaunchActivityBuilder().setTargetActivity(SECOND_ACTIVITY).execute();
+            // Launch an activity from a different UID into the first activity's task
+            getLaunchActivityBuilder().setTargetActivity(SECOND_ACTIVITY).execute();
 
-        waitAndAssertTopResumedActivity(SECOND_ACTIVITY, newDisplay.mId,
-                "Top activity must be the newly launched one");
-        frontStack = mAmWmState.getAmState().getStackById(frontStackId);
-        assertEquals("Secondary display must contain 1 task", 1, frontStack.getTasks().size());
+            waitAndAssertTopResumedActivity(SECOND_ACTIVITY, newDisplay.mId,
+                    "Top activity must be the newly launched one");
+            frontStack = mAmWmState.getAmState().getStackById(frontStackId);
+            assertEquals("Secondary display must contain 1 task", 1, frontStack.getTasks().size());
+        }
     }
 
     /**
@@ -552,36 +613,38 @@
      * doesn't have anything on the display.
      */
     @Test
-    public void testPermissionLaunchFromOwner() {
-        // Create new virtual display.
-        final ActivityDisplay newDisplay = createManagedVirtualDisplaySession().createDisplay();
-        mAmWmState.assertVisibility(VIRTUAL_DISPLAY_ACTIVITY, true /* visible */);
-        mAmWmState.assertFocusedActivity("Virtual display activity must be focused",
-                VIRTUAL_DISPLAY_ACTIVITY);
-        final int defaultDisplayFocusedStackId = mAmWmState.getAmState().getFocusedStackId();
-        ActivityStack frontStack =
-                mAmWmState.getAmState().getStackById(defaultDisplayFocusedStackId);
-        assertEquals("Top stack must remain on primary display",
-                DEFAULT_DISPLAY, frontStack.mDisplayId);
+    public void testPermissionLaunchFromOwner() throws Exception {
+        try (final VirtualDisplaySession virtualDisplaySession = new VirtualDisplaySession()) {
+            // Create new virtual display.
+            final ActivityDisplay newDisplay = virtualDisplaySession.createDisplay();
+            mAmWmState.assertVisibility(VIRTUAL_DISPLAY_ACTIVITY, true /* visible */);
+            mAmWmState.assertFocusedActivity("Virtual display activity must be focused",
+                    VIRTUAL_DISPLAY_ACTIVITY);
+            final int defaultDisplayFocusedStackId = mAmWmState.getAmState().getFocusedStackId();
+            ActivityStack frontStack =
+                    mAmWmState.getAmState().getStackById(defaultDisplayFocusedStackId);
+            assertEquals("Top stack must remain on primary display",
+                    DEFAULT_DISPLAY, frontStack.mDisplayId);
 
-        // Launch other activity with different uid on secondary display.
-        final String startCmd = "am start -n " + getActivityName(SECOND_ACTIVITY)
-                + " --display " + newDisplay.mId;
-        executeShellCommand(startCmd);
+            // Launch other activity with different uid on secondary display.
+            final String startCmd = "am start -n " + getActivityName(SECOND_ACTIVITY)
+                    + " --display " + newDisplay.mId;
+            executeShellCommand(startCmd);
 
-        waitAndAssertTopResumedActivity(SECOND_ACTIVITY, newDisplay.mId,
-                "Top activity must be the newly launched one");
+            waitAndAssertTopResumedActivity(SECOND_ACTIVITY, newDisplay.mId,
+                    "Top activity must be the newly launched one");
 
-        // Check that owner uid can launch its own activity on secondary display.
-        getLaunchActivityBuilder()
-                .setUseBroadcastReceiver(LAUNCH_BROADCAST_RECEIVER, LAUNCH_BROADCAST_ACTION)
-                .setNewTask(true)
-                .setMultipleTask(true)
-                .setDisplayId(newDisplay.mId)
-                .execute();
+            // Check that owner uid can launch its own activity on secondary display.
+            getLaunchActivityBuilder()
+                    .setUseBroadcastReceiver(LAUNCH_BROADCAST_RECEIVER, LAUNCH_BROADCAST_ACTION)
+                    .setNewTask(true)
+                    .setMultipleTask(true)
+                    .setDisplayId(newDisplay.mId)
+                    .execute();
 
-        waitAndAssertTopResumedActivity(TEST_ACTIVITY, newDisplay.mId,
-                "Top activity must be the newly launched one");
+            waitAndAssertTopResumedActivity(TEST_ACTIVITY, newDisplay.mId,
+                    "Top activity must be the newly launched one");
+        }
     }
 
     /**
@@ -589,54 +652,58 @@
      * that external display is not allowed.
      */
     @Test
-    public void testPermissionLaunchFromDifferentApp() {
-        // Create new virtual display.
-        final ActivityDisplay newDisplay = createManagedVirtualDisplaySession().createDisplay();
-        mAmWmState.assertVisibility(VIRTUAL_DISPLAY_ACTIVITY, true /* visible */);
-        mAmWmState.assertFocusedActivity("Virtual display activity must be focused",
-                VIRTUAL_DISPLAY_ACTIVITY);
-        final int defaultDisplayFocusedStackId = mAmWmState.getAmState().getFocusedStackId();
-        ActivityStack frontStack = mAmWmState.getAmState().getStackById(
-                defaultDisplayFocusedStackId);
-        assertEquals("Top stack must remain on primary display",
-                DEFAULT_DISPLAY, frontStack.mDisplayId);
+    public void testPermissionLaunchFromDifferentApp() throws Exception {
+        try (final VirtualDisplaySession virtualDisplaySession = new VirtualDisplaySession()) {
+            // Create new virtual display.
+            final ActivityDisplay newDisplay = virtualDisplaySession.createDisplay();
+            mAmWmState.assertVisibility(VIRTUAL_DISPLAY_ACTIVITY, true /* visible */);
+            mAmWmState.assertFocusedActivity("Virtual display activity must be focused",
+                    VIRTUAL_DISPLAY_ACTIVITY);
+            final int defaultDisplayFocusedStackId = mAmWmState.getAmState().getFocusedStackId();
+            ActivityStack frontStack = mAmWmState.getAmState().getStackById(
+                    defaultDisplayFocusedStackId);
+            assertEquals("Top stack must remain on primary display",
+                    DEFAULT_DISPLAY, frontStack.mDisplayId);
 
-        // Launch activity on new secondary display.
-        launchActivityOnDisplay(TEST_ACTIVITY, newDisplay.mId);
-        waitAndAssertTopResumedActivity(TEST_ACTIVITY, newDisplay.mId,
-                "Top activity must be the newly launched one");
+            // Launch activity on new secondary display.
+            launchActivityOnDisplay(TEST_ACTIVITY, newDisplay.mId);
+            waitAndAssertTopResumedActivity(TEST_ACTIVITY, newDisplay.mId,
+                    "Top activity must be the newly launched one");
 
-        separateTestJournal();
+            separateTestJournal();
 
-        // Launch other activity with different uid and check security exception is triggered.
-        getLaunchActivityBuilder()
-                .setUseBroadcastReceiver(SECOND_LAUNCH_BROADCAST_RECEIVER,
-                        SECOND_LAUNCH_BROADCAST_ACTION)
-                .setDisplayId(newDisplay.mId)
-                .setTargetActivity(THIRD_ACTIVITY)
-                .execute();
+            // Launch other activity with different uid and check security exception is triggered.
+            getLaunchActivityBuilder()
+                    .setUseBroadcastReceiver(SECOND_LAUNCH_BROADCAST_RECEIVER,
+                            SECOND_LAUNCH_BROADCAST_ACTION)
+                    .setDisplayId(newDisplay.mId)
+                    .setTargetActivity(THIRD_ACTIVITY)
+                    .execute();
 
-        assertSecurityExceptionFromActivityLauncher();
+            assertSecurityExceptionFromActivityLauncher();
 
-        mAmWmState.waitForValidState(TEST_ACTIVITY);
-        mAmWmState.assertFocusedActivity("Top activity must be the first one launched",
-                TEST_ACTIVITY);
+            mAmWmState.waitForValidState(TEST_ACTIVITY);
+            mAmWmState.assertFocusedActivity("Top activity must be the first one launched",
+                    TEST_ACTIVITY);
+        }
     }
 
     /**
      * Test that only private virtual display can show content with insecure keyguard.
      */
     @Test
-    public void testFlagShowWithInsecureKeyguardOnPublicVirtualDisplay() {
-        // Try to create new show-with-insecure-keyguard public virtual display.
-        final ActivityDisplay newDisplay = createManagedVirtualDisplaySession()
-                .setPublicDisplay(true)
-                .setCanShowWithInsecureKeyguard(true)
-                .setMustBeCreated(false)
-                .createDisplay();
+    public void testFlagShowWithInsecureKeyguardOnPublicVirtualDisplay() throws Exception {
+        try (final VirtualDisplaySession virtualDisplaySession = new VirtualDisplaySession()) {
+            // Try to create new show-with-insecure-keyguard public virtual display.
+            final ActivityDisplay newDisplay = virtualDisplaySession
+                    .setPublicDisplay(true)
+                    .setCanShowWithInsecureKeyguard(true)
+                    .setMustBeCreated(false)
+                    .createDisplay();
 
-        // Check that the display is not created.
-        assertNull(newDisplay);
+            // Check that the display is not created.
+            assertNull(newDisplay);
+        }
     }
 
     /**
@@ -644,40 +711,41 @@
      */
     @Test
     public void testSettingFlagWithoutInternalSystemPermission() throws Exception {
-        // The reason to use a trusted display is that we can guarantee the security exception
-        // is coming from lacking internal system permission.
-        final ActivityDisplay trustedDisplay = createManagedVirtualDisplaySession()
-                .setSimulateDisplay(true)
-                .createDisplay();
-        final WindowManager wm = mTargetContext.getSystemService(WindowManager.class);
+        try (final VirtualDisplaySession virtualDisplaySession = new VirtualDisplaySession()) {
+            // The reason to use a trusted display is that we can guarantee the security exception
+            // is coming from lacking internal system permission.
+            final ActivityDisplay trustedDisplay = virtualDisplaySession
+                    .setSimulateDisplay(true).createDisplay();
+            final WindowManager wm = mTargetContext.getSystemService(WindowManager.class);
 
-        // Verify setting system decorations flag without internal system permission.
-        try {
-            wm.setShouldShowSystemDecors(trustedDisplay.mId, true);
+            // Verify setting system decorations flag without internal system permission.
+            try {
+                wm.setShouldShowSystemDecors(trustedDisplay.mId, true);
 
-            // Unexpected result, restore flag to avoid affecting other tests.
-            wm.setShouldShowSystemDecors(trustedDisplay.mId, false);
-            TestUtils.waitUntil("Waiting for system decoration flag to be set",
-                    5 /* timeoutSecond */,
-                    () -> !wm.shouldShowSystemDecors(trustedDisplay.mId));
-            fail("Should not allow setting system decoration flag without internal system "
-                    + "permission");
-        } catch (SecurityException e) {
-            // Expected security exception.
-        }
+                // Unexpected result, restore flag to avoid affecting other tests.
+                wm.setShouldShowSystemDecors(trustedDisplay.mId, false);
+                TestUtils.waitUntil("Waiting for system decoration flag to be set",
+                        5 /* timeoutSecond */,
+                        () -> !wm.shouldShowSystemDecors(trustedDisplay.mId));
+                fail("Should not allow setting system decoration flag without internal system "
+                        + "permission");
+            } catch (SecurityException e) {
+                // Expected security exception.
+            }
 
-        // Verify setting show IME flag without internal system permission.
-        try {
-            wm.setShouldShowIme(trustedDisplay.mId, true);
+            // Verify setting show IME flag without internal system permission.
+            try {
+                wm.setShouldShowIme(trustedDisplay.mId, true);
 
-            // Unexpected result, restore flag to avoid affecting other tests.
-            wm.setShouldShowIme(trustedDisplay.mId, false);
-            TestUtils.waitUntil("Waiting for show IME flag to be set",
-                    5 /* timeoutSecond */,
-                    () -> !wm.shouldShowIme(trustedDisplay.mId));
-            fail("Should not allow setting show IME flag without internal system permission");
-        } catch (SecurityException e) {
-            // Expected security exception.
+                // Unexpected result, restore flag to avoid affecting other tests.
+                wm.setShouldShowIme(trustedDisplay.mId, false);
+                TestUtils.waitUntil("Waiting for show IME flag to be set",
+                        5 /* timeoutSecond */,
+                        () -> !wm.shouldShowIme(trustedDisplay.mId));
+                fail("Should not allow setting show IME flag without internal system permission");
+            } catch (SecurityException e) {
+                // Expected security exception.
+            }
         }
     }
 
@@ -685,28 +753,29 @@
      * Test getting system decoration flag and show IME flag without sufficient permissions.
      */
     @Test
-    public void testGettingFlagWithoutInternalSystemPermission() {
-        // The reason to use a trusted display is that we can guarantee the security exception
-        // is coming from lacking internal system permission.
-        final ActivityDisplay trustedDisplay = createManagedVirtualDisplaySession()
-                .setSimulateDisplay(true)
-                .createDisplay();
-        final WindowManager wm = mTargetContext.getSystemService(WindowManager.class);
+    public void testGettingFlagWithoutInternalSystemPermission() throws Exception {
+        try (final VirtualDisplaySession virtualDisplaySession = new VirtualDisplaySession()) {
+            // The reason to use a trusted display is that we can guarantee the security exception
+            // is coming from lacking internal system permission.
+            final ActivityDisplay trustedDisplay = virtualDisplaySession
+                    .setSimulateDisplay(true).createDisplay();
+            final WindowManager wm = mTargetContext.getSystemService(WindowManager.class);
 
-        // Verify getting system decorations flag without internal system permission.
-        try {
-            wm.shouldShowSystemDecors(trustedDisplay.mId);
-            fail("Only allow internal system to get system decoration flag");
-        } catch (SecurityException e) {
-            // Expected security exception.
-        }
+            // Verify getting system decorations flag without internal system permission.
+            try {
+                wm.shouldShowSystemDecors(trustedDisplay.mId);
+                fail("Only allow internal system to get system decoration flag");
+            } catch (SecurityException e) {
+                // Expected security exception.
+            }
 
-        // Verify getting show IME flag without internal system permission.
-        try {
-            wm.shouldShowIme(trustedDisplay.mId);
-            fail("Only allow internal system to get show IME flag");
-        } catch (SecurityException e) {
-            // Expected security exception.
+            // Verify getting show IME flag without internal system permission.
+            try {
+                wm.shouldShowIme(trustedDisplay.mId);
+                fail("Only allow internal system to get show IME flag");
+            } catch (SecurityException e) {
+                // Expected security exception.
+            }
         }
     }
 
@@ -715,43 +784,44 @@
      */
     @Test
     public void testSettingFlagToUntrustedDisplay() throws Exception {
-        final ActivityDisplay untrustedDisplay = createManagedVirtualDisplaySession()
-                .createDisplay();
-        final WindowManager wm = mTargetContext.getSystemService(WindowManager.class);
+        try (final VirtualDisplaySession virtualDisplaySession = new VirtualDisplaySession()) {
+            final ActivityDisplay untrustedDisplay = virtualDisplaySession.createDisplay();
+            final WindowManager wm = mTargetContext.getSystemService(WindowManager.class);
 
-        // Verify setting system decoration flag to an untrusted display.
-        getInstrumentation().getUiAutomation().adoptShellPermissionIdentity();
-        try {
-            wm.setShouldShowSystemDecors(untrustedDisplay.mId, true);
+            // Verify setting system decoration flag to an untrusted display.
+            getInstrumentation().getUiAutomation().adoptShellPermissionIdentity();
+            try {
+                wm.setShouldShowSystemDecors(untrustedDisplay.mId, true);
 
-            // Unexpected result, restore flag to avoid affecting other tests.
-            wm.setShouldShowSystemDecors(untrustedDisplay.mId, false);
-            TestUtils.waitUntil("Waiting for system decoration flag to be set",
-                    5 /* timeoutSecond */,
-                    () -> !wm.shouldShowSystemDecors(untrustedDisplay.mId));
-            fail("Should not allow setting system decoration flag to the untrusted virtual "
-                    + "display");
-        } catch (SecurityException e) {
-            // Expected security exception.
-        } finally {
-            getInstrumentation().getUiAutomation().dropShellPermissionIdentity();
-        }
+                // Unexpected result, restore flag to avoid affecting other tests.
+                wm.setShouldShowSystemDecors(untrustedDisplay.mId, false);
+                TestUtils.waitUntil("Waiting for system decoration flag to be set",
+                        5 /* timeoutSecond */,
+                        () -> !wm.shouldShowSystemDecors(untrustedDisplay.mId));
+                fail("Should not allow setting system decoration flag to the untrusted virtual "
+                        + "display");
+            } catch (SecurityException e) {
+                // Expected security exception.
+            } finally {
+                getInstrumentation().getUiAutomation().dropShellPermissionIdentity();
+            }
 
-        // Verify setting show IME flag to an untrusted display.
-        getInstrumentation().getUiAutomation().adoptShellPermissionIdentity();
-        try {
-            wm.setShouldShowIme(untrustedDisplay.mId, true);
+            // Verify setting show IME flag to an untrusted display.
+            getInstrumentation().getUiAutomation().adoptShellPermissionIdentity();
+            try {
+                wm.setShouldShowIme(untrustedDisplay.mId, true);
 
-            // Unexpected result, restore flag to avoid affecting other tests.
-            wm.setShouldShowIme(untrustedDisplay.mId, false);
-            TestUtils.waitUntil("Waiting for show IME flag to be set",
-                    5 /* timeoutSecond */,
-                    () -> !wm.shouldShowIme(untrustedDisplay.mId));
-            fail("Should not allow setting show IME flag to the untrusted virtual display");
-        } catch (SecurityException e) {
-            // Expected security exception.
-        } finally {
-            getInstrumentation().getUiAutomation().dropShellPermissionIdentity();
+                // Unexpected result, restore flag to avoid affecting other tests.
+                wm.setShouldShowIme(untrustedDisplay.mId, false);
+                TestUtils.waitUntil("Waiting for show IME flag to be set",
+                        5 /* timeoutSecond */,
+                        () -> !wm.shouldShowIme(untrustedDisplay.mId));
+                fail("Should not allow setting show IME flag to the untrusted virtual display");
+            } catch (SecurityException e) {
+                // Expected security exception.
+            } finally {
+                getInstrumentation().getUiAutomation().dropShellPermissionIdentity();
+            }
         }
     }
 
@@ -759,20 +829,21 @@
      * Test getting system decoration flag and show IME flag from the untrusted display.
      */
     @Test
-    public void testGettingFlagFromUntrustedDisplay() {
-        final ActivityDisplay untrustedDisplay = createManagedVirtualDisplaySession()
-                .createDisplay();
-        final WindowManager wm = mTargetContext.getSystemService(WindowManager.class);
+    public void testGettingFlagFromUntrustedDisplay() throws Exception {
+        try (final VirtualDisplaySession virtualDisplaySession = new VirtualDisplaySession()) {
+            final ActivityDisplay untrustedDisplay = virtualDisplaySession.createDisplay();
+            final WindowManager wm = mTargetContext.getSystemService(WindowManager.class);
 
-        // Verify getting system decoration flag from an untrusted display.
-        SystemUtil.runWithShellPermissionIdentity(() -> assertFalse(
-                "Display should not support showing system decorations",
-                wm.shouldShowSystemDecors(untrustedDisplay.mId)));
+            // Verify getting system decoration flag from an untrusted display.
+            SystemUtil.runWithShellPermissionIdentity(() -> assertFalse(
+                    "Display should not support showing system decorations",
+                    wm.shouldShowSystemDecors(untrustedDisplay.mId)));
 
-        // Verify getting show IME flag from an untrusted display.
-        SystemUtil.runWithShellPermissionIdentity(() -> assertFalse(
-                "Display should not support showing IME window",
-                wm.shouldShowIme(untrustedDisplay.mId)));
+            // Verify getting show IME flag from an untrusted display.
+            SystemUtil.runWithShellPermissionIdentity(() -> assertFalse(
+                    "Display should not support showing IME window",
+                    wm.shouldShowIme(untrustedDisplay.mId)));
+        }
     }
 
     /**
@@ -780,51 +851,52 @@
      */
     @Test
     public void testSettingFlagToTrustedDisplay() throws Exception {
-        final ActivityDisplay trustedDisplay = createManagedVirtualDisplaySession()
-                .setSimulateDisplay(true)
-                .createDisplay();
-        final WindowManager wm = mTargetContext.getSystemService(WindowManager.class);
+        try (final VirtualDisplaySession virtualDisplaySession = new VirtualDisplaySession()) {
+            final ActivityDisplay trustedDisplay = virtualDisplaySession
+                    .setSimulateDisplay(true).createDisplay();
+            final WindowManager wm = mTargetContext.getSystemService(WindowManager.class);
 
-        // Verify setting system decoration flag to a trusted display.
-        SystemUtil.runWithShellPermissionIdentity(() -> {
-            // Assume the display should not support system decorations by default.
-            assertFalse(wm.shouldShowSystemDecors(trustedDisplay.mId));
+            // Verify setting system decoration flag to a trusted display.
+            SystemUtil.runWithShellPermissionIdentity(() -> {
+                // Assume the display should not support system decorations by default.
+                assertFalse(wm.shouldShowSystemDecors(trustedDisplay.mId));
 
-            try {
-                wm.setShouldShowSystemDecors(trustedDisplay.mId, true);
-                TestUtils.waitUntil("Waiting for system decoration flag to be set",
-                        5 /* timeoutSecond */,
-                        () -> wm.shouldShowSystemDecors(trustedDisplay.mId));
+                try {
+                    wm.setShouldShowSystemDecors(trustedDisplay.mId, true);
+                    TestUtils.waitUntil("Waiting for system decoration flag to be set",
+                            5 /* timeoutSecond */,
+                            () -> wm.shouldShowSystemDecors(trustedDisplay.mId));
 
-                assertTrue(wm.shouldShowSystemDecors(trustedDisplay.mId));
-            } finally {
-                // Restore flag to avoid affecting other tests.
-                wm.setShouldShowSystemDecors(trustedDisplay.mId, false);
-                TestUtils.waitUntil("Waiting for system decoration flag to be set",
-                        5 /* timeoutSecond */,
-                        () -> !wm.shouldShowSystemDecors(trustedDisplay.mId));
-            }
-        });
+                    assertTrue(wm.shouldShowSystemDecors(trustedDisplay.mId));
+                } finally {
+                    // Restore flag to avoid affecting other tests.
+                    wm.setShouldShowSystemDecors(trustedDisplay.mId, false);
+                    TestUtils.waitUntil("Waiting for system decoration flag to be set",
+                            5 /* timeoutSecond */,
+                            () -> !wm.shouldShowSystemDecors(trustedDisplay.mId));
+                }
+            });
 
-        // Verify setting show IME flag to a trusted display.
-        SystemUtil.runWithShellPermissionIdentity(() -> {
-            // Assume the display should not show IME window by default.
-            assertFalse(wm.shouldShowIme(trustedDisplay.mId));
+            // Verify setting show IME flag to a trusted display.
+            SystemUtil.runWithShellPermissionIdentity(() -> {
+                // Assume the display should not show IME window by default.
+                assertFalse(wm.shouldShowIme(trustedDisplay.mId));
 
-            try {
-                wm.setShouldShowIme(trustedDisplay.mId, true);
-                TestUtils.waitUntil("Waiting for show IME flag to be set",
-                        5 /* timeoutSecond */,
-                        () -> wm.shouldShowIme(trustedDisplay.mId));
+                try {
+                    wm.setShouldShowIme(trustedDisplay.mId, true);
+                    TestUtils.waitUntil("Waiting for show IME flag to be set",
+                            5 /* timeoutSecond */,
+                            () -> wm.shouldShowIme(trustedDisplay.mId));
 
-                assertTrue(wm.shouldShowIme(trustedDisplay.mId));
-            } finally {
-                // Restore flag to avoid affecting other tests.
-                wm.setShouldShowIme(trustedDisplay.mId, false);
-                TestUtils.waitUntil("Waiting for show IME flag to be set",
-                        5 /* timeoutSecond */,
-                        () -> !wm.shouldShowIme(trustedDisplay.mId));
-            }
-        });
+                    assertTrue(wm.shouldShowIme(trustedDisplay.mId));
+                } finally {
+                    // Restore flag to avoid affecting other tests.
+                    wm.setShouldShowIme(trustedDisplay.mId, false);
+                    TestUtils.waitUntil("Waiting for show IME flag to be set",
+                            5 /* timeoutSecond */,
+                            () -> !wm.shouldShowIme(trustedDisplay.mId));
+                }
+            });
+        }
     }
 }
diff --git a/tests/framework/base/windowmanager/src/android/server/wm/MultiDisplaySystemDecorationTests.java b/tests/framework/base/windowmanager/src/android/server/wm/MultiDisplaySystemDecorationTests.java
index ca1c3b1..0a6b9c3 100644
--- a/tests/framework/base/windowmanager/src/android/server/wm/MultiDisplaySystemDecorationTests.java
+++ b/tests/framework/base/windowmanager/src/android/server/wm/MultiDisplaySystemDecorationTests.java
@@ -25,10 +25,11 @@
 import static android.server.wm.app.Components.TestLiveWallpaperKeys.COMPONENT;
 import static android.server.wm.app.Components.TestLiveWallpaperKeys.ENGINE_DISPLAY_ID;
 import static android.server.wm.BarTestUtils.assumeHasBars;
-import static android.server.wm.MockImeHelper.createManagedMockImeSession;
 import static android.view.Display.DEFAULT_DISPLAY;
 import static android.view.WindowManager.LayoutParams.TYPE_WALLPAPER;
 
+import static androidx.test.InstrumentationRegistry.getInstrumentation;
+
 import static com.android.cts.mockime.ImeEventStreamTestUtils.editorMatcher;
 import static com.android.cts.mockime.ImeEventStreamTestUtils.expectCommand;
 import static com.android.cts.mockime.ImeEventStreamTestUtils.expectEvent;
@@ -70,6 +71,7 @@
 import com.android.cts.mockime.ImeCommand;
 import com.android.cts.mockime.ImeEvent;
 import com.android.cts.mockime.ImeEventStream;
+import com.android.cts.mockime.ImeSettings;
 import com.android.cts.mockime.MockImeSession;
 
 import org.junit.Before;
@@ -108,48 +110,48 @@
      */
     @Test
     public void testWallpaperGetDisplayContext() throws Exception {
-        final ChangeWallpaperSession wallpaperSession = createManagedChangeWallpaperSession();
-        final VirtualDisplaySession virtualDisplaySession = createManagedVirtualDisplaySession();
+        try (final ChangeWallpaperSession wallpaperSession = new ChangeWallpaperSession();
+             final VirtualDisplaySession virtualDisplaySession = new VirtualDisplaySession()) {
 
-        TestJournalContainer.start();
+            TestJournalContainer.start();
 
-        final ActivityDisplay newDisplay = virtualDisplaySession
-                .setSimulateDisplay(true).setShowSystemDecorations(true).createDisplay();
+            final ActivityDisplay newDisplay = virtualDisplaySession
+                    .setSimulateDisplay(true).setShowSystemDecorations(true).createDisplay();
 
-        wallpaperSession.setWallpaperComponent(TEST_LIVE_WALLPAPER_SERVICE);
-        final String TARGET_ENGINE_DISPLAY_ID = ENGINE_DISPLAY_ID + newDisplay.mId;
-        final TestJournalProvider.TestJournal journal = TestJournalContainer.get(COMPONENT);
-        TestUtils.waitUntil("Waiting for wallpaper engine bounded", 5 /* timeoutSecond */,
-                () -> journal.extras.getBoolean(TARGET_ENGINE_DISPLAY_ID));
+            wallpaperSession.setWallpaperComponent(TEST_LIVE_WALLPAPER_SERVICE);
+            final String TARGET_ENGINE_DISPLAY_ID = ENGINE_DISPLAY_ID + newDisplay.mId;
+            final TestJournalProvider.TestJournal journal = TestJournalContainer.get(COMPONENT);
+            TestUtils.waitUntil("Waiting for wallpaper engine bounded", 5 /* timeoutSecond */,
+                    () -> journal.extras.getBoolean(TARGET_ENGINE_DISPLAY_ID));
+        }
     }
 
     /**
      * Tests that wallpaper shows on secondary displays.
      */
     @Test
-    public void testWallpaperShowOnSecondaryDisplays()  {
-        final ChangeWallpaperSession wallpaperSession = createManagedChangeWallpaperSession();
+    public void testWallpaperShowOnSecondaryDisplays() throws Exception {
+        try (final ChangeWallpaperSession wallpaperSession = new ChangeWallpaperSession();
+             final ExternalDisplaySession externalDisplaySession = new ExternalDisplaySession();
+             final VirtualDisplaySession virtualDisplaySession = new VirtualDisplaySession()) {
 
-        final ActivityDisplay untrustedDisplay = createManagedExternalDisplaySession()
-                .setPublicDisplay(true).setShowSystemDecorations(true).createVirtualDisplay();
+            final ActivityDisplay untrustedDisplay = externalDisplaySession
+                    .setPublicDisplay(true).setShowSystemDecorations(true).createVirtualDisplay();
 
-        final ActivityDisplay decoredSystemDisplay = createManagedVirtualDisplaySession()
-                .setSimulateDisplay(true).setShowSystemDecorations(true).createDisplay();
+            final ActivityDisplay decoredSystemDisplay = virtualDisplaySession
+                    .setSimulateDisplay(true).setShowSystemDecorations(true).createDisplay();
 
-        final Bitmap tmpWallpaper = wallpaperSession.getTestBitmap();
-        wallpaperSession.setImageWallpaper(tmpWallpaper);
+            final Bitmap tmpWallpaper = wallpaperSession.getTestBitmap();
+            wallpaperSession.setImageWallpaper(tmpWallpaper);
 
-        assertTrue("Wallpaper must be displayed on system owned display with system decor flag",
-                mAmWmState.waitForWithWmState(
-                        state -> isWallpaperOnDisplay(state, decoredSystemDisplay.mId),
-                        "wallpaper window to show"));
+            assertTrue("Wallpaper must be displayed on system owned display with system decor flag",
+                    mAmWmState.waitForWithWmState(
+                            state -> isWallpaperOnDisplay(state, decoredSystemDisplay.mId),
+                            "wallpaper window to show"));
 
-        assertFalse("Wallpaper must not be displayed on the untrusted display",
-                isWallpaperOnDisplay(mAmWmState.getWmState(), untrustedDisplay.mId));
-    }
-
-    private ChangeWallpaperSession createManagedChangeWallpaperSession() {
-        return mObjectTracker.manage(new ChangeWallpaperSession());
+            assertFalse("Wallpaper must not be displayed on the untrusted display",
+                    isWallpaperOnDisplay(mAmWmState.getWmState(), untrustedDisplay.mId));
+        }
     }
 
     private class ChangeWallpaperSession implements AutoCloseable {
@@ -169,18 +171,18 @@
             return mTestBitmap;
         }
 
-        public void setImageWallpaper(Bitmap bitmap) {
+        public void setImageWallpaper(Bitmap bitmap) throws Exception {
             SystemUtil.runWithShellPermissionIdentity(() ->
                     mWallpaperManager.setBitmap(bitmap));
         }
 
-        public void setWallpaperComponent(ComponentName componentName) {
+        public void setWallpaperComponent(ComponentName componentName) throws Exception {
             SystemUtil.runWithShellPermissionIdentity(() ->
                     mWallpaperManager.setWallpaperComponent(componentName));
         }
 
         @Override
-        public void close() {
+        public void close() throws Exception {
             SystemUtil.runWithShellPermissionIdentity(mWallpaperManager::clearWallpaper);
             if (mTestBitmap != null) {
                 mTestBitmap.recycle();
@@ -199,28 +201,32 @@
      * Test that navigation bar should show on display with system decoration.
      */
     @Test
-    public void testNavBarShowingOnDisplayWithDecor() {
+    public void testNavBarShowingOnDisplayWithDecor() throws Exception {
         assumeHasBars();
-        final ActivityDisplay newDisplay = createManagedExternalDisplaySession()
-                .setPublicDisplay(true).setShowSystemDecorations(true).createVirtualDisplay();
+        try (final ExternalDisplaySession externalDisplaySession = new ExternalDisplaySession()) {
+            final ActivityDisplay newDisplay = externalDisplaySession
+                    .setPublicDisplay(true).setShowSystemDecorations(true).createVirtualDisplay();
 
-        mAmWmState.waitAndAssertNavBarShownOnDisplay(newDisplay.mId);
+            mAmWmState.waitAndAssertNavBarShownOnDisplay(newDisplay.mId);
+        }
     }
 
     /**
      * Test that navigation bar should not show on display without system decoration.
      */
     @Test
-    public void testNavBarNotShowingOnDisplayWithoutDecor() {
+    public void testNavBarNotShowingOnDisplayWithoutDecor() throws Exception {
         assumeHasBars();
-        // Wait for system decoration showing and record current nav states.
-        mAmWmState.waitForHomeActivityVisible();
-        final List<WindowState> expected = mAmWmState.getWmState().getAllNavigationBarStates();
+        try (final ExternalDisplaySession externalDisplaySession = new ExternalDisplaySession()) {
+            // Wait for system decoration showing and record current nav states.
+            mAmWmState.waitForHomeActivityVisible();
+            final List<WindowState> expected = mAmWmState.getWmState().getAllNavigationBarStates();
 
-        createManagedExternalDisplaySession().setPublicDisplay(true)
-                .setShowSystemDecorations(false).createVirtualDisplay();
+            externalDisplaySession.setPublicDisplay(true)
+                    .setShowSystemDecorations(false).createVirtualDisplay();
 
-        waitAndAssertNavBarStatesAreTheSame(expected);
+            waitAndAssertNavBarStatesAreTheSame(expected);
+        }
     }
 
     /**
@@ -228,19 +234,21 @@
      * supports system decoration.
      */
     @Test
-    public void testNavBarNotShowingOnPrivateDisplay() {
+    public void testNavBarNotShowingOnPrivateDisplay() throws Exception {
         assumeHasBars();
-        // Wait for system decoration showing and record current nav states.
-        mAmWmState.waitForHomeActivityVisible();
-        final List<WindowState> expected = mAmWmState.getWmState().getAllNavigationBarStates();
+        try (final ExternalDisplaySession externalDisplaySession = new ExternalDisplaySession()) {
+            // Wait for system decoration showing and record current nav states.
+            mAmWmState.waitForHomeActivityVisible();
+            final List<WindowState> expected = mAmWmState.getWmState().getAllNavigationBarStates();
 
-        createManagedExternalDisplaySession().setPublicDisplay(false)
-                .setShowSystemDecorations(true).createVirtualDisplay();
+            externalDisplaySession.setPublicDisplay(false)
+                    .setShowSystemDecorations(true).createVirtualDisplay();
 
-        waitAndAssertNavBarStatesAreTheSame(expected);
+            waitAndAssertNavBarStatesAreTheSame(expected);
+        }
     }
 
-    private void waitAndAssertNavBarStatesAreTheSame(List<WindowState> expected) {
+    private void waitAndAssertNavBarStatesAreTheSame(List<WindowState> expected) throws Exception {
         // This is used to verify that we have nav bars shown on the same displays
         // as before the test.
         //
@@ -273,16 +281,18 @@
      * Tests launching a home activity on virtual display without system decoration support.
      */
     @Test
-    public void testLaunchHomeActivityOnSecondaryDisplayWithoutDecorations() {
-        createManagedHomeActivitySession(SECONDARY_HOME_ACTIVITY);
+    public void testLaunchHomeActivityOnSecondaryDisplayWithoutDecorations() throws Exception {
+        try (final HomeActivitySession homeSession =
+                     new HomeActivitySession(SECONDARY_HOME_ACTIVITY);
+             final ExternalDisplaySession externalDisplaySession = new ExternalDisplaySession()) {
+            // Create new virtual display without system decoration support.
+            final ActivityDisplay newDisplay = externalDisplaySession.createVirtualDisplay();
 
-        // Create new virtual display without system decoration support.
-        final ActivityDisplay newDisplay = createManagedExternalDisplaySession()
-                .createVirtualDisplay();
-
-        // Secondary home activity can't be launched on the display without system decoration
-        // support.
-        assertEquals("No stacks on newly launched virtual display", 0, newDisplay.mStacks.size());
+            // Secondary home activity can't be launched on the display without system decoration
+            // support.
+            assertEquals("No stacks on newly launched virtual display", 0,
+                    newDisplay.mStacks.size());
+        }
     }
 
     /**
@@ -290,19 +300,20 @@
      * support.
      */
     @Test
-    public void testLaunchSingleHomeActivityOnDisplayWithDecorations() {
-        createManagedHomeActivitySession(SINGLE_HOME_ACTIVITY);
+    public void testLaunchSingleHomeActivityOnDisplayWithDecorations() throws Exception {
+        try (final HomeActivitySession homeSession = new HomeActivitySession(SINGLE_HOME_ACTIVITY);
+             final ExternalDisplaySession externalDisplaySession = new ExternalDisplaySession()) {
+            // Create new virtual display with system decoration support.
+            final ActivityDisplay newDisplay
+                    = externalDisplaySession.setShowSystemDecorations(true).createVirtualDisplay();
 
-        // Create new virtual display with system decoration support.
-        final ActivityDisplay newDisplay = createManagedExternalDisplaySession()
-                .setShowSystemDecorations(true).createVirtualDisplay();
-
-        // If default home doesn't support multi-instance, default secondary home activity
-        // should be automatically launched on the new display.
-        waitAndAssertTopResumedActivity(getDefaultSecondaryHomeComponent(), newDisplay.mId,
-                "Activity launched on secondary display must be focused and on top");
-        assertEquals("Top activity must be home type", ACTIVITY_TYPE_HOME,
-                mAmWmState.getAmState().getFrontStackActivityType(newDisplay.mId));
+            // If default home doesn't support multi-instance, default secondary home activity
+            // should be automatically launched on the new display.
+            waitAndAssertTopResumedActivity(getDefaultSecondaryHomeComponent(), newDisplay.mId,
+                    "Activity launched on secondary display must be focused and on top");
+            assertEquals("Top activity must be home type", ACTIVITY_TYPE_HOME,
+                    mAmWmState.getAmState().getFrontStackActivityType(newDisplay.mId));
+        }
     }
 
     /**
@@ -310,19 +321,21 @@
      * system decoration support.
      */
     @Test
-    public void testLaunchSingleSecondaryHomeActivityOnDisplayWithDecorations() {
-        createManagedHomeActivitySession(SINGLE_SECONDARY_HOME_ACTIVITY);
+    public void testLaunchSingleSecondaryHomeActivityOnDisplayWithDecorations() throws Exception {
+        try (final HomeActivitySession homeSession =
+                     new HomeActivitySession(SINGLE_SECONDARY_HOME_ACTIVITY);
+             final ExternalDisplaySession externalDisplaySession = new ExternalDisplaySession()) {
+            // Create new virtual display with system decoration support.
+            final ActivityDisplay newDisplay
+                    = externalDisplaySession.setShowSystemDecorations(true).createVirtualDisplay();
 
-        // Create new virtual display with system decoration support.
-        final ActivityDisplay newDisplay = createManagedExternalDisplaySession()
-                .setShowSystemDecorations(true).createVirtualDisplay();
-
-        // If provided secondary home doesn't support multi-instance, default secondary home
-        // activity should be automatically launched on the new display.
-        waitAndAssertTopResumedActivity(getDefaultSecondaryHomeComponent(), newDisplay.mId,
-                "Activity launched on secondary display must be focused and on top");
-        assertEquals("Top activity must be home type", ACTIVITY_TYPE_HOME,
-                mAmWmState.getAmState().getFrontStackActivityType(newDisplay.mId));
+            // If provided secondary home doesn't support multi-instance, default secondary home
+            // activity should be automatically launched on the new display.
+            waitAndAssertTopResumedActivity(getDefaultSecondaryHomeComponent(), newDisplay.mId,
+                    "Activity launched on secondary display must be focused and on top");
+            assertEquals("Top activity must be home type", ACTIVITY_TYPE_HOME,
+                    mAmWmState.getAmState().getFrontStackActivityType(newDisplay.mId));
+        }
     }
 
     /**
@@ -330,20 +343,20 @@
      * support.
      */
     @Test
-    public void testLaunchHomeActivityOnDisplayWithDecorations() {
-        createManagedHomeActivitySession(HOME_ACTIVITY);
-        final ExternalDisplaySession externalDisplaySession = createManagedExternalDisplaySession();
+    public void testLaunchHomeActivityOnDisplayWithDecorations() throws Exception {
+        try (final HomeActivitySession homeSession = new HomeActivitySession(HOME_ACTIVITY);
+             final ExternalDisplaySession externalDisplaySession = new ExternalDisplaySession()) {
+            // Create new virtual display with system decoration support.
+            final ActivityDisplay newDisplay
+                    = externalDisplaySession.setShowSystemDecorations(true).createVirtualDisplay();
 
-        // Create new virtual display with system decoration support.
-        final ActivityDisplay newDisplay =
-                externalDisplaySession.setShowSystemDecorations(true).createVirtualDisplay();
-
-        // If default home doesn't have SECONDARY_HOME category, default secondary home
-        // activity should be automatically launched on the new display.
-        waitAndAssertTopResumedActivity(getDefaultSecondaryHomeComponent(), newDisplay.mId,
-                "Activity launched on secondary display must be focused and on top");
-        assertEquals("Top activity must be home type", ACTIVITY_TYPE_HOME,
-                mAmWmState.getAmState().getFrontStackActivityType(newDisplay.mId));
+            // If default home doesn't have SECONDARY_HOME category, default secondary home
+            // activity should be automatically launched on the new display.
+            waitAndAssertTopResumedActivity(getDefaultSecondaryHomeComponent(), newDisplay.mId,
+                    "Activity launched on secondary display must be focused and on top");
+            assertEquals("Top activity must be home type", ACTIVITY_TYPE_HOME,
+                    mAmWmState.getAmState().getFrontStackActivityType(newDisplay.mId));
+        }
     }
 
     /**
@@ -351,102 +364,111 @@
      * system decoration support.
      */
     @Test
-    public void testLaunchSecondaryHomeActivityOnDisplayWithDecorations() {
-        createManagedHomeActivitySession(SECONDARY_HOME_ACTIVITY);
-        final ExternalDisplaySession externalDisplaySession = createManagedExternalDisplaySession();
+    public void testLaunchSecondaryHomeActivityOnDisplayWithDecorations() throws Exception {
+        try (final HomeActivitySession homeSession =
+                     new HomeActivitySession(SECONDARY_HOME_ACTIVITY);
+             final ExternalDisplaySession externalDisplaySession = new ExternalDisplaySession()) {
+            // Create new virtual display with system decoration support.
+            final ActivityDisplay newDisplay
+                    = externalDisplaySession.setShowSystemDecorations(true).createVirtualDisplay();
 
-        // Create new virtual display with system decoration support.
-        final ActivityDisplay newDisplay =
-                externalDisplaySession.setShowSystemDecorations(true).createVirtualDisplay();
-
-        // Provided secondary home activity should be automatically launched on the new display.
-        waitAndAssertTopResumedActivity(SECONDARY_HOME_ACTIVITY, newDisplay.mId,
-                "Activity launched on secondary display must be focused and on top");
-        assertEquals("Top activity must be home type", ACTIVITY_TYPE_HOME,
-                mAmWmState.getAmState().getFrontStackActivityType(newDisplay.mId));
+            // Provided secondary home activity should be automatically launched on the new
+            // display.
+            waitAndAssertTopResumedActivity(SECONDARY_HOME_ACTIVITY, newDisplay.mId,
+                    "Activity launched on secondary display must be focused and on top");
+            assertEquals("Top activity must be home type", ACTIVITY_TYPE_HOME,
+                    mAmWmState.getAmState().getFrontStackActivityType(newDisplay.mId));
+        }
     }
 
     // IME related tests
     @Test
     public void testImeWindowCanSwitchToDifferentDisplays() throws Exception {
-        final MockImeSession mockImeSession = createManagedMockImeSession(this);
-        final TestActivitySession<ImeTestActivity> imeTestActivitySession =
-                createManagedTestActivitySession();
-        final TestActivitySession<ImeTestActivity2> imeTestActivitySession2 =
-                createManagedTestActivitySession();
+        try (final TestActivitySession<ImeTestActivity> imeTestActivitySession = new
+                TestActivitySession<>();
+             final TestActivitySession<ImeTestActivity2> imeTestActivitySession2 = new
+                     TestActivitySession<>();
+             final VirtualDisplaySession virtualDisplaySession = new VirtualDisplaySession();
 
-        // Create a virtual display and launch an activity on it.
-        final ActivityDisplay newDisplay = createManagedVirtualDisplaySession()
-                .setShowSystemDecorations(true)
-                .setSimulateDisplay(true)
-                .createDisplay();
-        imeTestActivitySession.launchTestActivityOnDisplaySync(ImeTestActivity.class,
-                newDisplay.mId);
+             // Leverage MockImeSession to ensure at least an IME exists as default.
+             final MockImeSession mockImeSession = MockImeSession.create(
+                     mContext, getInstrumentation().getUiAutomation(), new ImeSettings.Builder())) {
 
-        // Make the activity to show soft input.
-        final ImeEventStream stream = mockImeSession.openEventStream();
-        imeTestActivitySession.runOnMainSyncAndWait(
-                imeTestActivitySession.getActivity()::showSoftInput);
-        waitOrderedImeEventsThenAssertImeShown(stream, newDisplay.mId,
-                editorMatcher("onStartInput",
-                        imeTestActivitySession.getActivity().mEditText.getPrivateImeOptions()),
-                event -> "showSoftInput".equals(event.getEventName()));
+            // Create a virtual display and launch an activity on it.
+            final ActivityDisplay newDisplay = virtualDisplaySession.setShowSystemDecorations(true)
+                    .setSimulateDisplay(true).createDisplay();
+            imeTestActivitySession.launchTestActivityOnDisplaySync(ImeTestActivity.class,
+                    newDisplay.mId);
 
-        // Assert the configuration of the IME window is the same as the configuration of the
-        // virtual display.
-        assertImeWindowAndDisplayConfiguration(mAmWmState.getImeWindowState(), newDisplay);
+            // Make the activity to show soft input.
+            final ImeEventStream stream = mockImeSession.openEventStream();
+            imeTestActivitySession.runOnMainSyncAndWait(
+                    imeTestActivitySession.getActivity()::showSoftInput);
+            waitOrderedImeEventsThenAssertImeShown(stream, newDisplay.mId,
+                    editorMatcher("onStartInput",
+                            imeTestActivitySession.getActivity().mEditText.getPrivateImeOptions()),
+                    event -> "showSoftInput".equals(event.getEventName()));
 
-        // Launch another activity on the default display.
-        imeTestActivitySession2.launchTestActivityOnDisplaySync(
-                ImeTestActivity2.class, DEFAULT_DISPLAY);
+            // Assert the configuration of the IME window is the same as the configuration of the
+            // virtual display.
+            assertImeWindowAndDisplayConfiguration(mAmWmState.getImeWindowState(), newDisplay);
 
-        // Make the activity to show soft input.
-        imeTestActivitySession2.runOnMainSyncAndWait(
-                imeTestActivitySession2.getActivity()::showSoftInput);
-        waitOrderedImeEventsThenAssertImeShown(stream, DEFAULT_DISPLAY,
-                editorMatcher("onStartInput",
-                        imeTestActivitySession2.getActivity().mEditText.getPrivateImeOptions()),
-                event -> "showSoftInput".equals(event.getEventName()));
+            // Launch another activity on the default display.
+            imeTestActivitySession2.launchTestActivityOnDisplaySync(
+                    ImeTestActivity2.class, DEFAULT_DISPLAY);
 
-        // Assert the configuration of the IME window is the same as the configuration of the
-        // default display.
-        assertImeWindowAndDisplayConfiguration(mAmWmState.getImeWindowState(),
-                mAmWmState.getAmState().getDisplay(DEFAULT_DISPLAY));
+            // Make the activity to show soft input.
+            imeTestActivitySession2.runOnMainSyncAndWait(
+                    imeTestActivitySession2.getActivity()::showSoftInput);
+            waitOrderedImeEventsThenAssertImeShown(stream, DEFAULT_DISPLAY,
+                    editorMatcher("onStartInput",
+                            imeTestActivitySession2.getActivity().mEditText.getPrivateImeOptions()),
+                    event -> "showSoftInput".equals(event.getEventName()));
+
+            // Assert the configuration of the IME window is the same as the configuration of the
+            // default display.
+            assertImeWindowAndDisplayConfiguration(mAmWmState.getImeWindowState(),
+                    mAmWmState.getAmState().getDisplay(DEFAULT_DISPLAY));
+        }
     }
 
     @Test
     public void testImeApiForBug118341760() throws Exception {
         final long TIMEOUT_START_INPUT = TimeUnit.SECONDS.toMillis(5);
 
-        final MockImeSession mockImeSession = createManagedMockImeSession(this);
-        final TestActivitySession<ImeTestActivityWithBrokenContextWrapper> imeTestActivitySession =
-                createManagedTestActivitySession();
-        // Create a virtual display and launch an activity on it.
-        final ActivityDisplay newDisplay = createManagedVirtualDisplaySession()
-                .setShowSystemDecorations(true)
-                .setSimulateDisplay(true)
-                .createDisplay();
-        imeTestActivitySession.launchTestActivityOnDisplaySync(
-                ImeTestActivityWithBrokenContextWrapper.class, newDisplay.mId);
+        try (final VirtualDisplaySession virtualDisplaySession = new VirtualDisplaySession();
+             final TestActivitySession<ImeTestActivityWithBrokenContextWrapper>
+                     imeTestActivitySession = new TestActivitySession<>();
 
-        final ImeTestActivityWithBrokenContextWrapper activity =
-                imeTestActivitySession.getActivity();
-        final ImeEventStream stream = mockImeSession.openEventStream();
-        final String privateImeOption = activity.getEditText().getPrivateImeOptions();
-        expectEvent(stream, event -> {
-            if (!TextUtils.equals("onStartInput", event.getEventName())) {
-                return false;
-            }
-            final EditorInfo editorInfo = event.getArguments().getParcelable("editorInfo");
-            return TextUtils.equals(editorInfo.packageName, mContext.getPackageName())
-                    && TextUtils.equals(editorInfo.privateImeOptions, privateImeOption);
-        }, TIMEOUT_START_INPUT);
+             // Leverage MockImeSession to ensure at least an IME exists as default.
+             final MockImeSession mockImeSession = MockImeSession.create(
+                     mContext, getInstrumentation().getUiAutomation(), new ImeSettings.Builder())) {
 
-        imeTestActivitySession.runOnMainSyncAndWait(() -> {
-            final InputMethodManager imm = activity.getSystemService(InputMethodManager.class);
-            assertTrue("InputMethodManager.isActive() should work",
-                    imm.isActive(activity.getEditText()));
-        });
+            // Create a virtual display and launch an activity on it.
+            final ActivityDisplay newDisplay = virtualDisplaySession.setShowSystemDecorations(true)
+                    .setSimulateDisplay(true).createDisplay();
+            imeTestActivitySession.launchTestActivityOnDisplaySync(
+                    ImeTestActivityWithBrokenContextWrapper.class, newDisplay.mId);
+
+            final ImeTestActivityWithBrokenContextWrapper activity =
+                    imeTestActivitySession.getActivity();
+            final ImeEventStream stream = mockImeSession.openEventStream();
+            final String privateImeOption = activity.getEditText().getPrivateImeOptions();
+            expectEvent(stream, event -> {
+                if (!TextUtils.equals("onStartInput", event.getEventName())) {
+                    return false;
+                }
+                final EditorInfo editorInfo = event.getArguments().getParcelable("editorInfo");
+                return TextUtils.equals(editorInfo.packageName, mContext.getPackageName())
+                        && TextUtils.equals(editorInfo.privateImeOptions, privateImeOption);
+            }, TIMEOUT_START_INPUT);
+
+            imeTestActivitySession.runOnMainSyncAndWait(() -> {
+                final InputMethodManager imm = activity.getSystemService(InputMethodManager.class);
+                assertTrue("InputMethodManager.isActive() should work",
+                        imm.isActive(activity.getEditText()));
+            });
+        }
     }
 
     @Test
@@ -455,53 +477,55 @@
         // the Activity in the different display.
         assumeFalse(perDisplayFocusEnabled());
 
-        final MockImeSession mockImeSession = createManagedMockImeSession(this);
-        final TestActivitySession<ImeTestActivity> imeTestActivitySession =
-                createManagedTestActivitySession();
-        final TestActivitySession<ImeTestActivity2> imeTestActivitySession2 =
-                createManagedTestActivitySession();
+        try (final VirtualDisplaySession virtualDisplaySession = new VirtualDisplaySession();
+             final TestActivitySession<ImeTestActivity> imeTestActivitySession = new
+                     TestActivitySession<>();
+             final TestActivitySession<ImeTestActivity2> imeTestActivitySession2 = new
+                     TestActivitySession<>();
+             // Leverage MockImeSession to ensure at least an IME exists as default.
+             final MockImeSession mockImeSession1 = MockImeSession.create(
+                     mContext, getInstrumentation().getUiAutomation(), new ImeSettings.Builder())) {
 
-        // Create a virtual display and launch an activity on virtual & default display.
-        final ActivityDisplay newDisplay = createManagedVirtualDisplaySession()
-                .setShowSystemDecorations(true)
-                .setSimulateDisplay(true)
-                .createDisplay();
-        imeTestActivitySession.launchTestActivityOnDisplaySync(ImeTestActivity.class,
-                DEFAULT_DISPLAY);
-        imeTestActivitySession2.launchTestActivityOnDisplaySync(ImeTestActivity2.class,
-                newDisplay.mId);
+            // Create a virtual display and launch an activity on virtual & default display.
+            final ActivityDisplay newDisplay = virtualDisplaySession.setShowSystemDecorations(true)
+                    .setSimulateDisplay(true).createDisplay();
+            imeTestActivitySession.launchTestActivityOnDisplaySync(ImeTestActivity.class,
+                    DEFAULT_DISPLAY);
+            imeTestActivitySession2.launchTestActivityOnDisplaySync(ImeTestActivity2.class,
+                    newDisplay.mId);
 
-        final Display defDisplay = mAmWmState.getWmState().getDisplay(DEFAULT_DISPLAY);
-        final ImeEventStream stream = mockImeSession.openEventStream();
+            final Display defDisplay = mAmWmState.getWmState().getDisplay(DEFAULT_DISPLAY);
+            final ImeEventStream stream = mockImeSession1.openEventStream();
 
-        // Tap default display as top focused display & request focus on EditText to show
-        // soft input.
-        tapOnDisplayCenter(defDisplay.getDisplayId());
-        imeTestActivitySession.runOnMainSyncAndWait(
-                imeTestActivitySession.getActivity()::showSoftInput);
-        waitOrderedImeEventsThenAssertImeShown(stream, defDisplay.getDisplayId(),
-                editorMatcher("onStartInput",
-                        imeTestActivitySession.getActivity().mEditText.getPrivateImeOptions()),
-                event -> "showSoftInput".equals(event.getEventName()));
+            // Tap default display as top focused display & request focus on EditText to show
+            // soft input.
+            tapOnDisplayCenter(defDisplay.getDisplayId());
+            imeTestActivitySession.runOnMainSyncAndWait(
+                    imeTestActivitySession.getActivity()::showSoftInput);
+            waitOrderedImeEventsThenAssertImeShown(stream, defDisplay.getDisplayId(),
+                    editorMatcher("onStartInput",
+                            imeTestActivitySession.getActivity().mEditText.getPrivateImeOptions()),
+                    event -> "showSoftInput".equals(event.getEventName()));
 
-        // Tap virtual display as top focused display & request focus on EditText to show
-        // soft input.
-        tapOnDisplayCenter(newDisplay.mId);
-        imeTestActivitySession2.runOnMainSyncAndWait(
-                imeTestActivitySession2.getActivity()::showSoftInput);
-        waitOrderedImeEventsThenAssertImeShown(stream, newDisplay.mId,
-                editorMatcher("onStartInput",
-                        imeTestActivitySession2.getActivity().mEditText.getPrivateImeOptions()),
-                event -> "showSoftInput".equals(event.getEventName()));
+            // Tap virtual display as top focused display & request focus on EditText to show
+            // soft input.
+            tapOnDisplayCenter(newDisplay.mId);
+            imeTestActivitySession2.runOnMainSyncAndWait(
+                    imeTestActivitySession2.getActivity()::showSoftInput);
+            waitOrderedImeEventsThenAssertImeShown(stream, newDisplay.mId,
+                    editorMatcher("onStartInput",
+                            imeTestActivitySession2.getActivity().mEditText.getPrivateImeOptions()),
+                    event -> "showSoftInput".equals(event.getEventName()));
 
-        // Tap default display again to make sure the IME window will come back.
-        tapOnDisplayCenter(defDisplay.getDisplayId());
-        imeTestActivitySession.runOnMainSyncAndWait(
-                imeTestActivitySession.getActivity()::showSoftInput);
-        waitOrderedImeEventsThenAssertImeShown(stream, defDisplay.getDisplayId(),
-                editorMatcher("onStartInput",
-                        imeTestActivitySession.getActivity().mEditText.getPrivateImeOptions()),
-                event -> "showSoftInput".equals(event.getEventName()));
+            // Tap default display again to make sure the IME window will come back.
+            tapOnDisplayCenter(defDisplay.getDisplayId());
+            imeTestActivitySession.runOnMainSyncAndWait(
+                    imeTestActivitySession.getActivity()::showSoftInput);
+            waitOrderedImeEventsThenAssertImeShown(stream, defDisplay.getDisplayId(),
+                    editorMatcher("onStartInput",
+                            imeTestActivitySession.getActivity().mEditText.getPrivateImeOptions()),
+                    event -> "showSoftInput".equals(event.getEventName()));
+        }
     }
 
     /**
@@ -513,46 +537,49 @@
     public void testCrossDisplayBasicImeOperations() throws Exception {
         final long TIMEOUT = TimeUnit.SECONDS.toMillis(5);
 
-        final MockImeSession mockImeSession = createManagedMockImeSession(this);
-        final TestActivitySession<ImeTestActivity> imeTestActivitySession =
-                createManagedTestActivitySession();
+        try (final VirtualDisplaySession virtualDisplaySession  = new VirtualDisplaySession();
+             final TestActivitySession<ImeTestActivity>
+                     imeTestActivitySession = new TestActivitySession<>();
+             // Leverage MockImeSession to ensure at least a test Ime exists as default.
+             final MockImeSession mockImeSession = MockImeSession.create(
+                     mContext, getInstrumentation().getUiAutomation(), new ImeSettings.Builder())) {
 
-        // Create a virtual display by app and assume the display should not show IME window.
-        final ActivityDisplay newDisplay = createManagedVirtualDisplaySession()
-                .setPublicDisplay(true)
-                .createDisplay();
-        SystemUtil.runWithShellPermissionIdentity(
-                () -> assertFalse("Display should not support showing IME window",
-                        mTargetContext.getSystemService(WindowManager.class)
-                                .shouldShowIme(newDisplay.mId)));
+            // Create a virtual display by app and assume the display should not show IME window.
+            final ActivityDisplay newDisplay = virtualDisplaySession.setPublicDisplay(true)
+                    .createDisplay();
+            SystemUtil.runWithShellPermissionIdentity(
+                    () -> assertFalse("Display should not support showing IME window",
+                            mTargetContext.getSystemService(WindowManager.class)
+                                    .shouldShowIme(newDisplay.mId)));
 
-        // Launch Ime test activity in virtual display.
-        imeTestActivitySession.launchTestActivityOnDisplaySync(ImeTestActivity.class,
-                newDisplay.mId);
+            // Launch Ime test activity in virtual display.
+            imeTestActivitySession.launchTestActivityOnDisplaySync(ImeTestActivity.class,
+                    newDisplay.mId);
 
-        // Verify the activity to show soft input on the default display.
-        final ImeEventStream stream = mockImeSession.openEventStream();
-        final EditText editText = imeTestActivitySession.getActivity().mEditText;
-        imeTestActivitySession.runOnMainSyncAndWait(
-                imeTestActivitySession.getActivity()::showSoftInput);
-        waitOrderedImeEventsThenAssertImeShown(stream, DEFAULT_DISPLAY,
-                editorMatcher("onStartInput", editText.getPrivateImeOptions()),
-                event -> "showSoftInput".equals(event.getEventName()));
+            // Verify the activity to show soft input on the default display.
+            final ImeEventStream stream = mockImeSession.openEventStream();
+            final EditText editText = imeTestActivitySession.getActivity().mEditText;
+            imeTestActivitySession.runOnMainSyncAndWait(
+                    imeTestActivitySession.getActivity()::showSoftInput);
+            waitOrderedImeEventsThenAssertImeShown(stream, DEFAULT_DISPLAY,
+                    editorMatcher("onStartInput", editText.getPrivateImeOptions()),
+                    event -> "showSoftInput".equals(event.getEventName()));
 
-        // Commit text & make sure the input texts should be delivered to focused EditText on
-        // virtual display.
-        final String commitText = "test commit";
-        expectCommand(stream, mockImeSession.callCommitText(commitText, 1), TIMEOUT);
-        imeTestActivitySession.runOnMainAndAssertWithTimeout(
-                () -> TextUtils.equals(commitText, editText.getText()), TIMEOUT,
-                "The input text should be delivered");
+            // Commit text & make sure the input texts should be delivered to focused EditText on
+            // virtual display.
+            final String commitText = "test commit";
+            expectCommand(stream, mockImeSession.callCommitText(commitText, 1), TIMEOUT);
+            imeTestActivitySession.runOnMainAndAssertWithTimeout(
+                    () -> TextUtils.equals(commitText, editText.getText()), TIMEOUT,
+                    "The input text should be delivered");
 
-        // Since the IME and the IME target app are running in different displays,
-        // InputConnection#requestCursorUpdates() is not supported and it should return false.
-        // See InputMethodServiceTest#testOnUpdateCursorAnchorInfo() for the normal scenario.
-        final ImeCommand callCursorUpdates = mockImeSession.callRequestCursorUpdates(
-                InputConnection.CURSOR_UPDATE_IMMEDIATE);
-        assertFalse(expectCommand(stream, callCursorUpdates, TIMEOUT).getReturnBooleanValue());
+            // Since the IME and the IME target app are running in different displays,
+            // InputConnection#requestCursorUpdates() is not supported and it should return false.
+            // See InputMethodServiceTest#testOnUpdateCursorAnchorInfo() for the normal scenario.
+            final ImeCommand callCursorUpdates = mockImeSession.callRequestCursorUpdates(
+                    InputConnection.CURSOR_UPDATE_IMMEDIATE);
+            assertFalse(expectCommand(stream, callCursorUpdates, TIMEOUT).getReturnBooleanValue());
+        }
     }
 
     public static class ImeTestActivity extends Activity {
@@ -642,8 +669,7 @@
         assertEquals("Display bounds not the same", displayBounds, displayBoundsForIme);
     }
 
-    @SafeVarargs
-    final void waitOrderedImeEventsThenAssertImeShown(ImeEventStream stream, int displayId,
+    void waitOrderedImeEventsThenAssertImeShown(ImeEventStream stream, int displayId,
             Predicate<ImeEvent>... conditions) throws Exception {
         for (Predicate<ImeEvent> condition : conditions) {
             expectEvent(stream, condition, TimeUnit.SECONDS.toMillis(5) /* eventTimeout */);
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 f29b6c6..eb90555 100644
--- a/tests/framework/base/windowmanager/src/android/server/wm/MultiDisplayTestBase.java
+++ b/tests/framework/base/windowmanager/src/android/server/wm/MultiDisplayTestBase.java
@@ -55,7 +55,6 @@
 import android.server.wm.CommandSession.ActivitySession;
 import android.server.wm.CommandSession.ActivitySessionClient;
 import android.server.wm.settings.SettingsSession;
-import android.util.Pair;
 import android.util.Size;
 import android.view.WindowManager;
 
@@ -248,11 +247,6 @@
             ActivityLauncher::hasCaughtSecurityException);
     }
 
-    /** @see ObjectTracker#manage(AutoCloseable) */
-    protected VirtualDisplaySession createManagedVirtualDisplaySession() {
-        return mObjectTracker.manage(new VirtualDisplaySession());
-    }
-
     /**
      * This class should only be used when you need to test virtual display created by a
      * non-privileged app.
@@ -334,12 +328,12 @@
         }
 
         @Nullable
-        public ActivityDisplay createDisplay() {
+        public ActivityDisplay createDisplay() throws Exception {
             return createDisplays(1).stream().findFirst().orElse(null);
         }
 
         @NonNull
-        List<ActivityDisplay> createDisplays(int count) {
+        List<ActivityDisplay> createDisplays(int count) throws Exception {
             if (mSimulateDisplay) {
                 return simulateDisplay();
             } else {
@@ -353,7 +347,7 @@
         }
 
         @Override
-        public void close()  {
+        public void close() throws Exception {
             mOverlayDisplayDeviceSession.close();
             if (mVirtualDisplayCreated) {
                 destroyVirtualDisplays();
@@ -368,7 +362,7 @@
          * </pre>
          * @return {@link ActivityDisplay} of newly created display.
          */
-        private List<ActivityDisplay> simulateDisplay() {
+        private List<ActivityDisplay> simulateDisplay() throws Exception {
             // Create virtual display with custom density dpi and specified size.
             mOverlayDisplayDeviceSession.set(mSimulationDisplaySize + "/" + mDensityDpi);
             if (mShowSystemDecorations) {
@@ -488,7 +482,7 @@
         }
 
         @Override
-        public void close() {
+        public void close() throws Exception {
             super.close();
             mActivitySessionClient.close();
         }
@@ -557,7 +551,7 @@
         }
 
         @Override
-        public void close() {
+        public void close() throws Exception {
             // Need to restore display state before display is destroyed.
             restoreDisplayStates();
             super.close();
@@ -636,30 +630,11 @@
         return newDisplays;
     }
 
-    /** A clearer alias of {@link Pair#create(Object, Object)}. */
-    protected <K, V> Pair<K, V> pair(K k, V v) {
-        return new Pair<>(k, v);
-    }
-
-    protected void assertBothDisplaysHaveResumedActivities(
-            Pair<Integer, ComponentName> firstPair, Pair<Integer, ComponentName> secondPair) {
-        mAmWmState.assertResumedActivities("Both displays must have resumed activities",
-                mapping -> {
-                    mapping.put(firstPair.first, firstPair.second);
-                    mapping.put(secondPair.first, secondPair.second);
-                });
-    }
-
     /** Checks if the device supports multi-display. */
     protected boolean supportsMultiDisplay() {
         return hasDeviceFeature(FEATURE_ACTIVITIES_ON_SECONDARY_DISPLAYS);
     }
 
-    /** @see ObjectTracker#manage(AutoCloseable) */
-    protected ExternalDisplaySession createManagedExternalDisplaySession() {
-        return mObjectTracker.manage(new ExternalDisplaySession());
-    }
-
     /**
      * This class is used when you need to test virtual display created by a privileged app.
      *
@@ -695,7 +670,7 @@
         /**
          * Creates a private virtual display with insecure keyguard flags set.
          */
-        ActivityDisplay createVirtualDisplay() {
+        ActivityDisplay createVirtualDisplay() throws Exception {
             final List<ActivityDisplay> originalDS = getDisplaysStates();
             final int originalDisplayCount = originalDS.size();
 
@@ -732,7 +707,7 @@
         }
 
         @Override
-        public void close() {
+        public void close() throws Exception {
             if (mExternalDisplayHelper != null) {
                 mExternalDisplayHelper.releaseDisplay();
                 mExternalDisplayHelper = null;
@@ -750,7 +725,7 @@
         }
 
         @Override
-        public void close() {
+        public void close() throws Exception {
             setPrimaryDisplayState(true);
         }
 
diff --git a/tests/framework/base/windowmanager/src/android/server/wm/OverrideConfigTests.java b/tests/framework/base/windowmanager/src/android/server/wm/OverrideConfigTests.java
index 476a318..a13f881 100644
--- a/tests/framework/base/windowmanager/src/android/server/wm/OverrideConfigTests.java
+++ b/tests/framework/base/windowmanager/src/android/server/wm/OverrideConfigTests.java
@@ -36,24 +36,25 @@
 public class OverrideConfigTests extends ActivityManagerTestBase {
 
     @Test
-    public void testReceiveOverrideConfigFromRelayout() {
+    public void testReceiveOverrideConfigFromRelayout() throws Exception {
         assumeTrue("Device doesn't support freeform. Skipping test.", supportsFreeform());
 
         launchActivity(LOG_CONFIGURATION_ACTIVITY, WINDOWING_MODE_FREEFORM);
 
-        final RotationSession rotationSession = createManagedRotationSession();
-        rotationSession.set(ROTATION_0);
-        separateTestJournal();
-        resizeActivityTask(LOG_CONFIGURATION_ACTIVITY, 0, 0, 100, 100);
-        new ActivityLifecycleCounts(LOG_CONFIGURATION_ACTIVITY).assertCountWithRetry(
-                "Expected to observe configuration change when resizing",
-                countSpec(ActivityCallback.ON_CONFIGURATION_CHANGED, CountSpec.EQUALS, 1));
+        try (final RotationSession rotationSession = new RotationSession()) {
+            rotationSession.set(ROTATION_0);
+            separateTestJournal();
+            resizeActivityTask(LOG_CONFIGURATION_ACTIVITY, 0, 0, 100, 100);
+            new ActivityLifecycleCounts(LOG_CONFIGURATION_ACTIVITY).assertCountWithRetry(
+                    "Expected to observe configuration change when resizing",
+                    countSpec(ActivityCallback.ON_CONFIGURATION_CHANGED, CountSpec.EQUALS, 1));
 
-        separateTestJournal();
-        rotationSession.set(ROTATION_180);
-        new ActivityLifecycleCounts(LOG_CONFIGURATION_ACTIVITY).assertCountWithRetry(
-                "Not expected to observe configuration change after flip rotation",
-                countSpec(ActivityCallback.ON_CONFIGURATION_CHANGED, CountSpec.EQUALS, 0));
+            separateTestJournal();
+            rotationSession.set(ROTATION_180);
+            new ActivityLifecycleCounts(LOG_CONFIGURATION_ACTIVITY).assertCountWithRetry(
+                    "Not expected to observe configuration change after flip rotation",
+                    countSpec(ActivityCallback.ON_CONFIGURATION_CHANGED, CountSpec.EQUALS, 0));
+        }
     }
 }
 
diff --git a/tests/framework/base/windowmanager/src/android/server/wm/PinnedStackTests.java b/tests/framework/base/windowmanager/src/android/server/wm/PinnedStackTests.java
index 3e8c8d3..298276a 100644
--- a/tests/framework/base/windowmanager/src/android/server/wm/PinnedStackTests.java
+++ b/tests/framework/base/windowmanager/src/android/server/wm/PinnedStackTests.java
@@ -209,57 +209,59 @@
     }
 
     @Test
-    public void testPinnedStackDefaultBounds() {
+    public void testPinnedStackDefaultBounds() throws Exception {
         // Launch a PIP activity
         launchActivity(PIP_ACTIVITY, EXTRA_ENTER_PIP, "true");
         // Wait for animation complete since we are comparing bounds
         waitForEnterPipAnimationComplete(PIP_ACTIVITY);
 
-        final RotationSession rotationSession = createManagedRotationSession();
-        rotationSession.set(ROTATION_0);
-        waitForValidPinnedStackBounds(WindowManagerState::getDefaultPinnedStackBounds);
-        WindowManagerState wmState = mAmWmState.getWmState();
-        wmState.computeState();
-        Rect defaultPipBounds = wmState.getDefaultPinnedStackBounds();
-        Rect stableBounds = wmState.getStableBounds();
-        assertTrue(defaultPipBounds.width() > 0 && defaultPipBounds.height() > 0);
-        assertTrue(stableBounds.contains(defaultPipBounds));
+        try (final RotationSession rotationSession = new RotationSession()) {
+            rotationSession.set(ROTATION_0);
+            waitForValidPinnedStackBounds(WindowManagerState::getDefaultPinnedStackBounds);
+            WindowManagerState wmState = mAmWmState.getWmState();
+            wmState.computeState();
+            Rect defaultPipBounds = wmState.getDefaultPinnedStackBounds();
+            Rect stableBounds = wmState.getStableBounds();
+            assertTrue(defaultPipBounds.width() > 0 && defaultPipBounds.height() > 0);
+            assertTrue(stableBounds.contains(defaultPipBounds));
 
-        rotationSession.set(ROTATION_90);
-        waitForValidPinnedStackBounds(WindowManagerState::getDefaultPinnedStackBounds);
-        wmState = mAmWmState.getWmState();
-        wmState.computeState();
-        defaultPipBounds = wmState.getDefaultPinnedStackBounds();
-        stableBounds = wmState.getStableBounds();
-        assertTrue(defaultPipBounds.width() > 0 && defaultPipBounds.height() > 0);
-        assertTrue(stableBounds.contains(defaultPipBounds));
+            rotationSession.set(ROTATION_90);
+            waitForValidPinnedStackBounds(WindowManagerState::getDefaultPinnedStackBounds);
+            wmState = mAmWmState.getWmState();
+            wmState.computeState();
+            defaultPipBounds = wmState.getDefaultPinnedStackBounds();
+            stableBounds = wmState.getStableBounds();
+            assertTrue(defaultPipBounds.width() > 0 && defaultPipBounds.height() > 0);
+            assertTrue(stableBounds.contains(defaultPipBounds));
+        }
     }
 
     @Test
-    public void testPinnedStackMovementBounds() {
+    public void testPinnedStackMovementBounds() throws Exception {
         // Launch a PIP activity
         launchActivity(PIP_ACTIVITY, EXTRA_ENTER_PIP, "true");
         // Wait for animation complete since we are comparing bounds
         waitForEnterPipAnimationComplete(PIP_ACTIVITY);
 
-        final RotationSession rotationSession = createManagedRotationSession();
-        rotationSession.set(ROTATION_0);
-        waitForValidPinnedStackBounds(WindowManagerState::getPinnedStackMovementBounds);
-        WindowManagerState wmState = mAmWmState.getWmState();
-        wmState.computeState();
-        Rect pipMovementBounds = wmState.getPinnedStackMovementBounds();
-        Rect stableBounds = wmState.getStableBounds();
-        assertTrue(pipMovementBounds.width() > 0 && pipMovementBounds.height() > 0);
-        assertTrue(stableBounds.contains(pipMovementBounds));
+        try (final RotationSession rotationSession = new RotationSession()) {
+            rotationSession.set(ROTATION_0);
+            waitForValidPinnedStackBounds(WindowManagerState::getPinnedStackMovementBounds);
+            WindowManagerState wmState = mAmWmState.getWmState();
+            wmState.computeState();
+            Rect pipMovementBounds = wmState.getPinnedStackMovementBounds();
+            Rect stableBounds = wmState.getStableBounds();
+            assertTrue(pipMovementBounds.width() > 0 && pipMovementBounds.height() > 0);
+            assertTrue(stableBounds.contains(pipMovementBounds));
 
-        rotationSession.set(ROTATION_90);
-        waitForValidPinnedStackBounds(WindowManagerState::getPinnedStackMovementBounds);
-        wmState = mAmWmState.getWmState();
-        wmState.computeState();
-        pipMovementBounds = wmState.getPinnedStackMovementBounds();
-        stableBounds = wmState.getStableBounds();
-        assertTrue(pipMovementBounds.width() > 0 && pipMovementBounds.height() > 0);
-        assertTrue(stableBounds.contains(pipMovementBounds));
+            rotationSession.set(ROTATION_90);
+            waitForValidPinnedStackBounds(WindowManagerState::getPinnedStackMovementBounds);
+            wmState = mAmWmState.getWmState();
+            wmState.computeState();
+            pipMovementBounds = wmState.getPinnedStackMovementBounds();
+            stableBounds = wmState.getStableBounds();
+            assertTrue(pipMovementBounds.width() > 0 && pipMovementBounds.height() > 0);
+            assertTrue(stableBounds.contains(pipMovementBounds));
+        }
     }
 
     private void waitForValidPinnedStackBounds(Function<WindowManagerState, Rect> boundsFunc) {
@@ -302,7 +304,7 @@
     }
 
     @Test
-    public void testPinnedStackInBoundsAfterRotation() {
+    public void testPinnedStackInBoundsAfterRotation() throws Exception {
         // Launch an activity into the pinned stack
         launchActivity(PIP_ACTIVITY,
                 EXTRA_ENTER_PIP, "true",
@@ -311,15 +313,16 @@
         waitForEnterPipAnimationComplete(PIP_ACTIVITY);
 
         // Ensure that the PIP stack is fully visible in each orientation
-        final RotationSession rotationSession = createManagedRotationSession();
-        rotationSession.set(ROTATION_0);
-        assertPinnedStackActivityIsInDisplayBounds(PIP_ACTIVITY);
-        rotationSession.set(ROTATION_90);
-        assertPinnedStackActivityIsInDisplayBounds(PIP_ACTIVITY);
-        rotationSession.set(ROTATION_180);
-        assertPinnedStackActivityIsInDisplayBounds(PIP_ACTIVITY);
-        rotationSession.set(ROTATION_270);
-        assertPinnedStackActivityIsInDisplayBounds(PIP_ACTIVITY);
+        try (final RotationSession rotationSession = new RotationSession()) {
+            rotationSession.set(ROTATION_0);
+            assertPinnedStackActivityIsInDisplayBounds(PIP_ACTIVITY);
+            rotationSession.set(ROTATION_90);
+            assertPinnedStackActivityIsInDisplayBounds(PIP_ACTIVITY);
+            rotationSession.set(ROTATION_180);
+            assertPinnedStackActivityIsInDisplayBounds(PIP_ACTIVITY);
+            rotationSession.set(ROTATION_270);
+            assertPinnedStackActivityIsInDisplayBounds(PIP_ACTIVITY);
+        }
     }
 
     @Test
@@ -844,7 +847,7 @@
         }
 
         @Override
-        public void close() {
+        public void close() throws Exception {
             // Wait for the restored setting to apply before we continue on with the next test
             final CountDownLatch waitLock = new CountDownLatch(1);
             final Context context = getInstrumentation().getTargetContext();
@@ -856,44 +859,44 @@
                         }
                     });
             super.close();
-            try {
-                if (!waitLock.await(2, TimeUnit.SECONDS)) {
-                    Log.i(TAG, "TransitionAnimationScaleSession value not restored");
-                }
-            } catch (InterruptedException impossible) {}
+            if (!waitLock.await(2, TimeUnit.SECONDS)) {
+                Log.i(TAG, "TransitionAnimationScaleSession value not restored");
+            }
         }
     }
 
     @Test
-    public void testEnterPipInterruptedCallbacks() {
-        final TransitionAnimationScaleSession transitionAnimationScaleSession =
-                mObjectTracker.manage(new TransitionAnimationScaleSession());
-        // Slow down the transition animations for this test
-        transitionAnimationScaleSession.set(20f);
+    public void testEnterPipInterruptedCallbacks() throws Exception {
+        try (final TransitionAnimationScaleSession transitionAnimationScaleSession =
+                new TransitionAnimationScaleSession()) {
+            // Slow down the transition animations for this test
+            transitionAnimationScaleSession.set(20f);
 
-        // Launch a PiP activity
-        launchActivity(PIP_ACTIVITY, EXTRA_ENTER_PIP, "true");
-        // Wait until the PiP activity has moved into the pinned stack (happens before the
-        // transition has started)
-        waitForEnterPip(PIP_ACTIVITY);
-        assertPinnedStackExists();
+            // Launch a PiP activity
+            launchActivity(PIP_ACTIVITY, EXTRA_ENTER_PIP, "true");
+            // Wait until the PiP activity has moved into the pinned stack (happens before the
+            // transition has started)
+            waitForEnterPip(PIP_ACTIVITY);
+            assertPinnedStackExists();
 
-        // Relaunch the PiP activity back into fullscreen
-        separateTestJournal();
-        launchActivity(PIP_ACTIVITY);
-        // Wait until the PiP activity is reparented into the fullscreen stack (happens after
-        // the transition has finished)
-        waitForExitPipToFullscreen(PIP_ACTIVITY);
+            // Relaunch the PiP activity back into fullscreen
+            separateTestJournal();
+            launchActivity(PIP_ACTIVITY);
+            // Wait until the PiP activity is reparented into the fullscreen stack (happens after
+            // the transition has finished)
+            waitForExitPipToFullscreen(PIP_ACTIVITY);
 
-        // Ensure that we get the callbacks indicating that PiP/MW mode was cancelled, but no
-        // configuration change (since none was sent)
-        final ActivityLifecycleCounts lifecycleCounts = new ActivityLifecycleCounts(PIP_ACTIVITY);
-        assertEquals("onConfigurationChanged", 0,
-                lifecycleCounts.getCount(ActivityCallback.ON_CONFIGURATION_CHANGED));
-        assertEquals("onPictureInPictureModeChanged", 1,
-                lifecycleCounts.getCount(ActivityCallback.ON_PICTURE_IN_PICTURE_MODE_CHANGED));
-        assertEquals("onMultiWindowModeChanged", 1,
-                lifecycleCounts.getCount(ActivityCallback.ON_MULTI_WINDOW_MODE_CHANGED));
+            // Ensure that we get the callbacks indicating that PiP/MW mode was cancelled, but no
+            // configuration change (since none was sent)
+            final ActivityLifecycleCounts lifecycleCounts = new ActivityLifecycleCounts(
+                    PIP_ACTIVITY);
+            assertEquals("onConfigurationChanged", 0,
+                    lifecycleCounts.getCount(ActivityCallback.ON_CONFIGURATION_CHANGED));
+            assertEquals("onPictureInPictureModeChanged", 1,
+                    lifecycleCounts.getCount(ActivityCallback.ON_PICTURE_IN_PICTURE_MODE_CHANGED));
+            assertEquals("onMultiWindowModeChanged", 1,
+                    lifecycleCounts.getCount(ActivityCallback.ON_MULTI_WINDOW_MODE_CHANGED));
+        }
     }
 
     @Test
diff --git a/tests/framework/base/windowmanager/src/android/server/wm/PrereleaseSdkTest.java b/tests/framework/base/windowmanager/src/android/server/wm/PrereleaseSdkTest.java
index 8816bb5..4747e49 100644
--- a/tests/framework/base/windowmanager/src/android/server/wm/PrereleaseSdkTest.java
+++ b/tests/framework/base/windowmanager/src/android/server/wm/PrereleaseSdkTest.java
@@ -48,7 +48,10 @@
     }
 
     @After
-    public void tearDown() {
+    @Override
+    public void tearDown() throws Exception {
+        super.tearDown();
+
         // Ensure app process is stopped.
         stopTestPackage(MAIN_ACTIVITY.getPackageName());
     }
diff --git a/tests/framework/base/windowmanager/src/android/server/wm/SplitScreenTests.java b/tests/framework/base/windowmanager/src/android/server/wm/SplitScreenTests.java
index a410ee1..5d4efb5 100644
--- a/tests/framework/base/windowmanager/src/android/server/wm/SplitScreenTests.java
+++ b/tests/framework/base/windowmanager/src/android/server/wm/SplitScreenTests.java
@@ -353,7 +353,7 @@
     }
 
     @Test
-    public void testRotationWhenDocked() {
+    public void testRotationWhenDocked() throws Exception {
         launchActivitiesInSplitScreen(
                 getLaunchActivityBuilder().setTargetActivity(LAUNCHING_ACTIVITY),
                 getLaunchActivityBuilder().setTargetActivity(TEST_ACTIVITY));
@@ -364,27 +364,28 @@
 
         // Rotate device single steps (90°) 0-1-2-3.
         // Each time we compute the state we implicitly assert valid bounds.
-        final RotationSession rotationSession = createManagedRotationSession();
-        for (int i = 0; i < 4; i++) {
-            rotationSession.set(i);
+        try (final RotationSession rotationSession = new RotationSession()) {
+            for (int i = 0; i < 4; i++) {
+                rotationSession.set(i);
+                mAmWmState.computeState(LAUNCHING_ACTIVITY, TEST_ACTIVITY);
+            }
+            // Double steps (180°) We ended the single step at 3. So, we jump directly to 1 for
+            // double step. So, we are testing 3-1-3 for one side and 0-2-0 for the other side.
+            rotationSession.set(ROTATION_90);
+            mAmWmState.computeState(LAUNCHING_ACTIVITY, TEST_ACTIVITY);
+            rotationSession.set(ROTATION_270);
+            mAmWmState.computeState(LAUNCHING_ACTIVITY, TEST_ACTIVITY);
+            rotationSession.set(ROTATION_0);
+            mAmWmState.computeState(LAUNCHING_ACTIVITY, TEST_ACTIVITY);
+            rotationSession.set(ROTATION_180);
+            mAmWmState.computeState(LAUNCHING_ACTIVITY, TEST_ACTIVITY);
+            rotationSession.set(ROTATION_0);
             mAmWmState.computeState(LAUNCHING_ACTIVITY, TEST_ACTIVITY);
         }
-        // Double steps (180°) We ended the single step at 3. So, we jump directly to 1 for
-        // double step. So, we are testing 3-1-3 for one side and 0-2-0 for the other side.
-        rotationSession.set(ROTATION_90);
-        mAmWmState.computeState(LAUNCHING_ACTIVITY, TEST_ACTIVITY);
-        rotationSession.set(ROTATION_270);
-        mAmWmState.computeState(LAUNCHING_ACTIVITY, TEST_ACTIVITY);
-        rotationSession.set(ROTATION_0);
-        mAmWmState.computeState(LAUNCHING_ACTIVITY, TEST_ACTIVITY);
-        rotationSession.set(ROTATION_180);
-        mAmWmState.computeState(LAUNCHING_ACTIVITY, TEST_ACTIVITY);
-        rotationSession.set(ROTATION_0);
-        mAmWmState.computeState(LAUNCHING_ACTIVITY, TEST_ACTIVITY);
     }
 
     @Test
-    public void testRotationWhenDockedWhileLocked() {
+    public void testRotationWhenDockedWhileLocked() throws Exception {
         launchActivitiesInSplitScreen(
                 getLaunchActivityBuilder().setTargetActivity(LAUNCHING_ACTIVITY),
                 getLaunchActivityBuilder().setTargetActivity(TEST_ACTIVITY));
@@ -394,94 +395,99 @@
         mAmWmState.assertContainsStack("Must contain docked stack.",
                 WINDOWING_MODE_SPLIT_SCREEN_PRIMARY, ACTIVITY_TYPE_STANDARD);
 
-        final RotationSession rotationSession = createManagedRotationSession();
-        final LockScreenSession lockScreenSession = createManagedLockScreenSession();
-        for (int i = 0; i < 4; i++) {
-            lockScreenSession.sleepDevice();
-            // The display may not be rotated while device is locked.
-            rotationSession.set(i, false /* waitDeviceRotation */);
-            lockScreenSession.wakeUpDevice()
-                    .unlockDevice();
-            mAmWmState.computeState(LAUNCHING_ACTIVITY, TEST_ACTIVITY);
-        }
-    }
-
-    @Test
-    public void testMinimizedFromEachDockedSide() {
-        final RotationSession rotationSession = createManagedRotationSession();
-        for (int i = 0; i < 2; i++) {
-            rotationSession.set(i);
-            launchActivityInDockStackAndMinimize(TEST_ACTIVITY);
-            if (!mAmWmState.isScreenPortrait() && isTablet()) {
-                // Test minimize to the right only on tablets in landscape
-                removeStacksWithActivityTypes(ALL_ACTIVITY_TYPE_BUT_HOME);
-                launchActivityInDockStackAndMinimize(TEST_ACTIVITY,
-                        SPLIT_SCREEN_CREATE_MODE_BOTTOM_OR_RIGHT);
+        try (final RotationSession rotationSession = new RotationSession();
+             final LockScreenSession lockScreenSession = new LockScreenSession()) {
+            for (int i = 0; i < 4; i++) {
+                lockScreenSession.sleepDevice();
+                // The display may not be rotated while device is locked.
+                rotationSession.set(i, false /* waitDeviceRotation */);
+                lockScreenSession.wakeUpDevice()
+                        .unlockDevice();
+                mAmWmState.computeState(LAUNCHING_ACTIVITY, TEST_ACTIVITY);
             }
-            removeStacksWithActivityTypes(ALL_ACTIVITY_TYPE_BUT_HOME);
         }
     }
 
     @Test
-    public void testRotationWhileDockMinimized() {
+    public void testMinimizedFromEachDockedSide() throws Exception {
+        try (final RotationSession rotationSession = new RotationSession()) {
+            for (int i = 0; i < 2; i++) {
+                rotationSession.set(i);
+                launchActivityInDockStackAndMinimize(TEST_ACTIVITY);
+                if (!mAmWmState.isScreenPortrait() && isTablet()) {
+                    // Test minimize to the right only on tablets in landscape
+                    removeStacksWithActivityTypes(ALL_ACTIVITY_TYPE_BUT_HOME);
+                    launchActivityInDockStackAndMinimize(TEST_ACTIVITY,
+                            SPLIT_SCREEN_CREATE_MODE_BOTTOM_OR_RIGHT);
+                }
+                removeStacksWithActivityTypes(ALL_ACTIVITY_TYPE_BUT_HOME);
+            }
+        }
+    }
+
+    @Test
+    public void testRotationWhileDockMinimized() throws Exception {
         launchActivityInDockStackAndMinimize(TEST_ACTIVITY);
 
         // Rotate device single steps (90°) 0-1-2-3.
         // Each time we compute the state we implicitly assert valid bounds in minimized mode.
-        final RotationSession rotationSession = createManagedRotationSession();
-        for (int i = 0; i < 4; i++) {
-            rotationSession.set(i);
+        try (final RotationSession rotationSession = new RotationSession()) {
+            for (int i = 0; i < 4; i++) {
+                rotationSession.set(i);
+                mAmWmState.computeState(TEST_ACTIVITY);
+            }
+
+            // Double steps (180°) We ended the single step at 3. So, we jump directly to 1 for
+            // double step. So, we are testing 3-1-3 for one side and 0-2-0 for the other side in
+            // minimized mode.
+            rotationSession.set(ROTATION_90);
+            mAmWmState.computeState(TEST_ACTIVITY);
+            rotationSession.set(ROTATION_270);
+            mAmWmState.computeState(TEST_ACTIVITY);
+            rotationSession.set(ROTATION_0);
+            mAmWmState.computeState(TEST_ACTIVITY);
+            rotationSession.set(ROTATION_180);
+            mAmWmState.computeState(TEST_ACTIVITY);
+            rotationSession.set(ROTATION_0);
             mAmWmState.computeState(TEST_ACTIVITY);
         }
-
-        // Double steps (180°) We ended the single step at 3. So, we jump directly to 1 for
-        // double step. So, we are testing 3-1-3 for one side and 0-2-0 for the other side in
-        // minimized mode.
-        rotationSession.set(ROTATION_90);
-        mAmWmState.computeState(TEST_ACTIVITY);
-        rotationSession.set(ROTATION_270);
-        mAmWmState.computeState(TEST_ACTIVITY);
-        rotationSession.set(ROTATION_0);
-        mAmWmState.computeState(TEST_ACTIVITY);
-        rotationSession.set(ROTATION_180);
-        mAmWmState.computeState(TEST_ACTIVITY);
-        rotationSession.set(ROTATION_0);
-        mAmWmState.computeState(TEST_ACTIVITY);
     }
 
     @Test
-    public void testMinimizeAndUnminimizeThenGoingHome() {
+    public void testMinimizeAndUnminimizeThenGoingHome() throws Exception {
         // Rotate the screen to check that minimize, unminimize, dismiss the docked stack and then
         // going home has the correct app transition
-        final RotationSession rotationSession = createManagedRotationSession();
-        for (int rotation = ROTATION_0; rotation <= ROTATION_270; rotation++) {
-            launchActivityInDockStackAndMinimize(DOCKED_ACTIVITY);
-            // Set rotation after docked stack exists so the display can be rotated (home may
-            // be fixed-orientation if it is fullscreen).
-            rotationSession.set(rotation);
+        try (final RotationSession rotationSession = new RotationSession()) {
+            for (int rotation = ROTATION_0; rotation <= ROTATION_270; rotation++) {
+                launchActivityInDockStackAndMinimize(DOCKED_ACTIVITY);
+                // Set rotation after docked stack exists so the display can be rotated (home may
+                // be fixed-orientation if it is fullscreen).
+                rotationSession.set(rotation);
 
-            if (mIsHomeRecentsComponent) {
-                launchActivity(TEST_ACTIVITY, WINDOWING_MODE_FULLSCREEN_OR_SPLIT_SCREEN_SECONDARY);
-            } else {
-                // Unminimize the docked stack
-                pressAppSwitchButtonAndWaitForRecents();
-                waitForDockNotMinimized();
-                assertDockNotMinimized();
+                if (mIsHomeRecentsComponent) {
+                    launchActivity(TEST_ACTIVITY,
+                            WINDOWING_MODE_FULLSCREEN_OR_SPLIT_SCREEN_SECONDARY);
+                } else {
+                    // Unminimize the docked stack
+                    pressAppSwitchButtonAndWaitForRecents();
+                    waitForDockNotMinimized();
+                    assertDockNotMinimized();
+                }
+
+                // Dismiss the dock stack
+                setActivityTaskWindowingMode(DOCKED_ACTIVITY,
+                        WINDOWING_MODE_FULLSCREEN_OR_SPLIT_SCREEN_SECONDARY);
+                mAmWmState.computeState(DOCKED_ACTIVITY);
+
+                // Go home and check the app transition
+                assertNotEquals(TRANSIT_WALLPAPER_OPEN,
+                        mAmWmState.getWmState().getDefaultDisplayLastTransition());
+                launchHomeActivity();
+                mAmWmState.waitForAppTransitionIdleOnDisplay(DEFAULT_DISPLAY);
+
+                assertEquals(TRANSIT_WALLPAPER_OPEN,
+                        mAmWmState.getWmState().getDefaultDisplayLastTransition());
             }
-
-            // Dismiss the dock stack
-            setActivityTaskWindowingMode(DOCKED_ACTIVITY,
-                    WINDOWING_MODE_FULLSCREEN_OR_SPLIT_SCREEN_SECONDARY);
-            mAmWmState.computeState(DOCKED_ACTIVITY);
-
-            // Go home and check the app transition
-            assertNotEquals(TRANSIT_WALLPAPER_OPEN,
-                    mAmWmState.getWmState().getDefaultDisplayLastTransition());
-            launchHomeActivity();
-            mAmWmState.waitForAppTransitionIdleOnDisplay(DEFAULT_DISPLAY);
-
-            assertEquals(TRANSIT_WALLPAPER_OPEN,
-                    mAmWmState.getWmState().getDefaultDisplayLastTransition());
         }
     }
 
@@ -497,37 +503,41 @@
     }
 
     @Test
-    public void testDockedStackToMinimizeWhenUnlocked() {
+    public void testDockedStackToMinimizeWhenUnlocked() throws Exception {
         if (!mIsHomeRecentsComponent) {
             launchActivityInDockStackAndMinimize(TEST_ACTIVITY);
         } else {
             launchActivityInSplitScreenWithRecents(TEST_ACTIVITY);
         }
         mAmWmState.computeState(TEST_ACTIVITY);
-        createManagedLockScreenSession().sleepDevice()
-                .wakeUpDevice()
-                .unlockDevice();
-        mAmWmState.computeState(TEST_ACTIVITY);
-        assertDockMinimized();
+        try (final LockScreenSession lockScreenSession = new LockScreenSession()) {
+            lockScreenSession.sleepDevice()
+                    .wakeUpDevice()
+                    .unlockDevice();
+            mAmWmState.computeState(TEST_ACTIVITY);
+            assertDockMinimized();
+        }
     }
 
     @Test
-    public void testMinimizedStateWhenUnlockedAndUnMinimized() {
+    public void testMinimizedStateWhenUnlockedAndUnMinimized() throws Exception {
         launchActivityInDockStackAndMinimize(TEST_ACTIVITY);
 
-        createManagedLockScreenSession().sleepDevice()
-                .wakeUpDevice()
-                .unlockDevice();
-        mAmWmState.computeState(TEST_ACTIVITY);
+        try (final LockScreenSession lockScreenSession = new LockScreenSession()) {
+            lockScreenSession.sleepDevice()
+                    .wakeUpDevice()
+                    .unlockDevice();
+            mAmWmState.computeState(TEST_ACTIVITY);
 
-        // Unminimized back to splitscreen
-        if (mIsHomeRecentsComponent) {
-            launchActivity(RESIZEABLE_ACTIVITY,
-                    WINDOWING_MODE_FULLSCREEN_OR_SPLIT_SCREEN_SECONDARY);
-        } else {
-            pressAppSwitchButtonAndWaitForRecents();
+            // Unminimized back to splitscreen
+            if (mIsHomeRecentsComponent) {
+                launchActivity(RESIZEABLE_ACTIVITY,
+                        WINDOWING_MODE_FULLSCREEN_OR_SPLIT_SCREEN_SECONDARY);
+            } else {
+                pressAppSwitchButtonAndWaitForRecents();
+            }
+            mAmWmState.computeState(TEST_ACTIVITY);
         }
-        mAmWmState.computeState(TEST_ACTIVITY);
     }
 
     /**
diff --git a/tests/framework/base/windowmanager/src/android/server/wm/StartActivityTests.java b/tests/framework/base/windowmanager/src/android/server/wm/StartActivityTests.java
index 7f8bd8c..c4c6173 100644
--- a/tests/framework/base/windowmanager/src/android/server/wm/StartActivityTests.java
+++ b/tests/framework/base/windowmanager/src/android/server/wm/StartActivityTests.java
@@ -35,6 +35,7 @@
 import android.os.Bundle;
 import android.platform.test.annotations.Presubmit;
 import android.server.wm.CommandSession.ActivitySession;
+import android.server.wm.CommandSession.ActivitySessionClient;
 
 import androidx.test.rule.ActivityTestRule;
 
@@ -148,8 +149,10 @@
      */
     @Test
     public void testStartActivitiesInNewAndSameTask() {
-        final ActivitySession activity = createManagedActivityClientSession()
-                .startActivity(getLaunchActivityBuilder().setUseInstrumentation());
+        final ActivitySessionClient client = createActivitySessionClient();
+        final ActivitySession activity = client.startActivity(
+                getLaunchActivityBuilder().setUseInstrumentation()
+                        .setTargetActivity(TEST_ACTIVITY));
 
         final Intent[] intents = {
                 new Intent().setComponent(NO_RELAUNCH_ACTIVITY)
@@ -170,6 +173,7 @@
         final int callerTaskId = amState.getTaskByActivity(TEST_ACTIVITY).getTaskId();
         final int i0TaskId = amState.getTaskByActivity(intents[0].getComponent()).getTaskId();
         final int i1TaskId = amState.getTaskByActivity(intents[1].getComponent()).getTaskId();
+        client.close();
 
         assertNotEquals("The activities started by startActivities() should have a different task"
                 + " from their caller activity", callerTaskId, i0TaskId);
diff --git a/tests/framework/base/windowmanager/src/android/server/wm/ToastTest.java b/tests/framework/base/windowmanager/src/android/server/wm/ToastTest.java
index 821848e..39db6ff 100644
--- a/tests/framework/base/windowmanager/src/android/server/wm/ToastTest.java
+++ b/tests/framework/base/windowmanager/src/android/server/wm/ToastTest.java
@@ -38,7 +38,6 @@
 
 import com.android.compatibility.common.util.SystemUtil;
 
-import org.junit.After;
 import org.junit.Test;
 
 import java.util.Collections;
@@ -85,13 +84,14 @@
         mContext.registerReceiver(mAppCommunicator, filter);
     }
 
-    @After
-    public void tearDown() {
+    @Override
+    public void tearDown() throws Exception {
         mContext.unregisterReceiver(mAppCommunicator);
         SystemUtil.runWithShellPermissionIdentity(() -> {
             Settings.Global.putString(mContext.getContentResolver(), SETTING_HIDDEN_API_POLICY,
                     mPreviousHiddenApiPolicy);
         });
+        super.tearDown();
     }
 
     @Test
diff --git a/tests/framework/base/windowmanager/src/android/server/wm/VrDisplayTests.java b/tests/framework/base/windowmanager/src/android/server/wm/VrDisplayTests.java
index 9573cfa..a790a9a 100644
--- a/tests/framework/base/windowmanager/src/android/server/wm/VrDisplayTests.java
+++ b/tests/framework/base/windowmanager/src/android/server/wm/VrDisplayTests.java
@@ -70,14 +70,14 @@
     private static class VrModeSession implements AutoCloseable {
         private boolean applyVrModeChanges = !ActivityManagerTestBase.isUiModeLockedToVrHeadset();
 
-        void enablePersistentVrMode() {
+        void enablePersistentVrMode() throws Exception {
             if (!applyVrModeChanges) { return; }
             executeShellCommand("setprop vr_virtualdisplay true");
             executeShellCommand("vr set-persistent-vr-mode-enabled true");
         }
 
         @Override
-        public void close() {
+        public void close() throws Exception {
             if (!applyVrModeChanges) { return; }
             executeShellCommand("vr set-persistent-vr-mode-enabled false");
             executeShellCommand("setprop vr_virtualdisplay false");
@@ -96,7 +96,7 @@
                     Settings.Secure::putString);
         }
 
-        public void enableVrListener(ComponentName targetVrComponent) {
+        public void enableVrListener(ComponentName targetVrComponent) throws Exception {
             ComponentName component = new ComponentName(targetVrComponent.getPackageName(),
                     MockVrListenerService.class.getName());
             set(component.flattenToString());
@@ -107,97 +107,97 @@
      * Tests that any new activity launch in Vr mode is in Vr display.
      */
     @Test
-    public void testVrActivityLaunch() {
+    public void testVrActivityLaunch() throws Exception {
         assumeTrue(supportsMultiDisplay());
 
-        final VrModeSession vrModeSession = mObjectTracker.manage(new VrModeSession());
-        final EnableVrListenerSession enableVrListenerSession =
-                mObjectTracker.manage(new EnableVrListenerSession());
+        try (final VrModeSession vrModeSession = new VrModeSession();
+             final EnableVrListenerSession enableVrListenerSession =
+                     new EnableVrListenerSession()) {
+            // Put the device in persistent vr mode.
+            vrModeSession.enablePersistentVrMode();
+            enableVrListenerSession.enableVrListener(VR_TEST_ACTIVITY);
 
-        // Put the device in persistent vr mode.
-        vrModeSession.enablePersistentVrMode();
-        enableVrListenerSession.enableVrListener(VR_TEST_ACTIVITY);
+            // Launch the VR activity.
+            launchActivity(VR_TEST_ACTIVITY);
+            mAmWmState.computeState(VR_TEST_ACTIVITY);
+            mAmWmState.assertVisibility(VR_TEST_ACTIVITY, true /* visible */);
 
-        // Launch the VR activity.
-        launchActivity(VR_TEST_ACTIVITY);
-        mAmWmState.computeState(VR_TEST_ACTIVITY);
-        mAmWmState.assertVisibility(VR_TEST_ACTIVITY, true /* visible */);
+            // Launch the non-VR 2D activity and check where it ends up.
+            launchActivity(LAUNCHING_ACTIVITY);
+            mAmWmState.computeState(LAUNCHING_ACTIVITY);
 
-        // Launch the non-VR 2D activity and check where it ends up.
-        launchActivity(LAUNCHING_ACTIVITY);
-        mAmWmState.computeState(LAUNCHING_ACTIVITY);
+            // Ensure that the subsequent activity is visible
+            mAmWmState.assertVisibility(LAUNCHING_ACTIVITY, true /* visible */);
 
-        // Ensure that the subsequent activity is visible
-        mAmWmState.assertVisibility(LAUNCHING_ACTIVITY, true /* visible */);
+            // Check that activity is launched in focused stack on primary display.
+            mAmWmState.assertFocusedActivity("Launched activity must be focused",
+                    LAUNCHING_ACTIVITY);
+            final int focusedStackId = mAmWmState.getAmState().getFocusedStackId();
+            final ActivityManagerState.ActivityStack focusedStack
+                    = mAmWmState.getAmState().getStackById(focusedStackId);
+            assertEquals("Launched activity must be resumed in focused stack",
+                    getActivityName(LAUNCHING_ACTIVITY), focusedStack.mResumedActivity);
 
-        // Check that activity is launched in focused stack on primary display.
-        mAmWmState.assertFocusedActivity("Launched activity must be focused",
-                LAUNCHING_ACTIVITY);
-        final int focusedStackId = mAmWmState.getAmState().getFocusedStackId();
-        final ActivityManagerState.ActivityStack focusedStack
-                = mAmWmState.getAmState().getStackById(focusedStackId);
-        assertEquals("Launched activity must be resumed in focused stack",
-                getActivityName(LAUNCHING_ACTIVITY), focusedStack.mResumedActivity);
+            // Check if the launch activity is in Vr virtual display id.
+            final List<ActivityDisplay> reportedDisplays = getDisplaysStates();
+            final ActivityDisplay vrDisplay = getDisplayState(reportedDisplays,
+                    VR_VIRTUAL_DISPLAY_WIDTH, VR_VIRTUAL_DISPLAY_HEIGHT, VR_VIRTUAL_DISPLAY_DPI);
+            assertNotNull("Vr mode should have a virtual display", vrDisplay);
 
-        // Check if the launch activity is in Vr virtual display id.
-        final List<ActivityDisplay> reportedDisplays = getDisplaysStates();
-        final ActivityDisplay vrDisplay = getDisplayState(reportedDisplays,
-                VR_VIRTUAL_DISPLAY_WIDTH, VR_VIRTUAL_DISPLAY_HEIGHT, VR_VIRTUAL_DISPLAY_DPI);
-        assertNotNull("Vr mode should have a virtual display", vrDisplay);
-
-        // Check if the focused activity is on this virtual stack.
-        assertEquals("Launch in Vr mode should be in virtual stack", vrDisplay.mId,
-                focusedStack.mDisplayId);
+            // Check if the focused activity is on this virtual stack.
+            assertEquals("Launch in Vr mode should be in virtual stack", vrDisplay.mId,
+                    focusedStack.mDisplayId);
+        }
     }
 
     /**
      * Tests that any activity already present is re-launched in Vr display in vr mode.
      */
     @Test
-    public void testVrActivityReLaunch() {
+    public void testVrActivityReLaunch() throws Exception {
         assumeTrue(supportsMultiDisplay());
 
         // Launch a 2D activity.
         launchActivity(LAUNCHING_ACTIVITY);
 
-        final VrModeSession vrModeSession = mObjectTracker.manage(new VrModeSession());
-        final EnableVrListenerSession enableVrListenerSession =
-                mObjectTracker.manage(new EnableVrListenerSession());
+        try (final VrModeSession vrModeSession = new VrModeSession();
+             final EnableVrListenerSession enableVrListenerSession =
+                     new EnableVrListenerSession()) {
+            // Put the device in persistent vr mode.
+            vrModeSession.enablePersistentVrMode();
+            enableVrListenerSession.enableVrListener(VR_TEST_ACTIVITY);
 
-        // Put the device in persistent vr mode.
-        vrModeSession.enablePersistentVrMode();
-        enableVrListenerSession.enableVrListener(VR_TEST_ACTIVITY);
+            // Launch the VR activity.
+            launchActivity(VR_TEST_ACTIVITY);
+            mAmWmState.computeState(VR_TEST_ACTIVITY);
+            mAmWmState.assertVisibility(VR_TEST_ACTIVITY, true /* visible */);
 
-        // Launch the VR activity.
-        launchActivity(VR_TEST_ACTIVITY);
-        mAmWmState.computeState(VR_TEST_ACTIVITY);
-        mAmWmState.assertVisibility(VR_TEST_ACTIVITY, true /* visible */);
+            // Re-launch the non-VR 2D activity and check where it ends up.
+            launchActivity(LAUNCHING_ACTIVITY);
+            mAmWmState.computeState(LAUNCHING_ACTIVITY);
 
-        // Re-launch the non-VR 2D activity and check where it ends up.
-        launchActivity(LAUNCHING_ACTIVITY);
-        mAmWmState.computeState(LAUNCHING_ACTIVITY);
+            // Ensure that the subsequent activity is visible
+            mAmWmState.assertVisibility(LAUNCHING_ACTIVITY, true /* visible */);
 
-        // Ensure that the subsequent activity is visible
-        mAmWmState.assertVisibility(LAUNCHING_ACTIVITY, true /* visible */);
+            // Check that activity is launched in focused stack on primary display.
+            mAmWmState.assertFocusedActivity("Launched activity must be focused",
+                    LAUNCHING_ACTIVITY);
+            final int focusedStackId = mAmWmState.getAmState().getFocusedStackId();
+            final ActivityManagerState.ActivityStack focusedStack
+                    = mAmWmState.getAmState().getStackById(focusedStackId);
+            assertEquals("Launched activity must be resumed in focused stack",
+                    getActivityName(LAUNCHING_ACTIVITY), focusedStack.mResumedActivity);
 
-        // Check that activity is launched in focused stack on primary display.
-        mAmWmState.assertFocusedActivity("Launched activity must be focused",
-                LAUNCHING_ACTIVITY);
-        final int focusedStackId = mAmWmState.getAmState().getFocusedStackId();
-        final ActivityManagerState.ActivityStack focusedStack
-                = mAmWmState.getAmState().getStackById(focusedStackId);
-        assertEquals("Launched activity must be resumed in focused stack",
-                getActivityName(LAUNCHING_ACTIVITY), focusedStack.mResumedActivity);
+            // Check if the launch activity is in Vr virtual display id.
+            final List<ActivityDisplay> reportedDisplays = getDisplaysStates();
+            final ActivityDisplay vrDisplay = getDisplayState(reportedDisplays,
+                    VR_VIRTUAL_DISPLAY_WIDTH, VR_VIRTUAL_DISPLAY_HEIGHT, VR_VIRTUAL_DISPLAY_DPI);
+            assertNotNull("Vr mode should have a virtual display", vrDisplay);
 
-        // Check if the launch activity is in Vr virtual display id.
-        final List<ActivityDisplay> reportedDisplays = getDisplaysStates();
-        final ActivityDisplay vrDisplay = getDisplayState(reportedDisplays,
-                VR_VIRTUAL_DISPLAY_WIDTH, VR_VIRTUAL_DISPLAY_HEIGHT, VR_VIRTUAL_DISPLAY_DPI);
-        assertNotNull("Vr mode should have a virtual display", vrDisplay);
-
-        // Check if the focused activity is on this virtual stack.
-        assertEquals("Launch in Vr mode should be in virtual stack", vrDisplay.mId,
-                focusedStack.mDisplayId);
+            // Check if the focused activity is on this virtual stack.
+            assertEquals("Launch in Vr mode should be in virtual stack", vrDisplay.mId,
+                    focusedStack.mDisplayId);
+        }
     }
 
     /**
diff --git a/tests/framework/base/windowmanager/src/android/server/wm/WindowInsetsPolicyTest.java b/tests/framework/base/windowmanager/src/android/server/wm/WindowInsetsPolicyTest.java
index fa6212f..ba63060 100644
--- a/tests/framework/base/windowmanager/src/android/server/wm/WindowInsetsPolicyTest.java
+++ b/tests/framework/base/windowmanager/src/android/server/wm/WindowInsetsPolicyTest.java
@@ -21,6 +21,7 @@
 import static android.app.WindowConfiguration.WINDOWING_MODE_SPLIT_SCREEN_SECONDARY;
 import static android.content.res.Configuration.ORIENTATION_PORTRAIT;
 import static android.server.wm.app.Components.LAUNCHING_ACTIVITY;
+import static android.server.wm.app.Components.TEST_ACTIVITY;
 import static android.view.Display.DEFAULT_DISPLAY;
 import static android.view.Surface.ROTATION_0;
 import static android.view.Surface.ROTATION_90;
@@ -37,10 +38,12 @@
 import android.graphics.Insets;
 import android.os.Bundle;
 import android.platform.test.annotations.Presubmit;
+import android.util.Log;
 import android.view.View;
 import android.view.ViewGroup;
 import android.view.Window;
 import android.view.WindowInsets;
+import android.view.WindowManager;
 import android.view.WindowManager.LayoutParams;
 
 import androidx.test.rule.ActivityTestRule;
@@ -51,6 +54,7 @@
 import org.hamcrest.Matcher;
 import org.junit.Assert;
 import org.junit.Before;
+import org.junit.BeforeClass;
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.rules.ErrorCollector;
@@ -117,43 +121,44 @@
                 supportsSplitScreenMultiWindow());
 
         mAmWmState.computeState(new ComponentName[] {});
-        final boolean naturalOrientationPortrait =
+        boolean naturalOrientationPortrait =
                 mAmWmState.getWmState().getDisplay(DEFAULT_DISPLAY)
                         .mFullConfiguration.orientation == ORIENTATION_PORTRAIT;
 
-        final RotationSession rotationSession = createManagedRotationSession();
-        rotationSession.set(naturalOrientationPortrait ? ROTATION_90 : ROTATION_0);
+        try (final RotationSession rotationSession = new RotationSession()) {
+            rotationSession.set(naturalOrientationPortrait ? ROTATION_90 : ROTATION_0);
 
-        launchActivityInSplitScreenWithRecents(LAUNCHING_ACTIVITY);
-        final TestActivity activity = launchAndWait(mTestActivity);
-        mAmWmState.computeState(mTestActivityComponentName);
+            launchActivityInSplitScreenWithRecents(LAUNCHING_ACTIVITY);
+            final TestActivity activity = launchAndWait(mTestActivity);
+            mAmWmState.computeState(mTestActivityComponentName);
 
-        mAmWmState.assertContainsStack("Must contain fullscreen stack.",
-                WINDOWING_MODE_SPLIT_SCREEN_SECONDARY, ACTIVITY_TYPE_STANDARD);
-        mAmWmState.assertContainsStack("Must contain docked stack.",
-                WINDOWING_MODE_SPLIT_SCREEN_PRIMARY, ACTIVITY_TYPE_STANDARD);
+            mAmWmState.assertContainsStack("Must contain fullscreen stack.",
+                    WINDOWING_MODE_SPLIT_SCREEN_SECONDARY, ACTIVITY_TYPE_STANDARD);
+            mAmWmState.assertContainsStack("Must contain docked stack.",
+                    WINDOWING_MODE_SPLIT_SCREEN_PRIMARY, ACTIVITY_TYPE_STANDARD);
 
-        mAmWmState.computeState(LAUNCHING_ACTIVITY, mTestActivityComponentName);
+            mAmWmState.computeState(LAUNCHING_ACTIVITY, mTestActivityComponentName);
 
-        // Ensure that top insets are not consumed for LAYOUT_FULLSCREEN
-        WindowInsets insets = getOnMainSync(activity::getDispatchedInsets);
-        final WindowInsets rootInsets = getOnMainSync(activity::getRootInsets);
-        assertEquals("top inset must be dispatched in split screen",
-                rootInsets.getSystemWindowInsetTop(), insets.getSystemWindowInsetTop());
+            // Ensure that top insets are not consumed for LAYOUT_FULLSCREEN
+            WindowInsets insets = getOnMainSync(activity::getDispatchedInsets);
+            WindowInsets rootInsets = getOnMainSync(activity::getRootInsets);
+            assertEquals("top inset must be dispatched in split screen",
+                    rootInsets.getSystemWindowInsetTop(), insets.getSystemWindowInsetTop());
 
-        // Ensure that top insets are fully consumed for FULLSCREEN
-        final TestActivity fullscreenActivity = launchAndWait(mFullscreenTestActivity);
-        insets = getOnMainSync(fullscreenActivity::getDispatchedInsets);
-        assertEquals("top insets must be consumed if FULLSCREEN is set",
-                0, insets.getSystemWindowInsetTop());
+            // Ensure that top insets are fully consumed for FULLSCREEN
+            final TestActivity fullscreenActivity = launchAndWait(mFullscreenTestActivity);
+            insets = getOnMainSync(fullscreenActivity::getDispatchedInsets);
+            assertEquals("top insets must be consumed if FULLSCREEN is set",
+                    0, insets.getSystemWindowInsetTop());
 
-        // Ensure that top insets are fully consumed for FULLSCREEN when setting it over wm
-        // layout params
-        final TestActivity fullscreenWmFlagsActivity =
-                launchAndWait(mFullscreenWmFlagsTestActivity);
-        insets = getOnMainSync(fullscreenWmFlagsActivity::getDispatchedInsets);
-        assertEquals("top insets must be consumed if FULLSCREEN is set",
-                0, insets.getSystemWindowInsetTop());
+            // Ensure that top insets are fully consumed for FULLSCREEN when setting it over wm
+            // layout params
+            final TestActivity fullscreenWmFlagsActivity =
+                    launchAndWait(mFullscreenWmFlagsTestActivity);
+            insets = getOnMainSync(fullscreenWmFlagsActivity::getDispatchedInsets);
+            assertEquals("top insets must be consumed if FULLSCREEN is set",
+                    0, insets.getSystemWindowInsetTop());
+        }
     }
 
     private void commonAsserts(WindowInsets insets) {
diff --git a/tests/framework/base/windowmanager/src/android/server/wm/intent/IntentTestBase.java b/tests/framework/base/windowmanager/src/android/server/wm/intent/IntentTestBase.java
index 4256f9f..cdabec4 100644
--- a/tests/framework/base/windowmanager/src/android/server/wm/intent/IntentTestBase.java
+++ b/tests/framework/base/windowmanager/src/android/server/wm/intent/IntentTestBase.java
@@ -45,7 +45,9 @@
     }
 
     @After
+    @Override
     public void tearDown() throws Exception {
+        super.tearDown();
         cleanUp(activitiesUsedInTest());
     }
 
diff --git a/tests/framework/base/windowmanager/src/android/server/wm/lifecycle/ActivityLifecycleClientTestBase.java b/tests/framework/base/windowmanager/src/android/server/wm/lifecycle/ActivityLifecycleClientTestBase.java
index a615e8b..443addc 100644
--- a/tests/framework/base/windowmanager/src/android/server/wm/lifecycle/ActivityLifecycleClientTestBase.java
+++ b/tests/framework/base/windowmanager/src/android/server/wm/lifecycle/ActivityLifecycleClientTestBase.java
@@ -16,6 +16,7 @@
 
 package android.server.wm.lifecycle;
 
+import static android.app.Activity.RESULT_OK;
 import static android.content.Intent.FLAG_ACTIVITY_MULTIPLE_TASK;
 import static android.content.Intent.FLAG_ACTIVITY_NEW_TASK;
 import static android.server.wm.StateLogger.log;
@@ -49,9 +50,9 @@
 import android.content.pm.ActivityInfo;
 import android.content.res.Configuration;
 import android.os.Bundle;
+import android.os.Handler;
 import android.os.Looper;
 import android.server.wm.MultiDisplayTestBase;
-import android.server.wm.ObjectTracker;
 import android.server.wm.lifecycle.LifecycleLog.ActivityCallback;
 import android.transition.Transition;
 import android.transition.TransitionListenerAdapter;
@@ -120,7 +121,7 @@
     }
 
     /** Activity launch builder for lifecycle tests. */
-    class Launcher implements ObjectTracker.Consumable {
+    class Launcher {
         private int mFlags;
         private ActivityCallback mExpectedState;
         private List<String> mExtraFlags = new ArrayList<>();
@@ -137,7 +138,6 @@
          */
         Launcher(@NonNull Class<? extends Activity> activityClass) {
             mActivityClass = activityClass;
-            mObjectTracker.track(this);
         }
 
         /**
@@ -238,8 +238,11 @@
         }
 
         @Override
-        public boolean isConsumed() {
-            return mLaunchCalled;
+        protected void finalize() throws Throwable {
+            super.finalize();
+            if (!mLaunchCalled) {
+                throw new IllegalStateException("Activity launch builder created but not used!");
+            }
         }
     }
 
diff --git a/tests/framework/base/windowmanager/src/android/server/wm/lifecycle/ActivityLifecycleTests.java b/tests/framework/base/windowmanager/src/android/server/wm/lifecycle/ActivityLifecycleTests.java
index 20b8e01..9bffdde 100644
--- a/tests/framework/base/windowmanager/src/android/server/wm/lifecycle/ActivityLifecycleTests.java
+++ b/tests/framework/base/windowmanager/src/android/server/wm/lifecycle/ActivityLifecycleTests.java
@@ -497,50 +497,51 @@
                 state(becomingVisibleActivity, ON_STOP),
                 state(translucentActivity, ON_STOP));
 
-        final RotationSession rotationSession = createManagedRotationSession();
-        if (!supportsLockedUserRotation(
-                rotationSession, translucentActivity.getDisplay().getDisplayId())) {
-            return;
+        try (final RotationSession rotationSession = new RotationSession()) {
+            if (!supportsLockedUserRotation(
+                    rotationSession, translucentActivity.getDisplay().getDisplayId())) {
+                return;
+            }
+
+            getLifecycleLog().clear();
+
+            final int current = rotationSession.get();
+            // Set new rotation to cause a configuration change.
+            switch (current) {
+                case ROTATION_0:
+                case ROTATION_180:
+                    rotationSession.set(ROTATION_90);
+                    break;
+                case ROTATION_90:
+                case ROTATION_270:
+                    rotationSession.set(ROTATION_0);
+                    break;
+                default:
+                    fail("Unknown rotation:" + current);
+            }
+
+            // Assert that the top activity was relaunched.
+            waitAndAssertActivityStates(state(topOpaqueActivity, ON_RESUME));
+            LifecycleVerifier.assertRelaunchSequence(
+                    SecondActivity.class, getLifecycleLog(), ON_RESUME);
+
+            // Finish the top activity
+            getLifecycleLog().clear();
+            topOpaqueActivity.finish();
+
+            // Assert that the translucent activity and the activity visible behind it were
+            // relaunched.
+            waitAndAssertActivityStates(state(becomingVisibleActivity, ON_PAUSE),
+                    state(translucentActivity, ON_RESUME));
+
+            LifecycleVerifier.assertSequence(FirstActivity.class, getLifecycleLog(),
+                    Arrays.asList(ON_DESTROY, PRE_ON_CREATE, ON_CREATE, ON_START, ON_RESUME,
+                            ON_PAUSE), "becomingVisiblePaused");
+            final List<LifecycleLog.ActivityCallback> expectedSequence =
+                    Arrays.asList(ON_DESTROY, PRE_ON_CREATE, ON_CREATE, ON_START, ON_RESUME);
+            LifecycleVerifier.assertSequence(TranslucentActivity.class, getLifecycleLog(),
+                    expectedSequence, "becomingVisibleResumed");
         }
-
-        getLifecycleLog().clear();
-
-        final int current = rotationSession.get();
-        // Set new rotation to cause a configuration change.
-        switch (current) {
-            case ROTATION_0:
-            case ROTATION_180:
-                rotationSession.set(ROTATION_90);
-                break;
-            case ROTATION_90:
-            case ROTATION_270:
-                rotationSession.set(ROTATION_0);
-                break;
-            default:
-                fail("Unknown rotation:" + current);
-        }
-
-        // Assert that the top activity was relaunched.
-        waitAndAssertActivityStates(state(topOpaqueActivity, ON_RESUME));
-        LifecycleVerifier.assertRelaunchSequence(
-                SecondActivity.class, getLifecycleLog(), ON_RESUME);
-
-        // Finish the top activity
-        getLifecycleLog().clear();
-        topOpaqueActivity.finish();
-
-        // Assert that the translucent activity and the activity visible behind it were
-        // relaunched.
-        waitAndAssertActivityStates(state(becomingVisibleActivity, ON_PAUSE),
-                state(translucentActivity, ON_RESUME));
-
-        LifecycleVerifier.assertSequence(FirstActivity.class, getLifecycleLog(),
-                Arrays.asList(ON_DESTROY, PRE_ON_CREATE, ON_CREATE, ON_START, ON_RESUME,
-                        ON_PAUSE), "becomingVisiblePaused");
-        final List<LifecycleLog.ActivityCallback> expectedSequence =
-                Arrays.asList(ON_DESTROY, PRE_ON_CREATE, ON_CREATE, ON_START, ON_RESUME);
-        LifecycleVerifier.assertSequence(TranslucentActivity.class, getLifecycleLog(),
-                expectedSequence, "becomingVisibleResumed");
     }
 
     @Test
diff --git a/tests/framework/base/windowmanager/src/android/server/wm/lifecycle/ActivityLifecycleTopResumedStateTests.java b/tests/framework/base/windowmanager/src/android/server/wm/lifecycle/ActivityLifecycleTopResumedStateTests.java
index 2c3985e..e29b7eb 100644
--- a/tests/framework/base/windowmanager/src/android/server/wm/lifecycle/ActivityLifecycleTopResumedStateTests.java
+++ b/tests/framework/base/windowmanager/src/android/server/wm/lifecycle/ActivityLifecycleTopResumedStateTests.java
@@ -48,7 +48,6 @@
 import android.app.ActivityOptions;
 import android.content.ComponentName;
 import android.content.Intent;
-import android.os.SystemClock;
 import android.platform.test.annotations.Presubmit;
 import android.server.wm.ActivityManagerState;
 import android.server.wm.ActivityManagerState.ActivityStack;
@@ -794,186 +793,189 @@
         waitAndAssertTopResumedActivity(getComponentName(CallbackTrackingActivity.class),
                 DEFAULT_DISPLAY, "Activity launched on default display must be focused");
 
-        // Create new simulated display
-        final ActivityManagerState.ActivityDisplay newDisplay = createManagedVirtualDisplaySession()
-                .setSimulateDisplay(true)
-                .createDisplay();
+        try (final VirtualDisplaySession virtualDisplaySession = new VirtualDisplaySession()) {
+            // Create new simulated display
+            final ActivityManagerState.ActivityDisplay newDisplay
+                    = virtualDisplaySession.setSimulateDisplay(true).createDisplay();
 
-        // Launch another activity on new secondary display.
-        getLifecycleLog().clear();
-        launchOptions.setLaunchDisplayId(newDisplay.mId);
-        new Launcher(SingleTopActivity.class)
-                .setFlags(FLAG_ACTIVITY_NEW_TASK)
-                .setOptions(launchOptions)
-                .launch();
-        waitAndAssertTopResumedActivity(getComponentName(SingleTopActivity.class),
-                newDisplay.mId, "Activity launched on secondary display must be focused");
+            // Launch another activity on new secondary display.
+            getLifecycleLog().clear();
+            launchOptions.setLaunchDisplayId(newDisplay.mId);
+            new Launcher(SingleTopActivity.class)
+                    .setFlags(FLAG_ACTIVITY_NEW_TASK)
+                    .setOptions(launchOptions)
+                    .launch();
+            waitAndAssertTopResumedActivity(getComponentName(SingleTopActivity.class),
+                    newDisplay.mId, "Activity launched on secondary display must be focused");
 
-        getLifecycleLog().clear();
+            getLifecycleLog().clear();
 
-        // Tap on default display to switch the top activity
-        tapOnDisplayCenter(DEFAULT_DISPLAY);
+            // Tap on default display to switch the top activity
+            tapOnDisplayCenter(DEFAULT_DISPLAY);
 
-        // Wait and assert focus switch
-        waitAndAssertActivityTransitions(SingleTopActivity.class,
-                Arrays.asList(ON_TOP_POSITION_LOST), "tapOnFocusSwitch");
-        waitAndAssertActivityTransitions(CallbackTrackingActivity.class,
-                Arrays.asList(ON_TOP_POSITION_GAINED), "tapOnFocusSwitch");
-        LifecycleVerifier.assertEntireSequence(Arrays.asList(
-                transition(SingleTopActivity.class, ON_TOP_POSITION_LOST),
-                transition(CallbackTrackingActivity.class, ON_TOP_POSITION_GAINED)),
-                getLifecycleLog(), "Top activity must be switched on tap");
+            // Wait and assert focus switch
+            waitAndAssertActivityTransitions(SingleTopActivity.class,
+                    Arrays.asList(ON_TOP_POSITION_LOST), "tapOnFocusSwitch");
+            waitAndAssertActivityTransitions(CallbackTrackingActivity.class,
+                    Arrays.asList(ON_TOP_POSITION_GAINED), "tapOnFocusSwitch");
+            LifecycleVerifier.assertEntireSequence(Arrays.asList(
+                    transition(SingleTopActivity.class, ON_TOP_POSITION_LOST),
+                    transition(CallbackTrackingActivity.class, ON_TOP_POSITION_GAINED)),
+                    getLifecycleLog(), "Top activity must be switched on tap");
 
-        getLifecycleLog().clear();
+            getLifecycleLog().clear();
 
-        // Tap on new display to switch the top activity
-        tapOnDisplayCenter(newDisplay.mId);
+            // Tap on new display to switch the top activity
+            tapOnDisplayCenter(newDisplay.mId);
 
-        // Wait and assert focus switch
-        waitAndAssertActivityTransitions(CallbackTrackingActivity.class,
-                Arrays.asList(ON_TOP_POSITION_LOST), "tapOnFocusSwitch");
-        waitAndAssertActivityTransitions(SingleTopActivity.class,
-                Arrays.asList(ON_TOP_POSITION_GAINED), "tapOnFocusSwitch");
-        LifecycleVerifier.assertEntireSequence(Arrays.asList(
-                transition(CallbackTrackingActivity.class, ON_TOP_POSITION_LOST),
-                transition(SingleTopActivity.class, ON_TOP_POSITION_GAINED)),
-                getLifecycleLog(), "Top activity must be switched on tap");
+            // Wait and assert focus switch
+            waitAndAssertActivityTransitions(CallbackTrackingActivity.class,
+                    Arrays.asList(ON_TOP_POSITION_LOST), "tapOnFocusSwitch");
+            waitAndAssertActivityTransitions(SingleTopActivity.class,
+                    Arrays.asList(ON_TOP_POSITION_GAINED), "tapOnFocusSwitch");
+            LifecycleVerifier.assertEntireSequence(Arrays.asList(
+                    transition(CallbackTrackingActivity.class, ON_TOP_POSITION_LOST),
+                    transition(SingleTopActivity.class, ON_TOP_POSITION_GAINED)),
+                    getLifecycleLog(), "Top activity must be switched on tap");
+        }
     }
 
     @Test
-    public void testTopPositionSwitchAcrossDisplaysOnTapSlowDifferentProcess() {
+    public void testTopPositionSwitchAcrossDisplaysOnTapSlowDifferentProcess() throws Exception {
         assumeTrue(supportsMultiDisplay());
 
-        // Create new simulated display.
-        final ActivityManagerState.ActivityDisplay newDisplay = createManagedVirtualDisplaySession()
-                .setSimulateDisplay(true)
-                .createDisplay();
+        try (final VirtualDisplaySession virtualDisplaySession = new VirtualDisplaySession()) {
+            // Create new simulated display
+            final ActivityManagerState.ActivityDisplay newDisplay
+                    = virtualDisplaySession.setSimulateDisplay(true).createDisplay();
 
-        // Launch an activity on new secondary display.
-        final Class<? extends Activity> secondActivityClass =
-                SecondProcessCallbackTrackingActivity.class;
-        final ComponentName secondActivityComponent =
-                new ComponentName(mTargetContext, secondActivityClass);
+            // Launch an activity on new secondary display.
+            final Class<? extends Activity> secondActivityClass =
+                    SecondProcessCallbackTrackingActivity.class;
+            final ComponentName secondActivityComponent =
+                    new ComponentName(mTargetContext, secondActivityClass);
 
-        getLaunchActivityBuilder()
-                .setTargetActivity(secondActivityComponent)
-                .setUseInstrumentation()
-                .setDisplayId(newDisplay.mId)
-                .execute();
-        waitAndAssertActivityStates(state(secondActivityClass, ON_TOP_POSITION_GAINED));
+            getLaunchActivityBuilder()
+                    .setTargetActivity(secondActivityComponent)
+                    .setUseInstrumentation()
+                    .setDisplayId(newDisplay.mId)
+                    .execute();
+            waitAndAssertActivityStates(state(secondActivityClass, ON_TOP_POSITION_GAINED));
 
-        // Launch activity on default display, which will be slow to release top position.
-        getLifecycleLog().clear();
-        final ActivityOptions launchOptions = ActivityOptions.makeBasic();
-        launchOptions.setLaunchDisplayId(DEFAULT_DISPLAY);
-        final Class<? extends Activity> defaultActivityClass = SlowActivity.class;
-        final Intent defaultDisplaySlowIntent = new Intent(mContext, defaultActivityClass);
-        defaultDisplaySlowIntent.setFlags(FLAG_ACTIVITY_NEW_TASK | FLAG_ACTIVITY_MULTIPLE_TASK);
-        defaultDisplaySlowIntent.putExtra(SlowActivity.EXTRA_CONTROL_FLAGS,
-                SlowActivity.FLAG_SLOW_TOP_RESUME_RELEASE);
-        mTargetContext.startActivity(defaultDisplaySlowIntent, launchOptions.toBundle());
+            // Launch activity on default display, which will be slow to release top position
+            getLifecycleLog().clear();
+            final ActivityOptions launchOptions = ActivityOptions.makeBasic();
+            launchOptions.setLaunchDisplayId(DEFAULT_DISPLAY);
+            final Class defaultActivityClass = SlowActivity.class;
+            final Intent defaultDisplaySlowIntent = new Intent(mContext, defaultActivityClass);
+            defaultDisplaySlowIntent.setFlags(FLAG_ACTIVITY_NEW_TASK | FLAG_ACTIVITY_MULTIPLE_TASK);
+            defaultDisplaySlowIntent.putExtra(SlowActivity.EXTRA_CONTROL_FLAGS,
+                    SlowActivity.FLAG_SLOW_TOP_RESUME_RELEASE);
+            mTargetContext.startActivity(defaultDisplaySlowIntent, launchOptions.toBundle());
 
-        waitAndAssertTopResumedActivity(getComponentName(SlowActivity.class),
-                DEFAULT_DISPLAY, "Activity launched on default display must be focused");
+            waitAndAssertTopResumedActivity(getComponentName(SlowActivity.class),
+                    DEFAULT_DISPLAY, "Activity launched on default display must be focused");
 
-        // Wait and assert focus switch
-        waitAndAssertActivityStates(state(secondActivityClass, ON_TOP_POSITION_LOST),
-                state(defaultActivityClass, ON_TOP_POSITION_GAINED));
-        LifecycleVerifier.assertOrder(getLifecycleLog(), Arrays.asList(
-                transition(secondActivityClass, ON_TOP_POSITION_LOST),
-                transition(defaultActivityClass, ON_TOP_POSITION_GAINED)),
-                "launchOnDifferentDisplay");
+            // Wait and assert focus switch
+            waitAndAssertActivityStates(state(secondActivityClass, ON_TOP_POSITION_LOST),
+                    state(defaultActivityClass, ON_TOP_POSITION_GAINED));
+            LifecycleVerifier.assertOrder(getLifecycleLog(), Arrays.asList(
+                    transition(secondActivityClass, ON_TOP_POSITION_LOST),
+                    transition(defaultActivityClass, ON_TOP_POSITION_GAINED)),
+                    "launchOnDifferentDisplay");
 
-        // Tap on secondary display to switch the top activity.
-        getLifecycleLog().clear();
-        tapOnDisplayCenter(newDisplay.mId);
+            // Tap on secondary display to switch the top activity
+            getLifecycleLog().clear();
+            tapOnDisplayCenter(newDisplay.mId);
 
-        // Wait and assert top resumed position switch.
-        waitAndAssertActivityStates(state(secondActivityClass, ON_TOP_POSITION_GAINED),
-                state(defaultActivityClass, ON_TOP_POSITION_LOST));
-        LifecycleVerifier.assertOrder(getLifecycleLog(), Arrays.asList(
-                transition(defaultActivityClass, ON_TOP_POSITION_LOST),
-                transition(secondActivityClass, ON_TOP_POSITION_GAINED)),
-                "tapOnDifferentDisplay");
+            // Wait and assert top resumed position switch
+            waitAndAssertActivityStates(state(secondActivityClass, ON_TOP_POSITION_GAINED),
+                    state(defaultActivityClass, ON_TOP_POSITION_LOST));
+            LifecycleVerifier.assertOrder(getLifecycleLog(), Arrays.asList(
+                    transition(defaultActivityClass, ON_TOP_POSITION_LOST),
+                    transition(secondActivityClass, ON_TOP_POSITION_GAINED)),
+                    "tapOnDifferentDisplay");
 
-        // Tap on default display to switch the top activity again.
-        getLifecycleLog().clear();
-        tapOnDisplayCenter(DEFAULT_DISPLAY);
+            // Tap on default display to switch the top activity again
+            getLifecycleLog().clear();
+            tapOnDisplayCenter(DEFAULT_DISPLAY);
 
-        // Wait and assert top resumed position switch.
-        waitAndAssertActivityStates(state(secondActivityClass, ON_TOP_POSITION_LOST),
-                state(defaultActivityClass, ON_TOP_POSITION_GAINED));
-        LifecycleVerifier.assertOrder(getLifecycleLog(), Arrays.asList(
-                transition(secondActivityClass, ON_TOP_POSITION_LOST),
-                transition(defaultActivityClass, ON_TOP_POSITION_GAINED)),
-                "tapOnDifferentDisplay");
+            // Wait and assert top resumed position switch
+            waitAndAssertActivityStates(state(secondActivityClass, ON_TOP_POSITION_LOST),
+                    state(defaultActivityClass, ON_TOP_POSITION_GAINED));
+            LifecycleVerifier.assertOrder(getLifecycleLog(), Arrays.asList(
+                    transition(secondActivityClass, ON_TOP_POSITION_LOST),
+                    transition(defaultActivityClass, ON_TOP_POSITION_GAINED)),
+                    "tapOnDifferentDisplay");
+        }
     }
 
     @Test
-    public void testTopPositionSwitchAcrossDisplaysOnTapTimeoutDifferentProcess() {
+    public void testTopPositionSwitchAcrossDisplaysOnTapTimeoutDifferentProcess() throws Exception {
         assumeTrue(supportsMultiDisplay());
 
-        // Create new simulated display.
-        final ActivityManagerState.ActivityDisplay newDisplay = createManagedVirtualDisplaySession()
-                .setSimulateDisplay(true)
-                .createDisplay();
+        try (final VirtualDisplaySession virtualDisplaySession = new VirtualDisplaySession()) {
+            // Create new simulated display
+            final ActivityManagerState.ActivityDisplay newDisplay
+                    = virtualDisplaySession.setSimulateDisplay(true).createDisplay();
 
-        // Launch an activity on new secondary display.
-        final Class<? extends Activity> secondActivityClass =
-                SecondProcessCallbackTrackingActivity.class;
-        final ComponentName secondActivityComponent =
-                new ComponentName(mTargetContext, secondActivityClass);
+            // Launch an activity on new secondary display.
+            final Class<? extends Activity> secondActivityClass =
+                    SecondProcessCallbackTrackingActivity.class;
+            final ComponentName secondActivityComponent =
+                    new ComponentName(mTargetContext, secondActivityClass);
 
-        getLaunchActivityBuilder()
-                .setTargetActivity(secondActivityComponent)
-                .setUseInstrumentation()
-                .setDisplayId(newDisplay.mId)
-                .execute();
-        waitAndAssertActivityStates(state(secondActivityClass, ON_TOP_POSITION_GAINED));
+            getLaunchActivityBuilder()
+                    .setTargetActivity(secondActivityComponent)
+                    .setUseInstrumentation()
+                    .setDisplayId(newDisplay.mId)
+                    .execute();
+            waitAndAssertActivityStates(state(secondActivityClass, ON_TOP_POSITION_GAINED));
 
-        // Launch activity on default display, which will be slow to release top position.
-        getLifecycleLog().clear();
-        final ActivityOptions launchOptions = ActivityOptions.makeBasic();
-        launchOptions.setLaunchDisplayId(DEFAULT_DISPLAY);
-        final Class<? extends Activity> defaultActivityClass = SlowActivity.class;
-        final Intent defaultDisplaySlowIntent = new Intent(mContext, defaultActivityClass);
-        defaultDisplaySlowIntent.setFlags(FLAG_ACTIVITY_NEW_TASK | FLAG_ACTIVITY_MULTIPLE_TASK);
-        defaultDisplaySlowIntent.putExtra(SlowActivity.EXTRA_CONTROL_FLAGS,
-                SlowActivity.FLAG_TIMEOUT_TOP_RESUME_RELEASE);
-        mTargetContext.startActivity(defaultDisplaySlowIntent, launchOptions.toBundle());
+            // Launch activity on default display, which will be slow to release top position
+            getLifecycleLog().clear();
+            final ActivityOptions launchOptions = ActivityOptions.makeBasic();
+            launchOptions.setLaunchDisplayId(DEFAULT_DISPLAY);
+            final Class defaultActivityClass = SlowActivity.class;
+            final Intent defaultDisplaySlowIntent = new Intent(mContext, defaultActivityClass);
+            defaultDisplaySlowIntent.setFlags(FLAG_ACTIVITY_NEW_TASK | FLAG_ACTIVITY_MULTIPLE_TASK);
+            defaultDisplaySlowIntent.putExtra(SlowActivity.EXTRA_CONTROL_FLAGS,
+                    SlowActivity.FLAG_TIMEOUT_TOP_RESUME_RELEASE);
+            mTargetContext.startActivity(defaultDisplaySlowIntent, launchOptions.toBundle());
 
-        waitAndAssertTopResumedActivity(getComponentName(SlowActivity.class),
-                DEFAULT_DISPLAY, "Activity launched on default display must be focused");
+            waitAndAssertTopResumedActivity(getComponentName(SlowActivity.class),
+                    DEFAULT_DISPLAY, "Activity launched on default display must be focused");
 
-        // Wait and assert focus switch.
-        waitAndAssertActivityStates(state(secondActivityClass, ON_TOP_POSITION_LOST),
-                state(defaultActivityClass, ON_TOP_POSITION_GAINED));
-        LifecycleVerifier.assertOrder(getLifecycleLog(), Arrays.asList(
-                transition(secondActivityClass, ON_TOP_POSITION_LOST),
-                transition(defaultActivityClass, ON_TOP_POSITION_GAINED)),
-                "launchOnDifferentDisplay");
+            // Wait and assert focus switch
+            waitAndAssertActivityStates(state(secondActivityClass, ON_TOP_POSITION_LOST),
+                    state(defaultActivityClass, ON_TOP_POSITION_GAINED));
+            LifecycleVerifier.assertOrder(getLifecycleLog(), Arrays.asList(
+                    transition(secondActivityClass, ON_TOP_POSITION_LOST),
+                    transition(defaultActivityClass, ON_TOP_POSITION_GAINED)),
+                    "launchOnDifferentDisplay");
 
-        // Tap on secondary display to switch the top activity.
-        getLifecycleLog().clear();
-        tapOnDisplayCenter(newDisplay.mId);
+            // Tap on secondary display to switch the top activity
+            getLifecycleLog().clear();
+            tapOnDisplayCenter(newDisplay.mId);
 
-        // Wait and assert top resumed position switch. Because of timeout top position gain
-        // will appear before top position loss handling is finished.
-        waitAndAssertActivityStates(state(secondActivityClass, ON_TOP_POSITION_GAINED),
-                state(defaultActivityClass, ON_TOP_POSITION_LOST));
-        LifecycleVerifier.assertOrder(getLifecycleLog(), Arrays.asList(
-                transition(secondActivityClass, ON_TOP_POSITION_GAINED),
-                transition(defaultActivityClass, ON_TOP_POSITION_LOST)),
-                "tapOnDifferentDisplay");
+            // Wait and assert top resumed position switch. Because of timeout top position gain
+            // will appear before top position loss handling is finished.
+            waitAndAssertActivityStates(state(secondActivityClass, ON_TOP_POSITION_GAINED),
+                    state(defaultActivityClass, ON_TOP_POSITION_LOST));
+            LifecycleVerifier.assertOrder(getLifecycleLog(), Arrays.asList(
+                    transition(secondActivityClass, ON_TOP_POSITION_GAINED),
+                    transition(defaultActivityClass, ON_TOP_POSITION_LOST)),
+                    "tapOnDifferentDisplay");
 
-        // Wait 5 seconds more to make sure that no new messages received after top resumed state
-        // released by the slow activity
-        getLifecycleLog().clear();
-        SystemClock.sleep(5000);
-        LifecycleVerifier.assertEmptySequence(defaultActivityClass, getLifecycleLog(),
-                "topStateLossTimeout");
-        LifecycleVerifier.assertEmptySequence(secondActivityClass, getLifecycleLog(),
-                "topStateLossTimeout");
+            // Wait 5 seconds more to make sure that no new messages received after top resumed state
+            // released by the slow activity
+            getLifecycleLog().clear();
+            Thread.sleep(5000);
+            LifecycleVerifier.assertEmptySequence(defaultActivityClass, getLifecycleLog(),
+                    "topStateLossTimeout");
+            LifecycleVerifier.assertEmptySequence(secondActivityClass, getLifecycleLog(),
+                    "topStateLossTimeout");
+        }
     }
 
     @FlakyTest(bugId=137329632)
@@ -988,36 +990,37 @@
         waitAndAssertTopResumedActivity(getComponentName(CallbackTrackingActivity.class),
                 DEFAULT_DISPLAY, "Activity launched on default display must be focused");
 
-        // Create new simulated display.
-        final ActivityManagerState.ActivityDisplay newDisplay = createManagedVirtualDisplaySession()
-                .setSimulateDisplay(true)
-                .createDisplay();
+        try (final VirtualDisplaySession virtualDisplaySession = new VirtualDisplaySession()) {
+            // Create new simulated display
+            final ActivityManagerState.ActivityDisplay newDisplay
+                    = virtualDisplaySession.setSimulateDisplay(true).createDisplay();
 
-        // Launch another activity on new secondary display.
-        getLifecycleLog().clear();
-        final ActivityOptions launchOptions = ActivityOptions.makeBasic();
-        launchOptions.setLaunchDisplayId(newDisplay.mId);
-        new Launcher(SingleTopActivity.class)
-                .setFlags(FLAG_ACTIVITY_NEW_TASK)
-                .setOptions(launchOptions)
-                .launch();
-        waitAndAssertTopResumedActivity(getComponentName(SingleTopActivity.class),
-                newDisplay.mId, "Activity launched on secondary display must be focused");
-        // An activity is launched on the new display, so the activity on default display should
-        // lose the top state.
-        LifecycleVerifier.assertSequence(CallbackTrackingActivity.class, getLifecycleLog(),
-                Arrays.asList(ON_TOP_POSITION_LOST), "launchFocusSwitch");
+            // Launch another activity on new secondary display.
+            getLifecycleLog().clear();
+            final ActivityOptions launchOptions = ActivityOptions.makeBasic();
+            launchOptions.setLaunchDisplayId(newDisplay.mId);
+            new Launcher(SingleTopActivity.class)
+                    .setFlags(FLAG_ACTIVITY_NEW_TASK)
+                    .setOptions(launchOptions)
+                    .launch();
+            waitAndAssertTopResumedActivity(getComponentName(SingleTopActivity.class),
+                    newDisplay.mId, "Activity launched on secondary display must be focused");
+            // An activity is launched on the new display, so the activity on default display should
+            // lose the top state.
+            LifecycleVerifier.assertSequence(CallbackTrackingActivity.class, getLifecycleLog(),
+                    Arrays.asList(ON_TOP_POSITION_LOST), "launchFocusSwitch");
 
-        // Finish the activity on the default display.
-        getLifecycleLog().clear();
-        callbackTrackingActivity.finish();
+            // Finish the activity on the default display
+            getLifecycleLog().clear();
+            callbackTrackingActivity.finish();
 
-        // Verify that activity was actually destroyed.
-        waitAndAssertActivityStates(state(CallbackTrackingActivity.class, ON_DESTROY));
-        // Verify that the original focused display is not affected by the finished activity on
-        // non-focused display.
-        LifecycleVerifier.assertEmptySequence(SingleTopActivity.class, getLifecycleLog(),
-                "destructionOnDifferentDisplay");
+            // Verify that activity was actually destroyed
+            waitAndAssertActivityStates(state(CallbackTrackingActivity.class, ON_DESTROY));
+            // Verify that the original focused display is not affected by the finished activity on
+            // non-focused display.
+            LifecycleVerifier.assertEmptySequence(SingleTopActivity.class, getLifecycleLog(),
+                    "destructionOnDifferentDisplay");
+        }
     }
 
     @FlakyTest(bugId=137329632)
@@ -1033,38 +1036,40 @@
         waitAndAssertTopResumedActivity(getComponentName(CallbackTrackingActivity.class),
                 DEFAULT_DISPLAY, "Activity launched on default display must be focused");
 
-        // Create new simulated display.
-        final ActivityManagerState.ActivityDisplay newDisplay = createManagedVirtualDisplaySession()
-                .setSimulateDisplay(true)
-                .createDisplay();
+        try (final VirtualDisplaySession virtualDisplaySession = new VirtualDisplaySession()) {
+            // Create new simulated display
+            final ActivityManagerState.ActivityDisplay newDisplay
+                    = virtualDisplaySession.setSimulateDisplay(true).createDisplay();
 
-        // Launch another activity on new secondary display.
-        getLifecycleLog().clear();
-        final ActivityOptions launchOptions = ActivityOptions.makeBasic();
-        launchOptions.setLaunchDisplayId(newDisplay.mId);
-        new Launcher(SingleTopActivity.class)
-                .setFlags(FLAG_ACTIVITY_NEW_TASK)
-                .setOptions(launchOptions)
-                .launch();
-        waitAndAssertTopResumedActivity(getComponentName(SingleTopActivity.class),
-                newDisplay.mId, "Activity launched on secondary display must be focused");
+            // Launch another activity on new secondary display.
+            getLifecycleLog().clear();
+            final ActivityOptions launchOptions = ActivityOptions.makeBasic();
+            launchOptions.setLaunchDisplayId(newDisplay.mId);
+            new Launcher(SingleTopActivity.class)
+                    .setFlags(FLAG_ACTIVITY_NEW_TASK)
+                    .setOptions(launchOptions)
+                    .launch();
+            waitAndAssertTopResumedActivity(getComponentName(SingleTopActivity.class),
+                    newDisplay.mId, "Activity launched on secondary display must be focused");
 
-        // Bring the focus back.
-        final Intent sameInstanceIntent = new Intent(mContext, CallbackTrackingActivity.class);
-        sameInstanceIntent.setFlags(FLAG_ACTIVITY_REORDER_TO_FRONT);
-        bottomActivity.startActivity(sameInstanceIntent, null);
-        waitAndAssertActivityStates(state(CallbackTrackingActivity.class, ON_TOP_POSITION_GAINED));
+            // Bring the focus back
+            final Intent sameInstanceIntent = new Intent(mContext, CallbackTrackingActivity.class);
+            sameInstanceIntent.setFlags(FLAG_ACTIVITY_REORDER_TO_FRONT);
+            bottomActivity.startActivity(sameInstanceIntent, null);
+            waitAndAssertActivityStates(
+                    state(CallbackTrackingActivity.class, ON_TOP_POSITION_GAINED));
 
-        // Finish the focused activity
-        getLifecycleLog().clear();
-        callbackTrackingActivity.finish();
+            // Finish the focused activity
+            getLifecycleLog().clear();
+            callbackTrackingActivity.finish();
 
-        // Verify that lifecycle of the activity on a different display did not change.
-        // Top resumed state will be given to home activity on that display.
-        waitAndAssertActivityStates(state(CallbackTrackingActivity.class, ON_DESTROY),
-                state(SecondActivity.class, ON_RESUME));
-        LifecycleVerifier.assertEmptySequence(SingleTopActivity.class, getLifecycleLog(),
-                "destructionOnDifferentDisplay");
+            // Verify that lifecycle of the activity on a different display did not change.
+            // Top resumed state will be given to home activity on that display.
+            waitAndAssertActivityStates(state(CallbackTrackingActivity.class, ON_DESTROY),
+                    state(SecondActivity.class, ON_RESUME));
+            LifecycleVerifier.assertEmptySequence(SingleTopActivity.class, getLifecycleLog(),
+                    "destructionOnDifferentDisplay");
+        }
     }
 
     @Test
diff --git a/tests/framework/base/windowmanager/util/src/android/server/wm/ActivityAndWindowManagersState.java b/tests/framework/base/windowmanager/util/src/android/server/wm/ActivityAndWindowManagersState.java
index 8076b7a..512d860 100644
--- a/tests/framework/base/windowmanager/util/src/android/server/wm/ActivityAndWindowManagersState.java
+++ b/tests/framework/base/windowmanager/util/src/android/server/wm/ActivityAndWindowManagersState.java
@@ -31,6 +31,7 @@
 import static android.util.DisplayMetrics.DENSITY_DEFAULT;
 import static android.view.Display.DEFAULT_DISPLAY;
 
+import static com.google.common.truth.Truth.assertThat;
 import static com.google.common.truth.Truth.assertWithMessage;
 
 import static org.hamcrest.Matchers.greaterThan;
@@ -48,19 +49,20 @@
 
 import android.content.ComponentName;
 import android.graphics.Rect;
+import android.os.SystemClock;
 import android.server.wm.ActivityManagerState.ActivityStack;
 import android.server.wm.ActivityManagerState.ActivityTask;
 import android.server.wm.WindowManagerState.Display;
 import android.server.wm.WindowManagerState.WindowStack;
 import android.server.wm.WindowManagerState.WindowState;
 import android.server.wm.WindowManagerState.WindowTask;
+import android.util.Log;
 import android.util.SparseArray;
 
 import java.util.Arrays;
 import java.util.List;
 import java.util.Objects;
 import java.util.function.BiPredicate;
-import java.util.function.Consumer;
 import java.util.function.Predicate;
 import java.util.stream.Collectors;
 
@@ -286,6 +288,11 @@
                 "activity window to be gone");
     }
 
+    @Deprecated
+    void waitForFocusedStack(int stackId) {
+        waitForWithAmState(state -> state.getFocusedStackId() == stackId, "focused stack");
+    }
+
     void waitForFocusedStack(int windowingMode, int activityType) {
         waitForWithAmState(state ->
                         (activityType == ACTIVITY_TYPE_UNDEFINED
@@ -600,9 +607,7 @@
 
     /** Asserts that each display has correct resumed activity. */
     public void assertResumedActivities(final String msg,
-            Consumer<SparseArray<ComponentName>> resumedActivitiesMapping) {
-        final SparseArray<ComponentName> resumedActivities = new SparseArray<>();
-        resumedActivitiesMapping.accept(resumedActivities);
+            SparseArray<ComponentName> resumedActivities) {
         for (int i = 0; i < resumedActivities.size(); i++) {
             final int displayId = resumedActivities.keyAt(i);
             final ComponentName activityComponent = resumedActivities.valueAt(i);
@@ -930,11 +935,11 @@
         }
     }
 
-    public void assertActivityDisplayed(final ComponentName activityName) {
+    public void assertActivityDisplayed(final ComponentName activityName) throws Exception {
         assertWindowDisplayed(getWindowName(activityName));
     }
 
-    public void assertWindowDisplayed(final String windowName) {
+    public void assertWindowDisplayed(final String windowName) throws Exception {
         waitForValidState(WaitForValidActivityState.forWindow(windowName));
         assertTrue(windowName + "is visible", getWmState().isWindowVisible(windowName));
     }
diff --git a/tests/framework/base/windowmanager/util/src/android/server/wm/ActivityManagerTestBase.java b/tests/framework/base/windowmanager/util/src/android/server/wm/ActivityManagerTestBase.java
index 98bbcf2..5468226 100644
--- a/tests/framework/base/windowmanager/util/src/android/server/wm/ActivityManagerTestBase.java
+++ b/tests/framework/base/windowmanager/util/src/android/server/wm/ActivityManagerTestBase.java
@@ -142,13 +142,13 @@
 
 import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
+import androidx.test.rule.ActivityTestRule;
 
 import com.android.compatibility.common.util.SystemUtil;
 
+import org.junit.After;
 import org.junit.Before;
 import org.junit.Rule;
-import org.junit.rules.ErrorCollector;
-import org.junit.rules.RuleChain;
 import org.junit.rules.TestRule;
 import org.junit.runner.Description;
 import org.junit.runners.model.Statement;
@@ -162,6 +162,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.UUID;
+import java.util.concurrent.Callable;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.function.BooleanSupplier;
 import java.util.function.Consumer;
@@ -213,16 +214,18 @@
     protected ActivityManager mAm;
     protected ActivityTaskManager mAtm;
 
-    /** The tracker to manage objects (especially {@link AutoCloseable}) in a test method. */
-    protected final ObjectTracker mObjectTracker = new ObjectTracker();
+    /**
+     * Callable to clear launch params for all test packages.
+     */
+    private final Callable<Void> mClearLaunchParamsCallable = () -> {
+        mAtm.clearLaunchParamsForPackages(TEST_PACKAGES);
+        return null;
+    };
 
-    /** The last rule to handle all errors. */
-    private final ErrorCollector mPostAssertionRule = new PostAssertionRule();
-
-    /** The necessary procedures of set up and tear down. */
     @Rule
-    public final TestRule mBaseRule = RuleChain.outerRule(mPostAssertionRule)
-            .around(new WrapperRule(this::setUpBase, this::tearDownBase));
+    public final ActivityTestRule<SideActivity> mSideActivityRule =
+            new ActivityTestRule<>(SideActivity.class, true /* initialTouchMode */,
+                    false /* launchActivity */);
 
     /**
      * @return the am command to start the given activity with the following extra key/value pairs.
@@ -390,8 +393,7 @@
                         .addMonitor((String) null, null, false);
                 mContext.startActivity(intent.addFlags(FLAG_ACTIVITY_NEW_TASK), bundle);
                 // Wait for activity launch with timeout.
-                mTestActivity = (T) getInstrumentation().waitForMonitorWithTimeout(monitor,
-                        ACTIVITY_LAUNCH_TIMEOUT);
+                mTestActivity = (T) monitor.waitForActivityWithTimeout(ACTIVITY_LAUNCH_TIMEOUT);
                 assertNotNull(mTestActivity);
                 // Check activity is launched and resumed.
                 final ComponentName testActivityName = mTestActivity.getComponentName();
@@ -434,7 +436,7 @@
         }
 
         @Override
-        public void close() {
+        public void close() throws Exception {
             if (mTestActivity != null && mFinishAfterClose) {
                 mTestActivity.finishAndRemoveTask();
             }
@@ -443,27 +445,21 @@
 
     @Before
     public void setUp() throws Exception {
+        mContext = getInstrumentation().getContext();
+        mAm = mContext.getSystemService(ActivityManager.class);
+        mAtm = mContext.getSystemService(ActivityTaskManager.class);
+
         pressWakeupButton();
         pressUnlockButton();
         launchHomeActivityNoWait();
         removeStacksWithActivityTypes(ALL_ACTIVITY_TYPE_BUT_HOME);
 
         // Clear launch params for all test packages to make sure each test is run in a clean state.
-        SystemUtil.runWithShellPermissionIdentity(
-                () -> mAtm.clearLaunchParamsForPackages(TEST_PACKAGES));
+        SystemUtil.callWithShellPermissionIdentity(mClearLaunchParamsCallable);
     }
 
-    /** It always executes before {@link Before}. */
-    private void setUpBase() {
-        mContext = getInstrumentation().getContext();
-        mAm = mContext.getSystemService(ActivityManager.class);
-        mAtm = mContext.getSystemService(ActivityTaskManager.class);
-    }
-
-    /** It always executes after {@link After}. */
-    private void tearDownBase() {
-        mObjectTracker.tearDown(mPostAssertionRule::addError);
-
+    @After
+    public void tearDown() throws Exception {
         // Synchronous execution of removeStacksWithActivityTypes() ensures that all activities but
         // home are cleaned up from the stack at the end of each test. Am force stop shell commands
         // might be asynchronous and could interrupt the stack cleanup process if executed first.
@@ -472,6 +468,18 @@
         stopTestPackage(SECOND_TEST_PACKAGE);
         stopTestPackage(THIRD_TEST_PACKAGE);
         launchHomeActivityNoWait();
+
+        try {
+            // Skip empty stack/task check if a leakage was already found in previous test, or
+            // all tests afterward would also fail (since the leakage is always there) and fire
+            // unnecessary false alarms.
+            if (!sStackTaskLeakFound) {
+                mAmWmState.assertEmptyStackOrTask();
+            }
+        } catch (Throwable t) {
+            sStackTaskLeakFound = true;
+            throw t;
+        }
     }
 
     /**
@@ -755,10 +763,9 @@
             if (showSideActivity) {
                 if (isHomeRecentsComponent) {
                     // Launch Placeholder Side Activity
-                    final ComponentName sideActivityName =
-                            new ComponentName(mContext, SideActivity.class);
-                    launchActivityNoWait(sideActivityName);
-                    mAmWmState.waitForActivityState(sideActivityName, STATE_RESUMED);
+                    final Activity sideActivity = mSideActivityRule.launchActivity(
+                            new Intent());
+                    mAmWmState.waitForActivityState(sideActivity.getComponentName(), STATE_RESUMED);
                 }
 
                 // There are two cases when showSideActivity == true:
@@ -929,7 +936,7 @@
         mAmWmState.assertFocusedStack("Top activity's stack must also be on top", frontStackId);
         mAmWmState.assertVisibility(activityName, true /* visible */);
     }
-
+    
     // TODO: Switch to using a feature flag, when available.
     protected static boolean isUiModeLockedToVrHeadset() {
         final String output = runCommandAndPrintOutput("dumpsys uimode");
@@ -1017,31 +1024,6 @@
                 .getBoolean(android.R.bool.config_perDisplayFocusEnabled);
     }
 
-    /** @see ObjectTracker#manage(AutoCloseable) */
-    protected HomeActivitySession createManagedHomeActivitySession(ComponentName homeActivity) {
-        return mObjectTracker.manage(new HomeActivitySession(homeActivity));
-    }
-
-    /** @see ObjectTracker#manage(AutoCloseable) */
-    protected ActivitySessionClient createManagedActivityClientSession() {
-        return mObjectTracker.manage(new ActivitySessionClient(mContext));
-    }
-
-    /** @see ObjectTracker#manage(AutoCloseable) */
-    protected LockScreenSession createManagedLockScreenSession() {
-        return mObjectTracker.manage(new LockScreenSession());
-    }
-
-    /** @see ObjectTracker#manage(AutoCloseable) */
-    protected RotationSession createManagedRotationSession() {
-        return mObjectTracker.manage(new RotationSession());
-    }
-
-    /** @see ObjectTracker#manage(AutoCloseable) */
-    protected <T extends Activity> TestActivitySession<T> createManagedTestActivitySession() {
-        return new TestActivitySession<T>();
-    }
-
     /**
      * Test @Rule class that disables screen doze settings before each test method running and
      * restoring to initial values after test method finished.
@@ -1096,7 +1078,7 @@
         private ComponentName mOrigHome;
         private ComponentName mSessionHome;
 
-        HomeActivitySession(ComponentName sessionHome) {
+        public HomeActivitySession(ComponentName sessionHome) {
             mSessionHome = sessionHome;
             mPackageManager = mContext.getPackageManager();
 
@@ -1352,7 +1334,7 @@
         }
 
         @Override
-        public void close() {
+        public void close() throws Exception {
             mThread.quitSafely();
             mUserRotation.close();
             // Restore accelerometer_rotation preference.
@@ -1399,7 +1381,8 @@
      * @return {@code true} if test device respects settings of locked user rotation mode;
      * {@code false} if not.
      */
-    protected boolean supportsLockedUserRotation(RotationSession session, int displayId) {
+    protected boolean supportsLockedUserRotation(RotationSession session, int displayId)
+            throws Exception {
         final int origRotation = getDeviceRotation(displayId);
         // Use the same orientation as target rotation to avoid affect of app-requested orientation.
         final int targetRotation = (origRotation + 2) % 4;
@@ -2119,63 +2102,6 @@
         }
     }
 
-    /**
-     * The actions which wraps a test method. It is used to set necessary rules that cannot be
-     * overridden by subclasses. It executes in the outer scope of {@link Before} and {@link After}.
-     */
-    protected class WrapperRule implements TestRule {
-        private final Runnable mBefore;
-        private final Runnable mAfter;
-
-        protected WrapperRule(Runnable before, Runnable after) {
-            mBefore = before;
-            mAfter = after;
-        }
-
-        @Override
-        public Statement apply(final Statement base, final Description description) {
-            return new Statement() {
-                @Override
-                public void evaluate()  {
-                    if (mBefore != null) {
-                        mBefore.run();
-                    }
-                    try {
-                        base.evaluate();
-                    } catch (Throwable e) {
-                        mPostAssertionRule.addError(e);
-                    } finally {
-                        if (mAfter != null) {
-                            mAfter.run();
-                        }
-                    }
-                }
-            };
-        }
-    }
-
-    /**
-     * The post assertion to ensure all test methods don't violate the generic rule. It is also used
-     * to collect multiple errors.
-     */
-    private class PostAssertionRule extends ErrorCollector {
-        @Override
-        protected void verify() throws Throwable {
-            if (!sStackTaskLeakFound) {
-                // Skip empty stack/task check if a leakage was already found in previous test, or
-                // all tests afterward would also fail (since the leakage is always there) and fire
-                // unnecessary false alarms.
-                try {
-                    mAmWmState.assertEmptyStackOrTask();
-                } catch (Throwable t) {
-                    sStackTaskLeakFound = true;
-                    addError(t);
-                }
-            }
-            super.verify();
-        }
-    }
-
     // Activity used in place of recents when home is the recents component.
     public static class SideActivity extends Activity {
     }
diff --git a/tests/framework/base/windowmanager/util/src/android/server/wm/ObjectTracker.java b/tests/framework/base/windowmanager/util/src/android/server/wm/ObjectTracker.java
deleted file mode 100644
index 28434cf..0000000
--- a/tests/framework/base/windowmanager/util/src/android/server/wm/ObjectTracker.java
+++ /dev/null
@@ -1,134 +0,0 @@
-/*
- * Copyright (C) 2019 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 android.server.wm;
-
-import android.os.Build;
-
-import androidx.annotation.NonNull;
-import androidx.annotation.Nullable;
-
-import java.util.ArrayList;
-import java.util.LinkedList;
-import java.util.function.Consumer;
-
-/** A helper utility to track or manage object after a test method is done. */
-public class ObjectTracker {
-    private static final boolean DEBUG = "eng".equals(Build.TYPE);
-    private LinkedList<AutoCloseable> mAutoCloseables;
-    private LinkedList<ConsumableEntry> mConsumables;
-
-    /** The interface used for tracking whether an object is consumed. */
-    public interface Consumable {
-        boolean isConsumed();
-    }
-
-    private static class ConsumableEntry {
-        @NonNull
-        final Consumable mConsumable;
-        @Nullable
-        final Throwable mStackTrace;
-
-        ConsumableEntry(Consumable consumable, Throwable stackTrace) {
-            mConsumable = consumable;
-            mStackTrace = stackTrace;
-        }
-    }
-
-    /**
-     * If a {@link AutoCloseable} should be closed at the end of test, or it is not important when
-     * to close, then we can use this method to manage the {@link AutoCloseable}. Then the extra
-     * indents of try-with-resource can be eliminated. If the caller want to close the object
-     * manually, it should use {@link ObjectTracker#close} to cancel the management.
-     */
-    public <T extends AutoCloseable> T manage(@NonNull T autoCloseable) {
-        if (mAutoCloseables == null) {
-            mAutoCloseables = new LinkedList<>();
-        }
-        mAutoCloseables.add(autoCloseable);
-        return autoCloseable;
-    }
-
-    /**
-     * Closes the {@link AutoCloseable} and remove from the managed list so it won't be closed twice
-     * when leaving a test method.
-     */
-    public void close(@NonNull AutoCloseable autoCloseable) {
-        if (mAutoCloseables == null) {
-            return;
-        }
-        mAutoCloseables.remove(autoCloseable);
-        try {
-            autoCloseable.close();
-        } catch (Throwable e) {
-            throw new AssertionError("Failed to close " + autoCloseable, e);
-        }
-    }
-
-    /** Tracks the {@link Consumable} to avoid misusing of the object that should do something. */
-    public void track(@NonNull Consumable consumable) {
-        if (mConsumables == null) {
-            mConsumables = new LinkedList<>();
-        }
-        mConsumables.add(new ConsumableEntry(consumable,
-                DEBUG ? new Throwable().fillInStackTrace() : null));
-    }
-
-    /**
-     * Cleans up the managed object and make sure all tracked {@link Consumable} are consumed.
-     * This method must be called after each test method.
-     */
-    public void tearDown(@NonNull Consumer<Throwable> errorConsumer) {
-        ArrayList<Throwable> errors = null;
-        if (mAutoCloseables != null) {
-            while (!mAutoCloseables.isEmpty()) {
-                final AutoCloseable autoCloseable = mAutoCloseables.removeLast();
-                try {
-                    autoCloseable.close();
-                } catch (Throwable t) {
-                    StateLogger.logE("Failed to close " + autoCloseable, t);
-                    if (errors == null) {
-                        errors = new ArrayList<>();
-                    }
-                    errors.add(t);
-                }
-            }
-        }
-
-        if (mConsumables != null) {
-            while (!mConsumables.isEmpty()) {
-                final ConsumableEntry entry = mConsumables.removeFirst();
-                if (!entry.mConsumable.isConsumed()) {
-                    StateLogger.logE("Found unconsumed object " + entry.mConsumable
-                            + " that was created from:", entry.mStackTrace);
-                    final Throwable t =
-                            new IllegalStateException("Unconsumed object " + entry.mConsumable);
-                    if (entry.mStackTrace != null) {
-                        t.initCause(entry.mStackTrace);
-                    }
-                    if (errors == null) {
-                        errors = new ArrayList<>();
-                    }
-                    errors.add(t);
-                }
-            }
-        }
-
-        if (errors != null) {
-            errors.forEach(errorConsumer);
-        }
-    }
-}
diff --git a/tests/framework/base/windowmanager/util/src/android/server/wm/settings/SettingsSession.java b/tests/framework/base/windowmanager/util/src/android/server/wm/settings/SettingsSession.java
index 710b920..e5194c8 100644
--- a/tests/framework/base/windowmanager/util/src/android/server/wm/settings/SettingsSession.java
+++ b/tests/framework/base/windowmanager/util/src/android/server/wm/settings/SettingsSession.java
@@ -124,7 +124,7 @@
     }
 
     @Override
-    public void close() {
+    public void close() throws Exception {
         if (mHasInitialValue) {
             put(mUri, mSetter, mInitialValue);
             if (DEBUG) {