Adopt upstream Mockito's solution to mem leak.

The upstream solution essentially disallows sharing mock between tests,
so the mock of PowerManagerWrapper is moved to SystemServicesTestRule so
that TestWindowManagerPolicy can use that to create window. It also
makes the life of PowerManagerWrapper mock the same as the
TestWindowManagerPolicy, which is used in both ActivityTestBase and
WindowTestBase but PowerManagerWrapper was only available in
WindowTestBase.

Also, Track mocks created in ActivityStarterTests#testStartActivityPreconditions
and clear them in the middle of execution because the memory usage could be
~0.8G just to run ActivityStarterTests only because of that.

Bug: 123984854
Test: "tradefed.sh run commandAndExit WmTests --include-annotation
android.platform.test.annotations.Presubmit --exclude-annotation
androidx.test.filters.FlakyTest" suffers from a lot fewer OOME.

Change-Id: I274d711fb52084d4c67547374360e373ad1e28d8
diff --git a/services/tests/wmtests/src/com/android/server/wm/ActivityStarterTests.java b/services/tests/wmtests/src/com/android/server/wm/ActivityStarterTests.java
index 40a029c..81fbfe4 100644
--- a/services/tests/wmtests/src/com/android/server/wm/ActivityStarterTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/ActivityStarterTests.java
@@ -79,6 +79,7 @@
 import androidx.test.filters.SmallTest;
 
 import com.android.server.wm.LaunchParamsController.LaunchParamsModifier;
+import com.android.server.wm.utils.MockTracker;
 
 import org.junit.Before;
 import org.junit.Test;
@@ -186,6 +187,19 @@
         verifyStartActivityPreconditions(preconditions, 0 /*launchFlags*/, expectedResult);
     }
 
+    private void verifyStartActivityPreconditions(int preconditions, int launchFlags,
+            int expectedResult) {
+        // We track mocks created here because this is used in a single test
+        // (testStartActivityPreconditions) as a specific case, and mocks created inside it won't be
+        // used for other cases. To avoid extensive memory usage, we clean up all used mocks after
+        // each case. This is necessary because usually we only clean up mocks after a test
+        // finishes, but this test creates too many mocks that the intermediate memory usage can be
+        // ~0.8 GiB and thus very susceptible to OutOfMemoryException.
+        try (MockTracker tracker = new MockTracker()) {
+            verifyStartActivityPreconditionsUntracked(preconditions, launchFlags, expectedResult);
+        }
+    }
+
     /**
      * Excercises how the {@link ActivityStarter} reacts to various preconditions. The caller
      * provides a bitmask of all the set conditions (such as {@link #PRECONDITION_NO_CALLER_APP})
@@ -197,7 +211,7 @@
      * @param launchFlags The launch flags to be provided by the launch {@link Intent}.
      * @param expectedResult The expected result from the launch.
      */
-    private void verifyStartActivityPreconditions(int preconditions, int launchFlags,
+    private void verifyStartActivityPreconditionsUntracked(int preconditions, int launchFlags,
             int expectedResult) {
         final ActivityTaskManagerService service = mService;
         final IPackageManager packageManager = mock(IPackageManager.class);
diff --git a/services/tests/wmtests/src/com/android/server/wm/ActivityTestsBase.java b/services/tests/wmtests/src/com/android/server/wm/ActivityTestsBase.java
index 2e7218d..a5dc241 100644
--- a/services/tests/wmtests/src/com/android/server/wm/ActivityTestsBase.java
+++ b/services/tests/wmtests/src/com/android/server/wm/ActivityTestsBase.java
@@ -45,9 +45,7 @@
 import android.view.DisplayInfo;
 
 import com.android.server.AttributeCache;
-import com.android.server.wm.utils.MockTracker;
 
-import org.junit.After;
 import org.junit.Before;
 import org.junit.BeforeClass;
 import org.junit.Rule;
@@ -70,8 +68,6 @@
     RootActivityContainer mRootActivityContainer;
     ActivityStackSupervisor mSupervisor;
 
-    private MockTracker mMockTracker;
-
     // Default package name
     static final String DEFAULT_COMPONENT_PACKAGE_NAME = "com.foo";
 
@@ -85,19 +81,11 @@
 
     @Before
     public void setUpBase() {
-        mMockTracker = new MockTracker();
-
         mService = mSystemServicesTestRule.getActivityTaskManagerService();
         mSupervisor = mService.mStackSupervisor;
         mRootActivityContainer = mService.mRootActivityContainer;
     }
 
-    @After
-    public void tearDownBase() {
-        mMockTracker.close();
-        mMockTracker = null;
-    }
-
     /** Creates a {@link TestActivityDisplay}. */
     TestActivityDisplay createNewActivityDisplay() {
         return TestActivityDisplay.create(mSupervisor);
diff --git a/services/tests/wmtests/src/com/android/server/wm/SystemServicesTestRule.java b/services/tests/wmtests/src/com/android/server/wm/SystemServicesTestRule.java
index c42098e..fc02ae7 100644
--- a/services/tests/wmtests/src/com/android/server/wm/SystemServicesTestRule.java
+++ b/services/tests/wmtests/src/com/android/server/wm/SystemServicesTestRule.java
@@ -75,11 +75,11 @@
 import com.android.server.policy.WindowManagerPolicy;
 import com.android.server.statusbar.StatusBarManagerInternal;
 import com.android.server.uri.UriGrantsManagerInternal;
-import com.android.server.wm.utils.MockTracker;
 
 import org.junit.rules.TestRule;
 import org.junit.runner.Description;
 import org.junit.runners.model.Statement;
+import org.mockito.Mockito;
 import org.mockito.quality.Strictness;
 
 import java.io.File;
@@ -99,25 +99,19 @@
     private final AtomicBoolean mCurrentMessagesProcessed = new AtomicBoolean(false);
 
     private Context mContext;
-    private MockTracker mMockTracker;
     private StaticMockitoSession mMockitoSession;
     ServiceThread mHandlerThread;
     private ActivityManagerService mAmService;
     private ActivityTaskManagerService mAtmService;
     private WindowManagerService mWmService;
     private TestWindowManagerPolicy mWMPolicy;
+    private WindowState.PowerManagerWrapper mPowerManagerWrapper;
     private InputManagerService mImService;
     /**
      * Spied {@link SurfaceControl.Transaction} class than can be used to verify calls.
      */
     SurfaceControl.Transaction mTransaction;
 
-    /** {@link MockTracker} to track mocks created by {@link SystemServicesTestRule}. */
-    private static class Tracker extends MockTracker {
-        // This empty extended class is necessary since Mockito distinguishes a listener by it
-        // class.
-    }
-
     @Override
     public Statement apply(Statement base, Description description) {
         return new Statement() {
@@ -134,23 +128,17 @@
     }
 
     private void setUp() {
-        try {
-            mMockTracker = new Tracker();
+        mMockitoSession = mockitoSession()
+                .spyStatic(LocalServices.class)
+                .mockStatic(LockGuard.class)
+                .mockStatic(Watchdog.class)
+                .strictness(Strictness.LENIENT)
+                .startMocking();
 
-            mMockitoSession = mockitoSession()
-                    .spyStatic(LocalServices.class)
-                    .mockStatic(LockGuard.class)
-                    .mockStatic(Watchdog.class)
-                    .strictness(Strictness.LENIENT)
-                    .startMocking();
-
-            setUpSystemCore();
-            setUpLocalServices();
-            setUpActivityTaskManagerService();
-            setUpWindowManagerService();
-        } finally {
-            mMockTracker.stopTracking();
-        }
+        setUpSystemCore();
+        setUpLocalServices();
+        setUpActivityTaskManagerService();
+        setUpWindowManagerService();
     }
 
     private void setUpSystemCore() {
@@ -260,7 +248,9 @@
     }
 
     private void setUpWindowManagerService() {
-        mWMPolicy = new TestWindowManagerPolicy(this::getWindowManagerService);
+        mPowerManagerWrapper = mock(WindowState.PowerManagerWrapper.class);
+        mWMPolicy = new TestWindowManagerPolicy(this::getWindowManagerService,
+                mPowerManagerWrapper);
         mWmService = WindowManagerService.main(
                 mContext, mImService, false, false, mWMPolicy, mAtmService, StubTransaction::new);
         spyOn(mWmService);
@@ -305,9 +295,12 @@
         DeviceConfig.removeOnPropertiesChangedListener(mWmService.mPropertiesChangedListener);
         mWmService = null;
         mWMPolicy = null;
+        mPowerManagerWrapper = null;
 
         tearDownLocalServices();
         tearDownSystemCore();
+
+        Mockito.framework().clearInlineMocks();
     }
 
     private void tearDownSystemCore() {
@@ -316,11 +309,6 @@
             mMockitoSession = null;
         }
 
-        if (mMockTracker != null) {
-            mMockTracker.close();
-            mMockTracker = null;
-        }
-
         if (mHandlerThread != null) {
             // Make sure there are no running messages and then quit the thread so the next test
             // won't be affected.
@@ -352,6 +340,10 @@
         return mAtmService;
     }
 
+    WindowState.PowerManagerWrapper getPowerManagerWrapper() {
+        return mPowerManagerWrapper;
+    }
+
     void cleanupWindowManagerHandlers() {
         final WindowManagerService wm = getWindowManagerService();
         if (wm == null) {
diff --git a/services/tests/wmtests/src/com/android/server/wm/TestWindowManagerPolicy.java b/services/tests/wmtests/src/com/android/server/wm/TestWindowManagerPolicy.java
index 2e5ce69..bb89446 100644
--- a/services/tests/wmtests/src/com/android/server/wm/TestWindowManagerPolicy.java
+++ b/services/tests/wmtests/src/com/android/server/wm/TestWindowManagerPolicy.java
@@ -40,12 +40,14 @@
 import com.android.internal.policy.IKeyguardDismissCallback;
 import com.android.internal.policy.IShortcutService;
 import com.android.server.policy.WindowManagerPolicy;
+import com.android.server.wm.WindowState.PowerManagerWrapper;
 
 import java.io.PrintWriter;
 import java.util.function.Supplier;
 
 class TestWindowManagerPolicy implements WindowManagerPolicy {
     private final Supplier<WindowManagerService> mWmSupplier;
+    private final PowerManagerWrapper mPowerManagerWrapper;
 
     int mRotationToReport = 0;
     boolean mKeyguardShowingAndNotOccluded = false;
@@ -53,8 +55,10 @@
 
     private Runnable mRunnableWhenAddingSplashScreen;
 
-    TestWindowManagerPolicy(Supplier<WindowManagerService> wmSupplier) {
+    TestWindowManagerPolicy(Supplier<WindowManagerService> wmSupplier,
+            PowerManagerWrapper powerManagerWrapper) {
         mWmSupplier = wmSupplier;
+        mPowerManagerWrapper = powerManagerWrapper;
     }
 
     @Override
@@ -121,7 +125,7 @@
             doReturn(mock(IBinder.class)).when(iWindow).asBinder();
             window = WindowTestsBase.createWindow(null, TYPE_APPLICATION_STARTING, atoken,
                     "Starting window", 0 /* ownerId */, false /* internalWindows */, wm,
-                    mock(Session.class), iWindow);
+                    mock(Session.class), iWindow, mPowerManagerWrapper);
             atoken.startingWindow = window;
         }
         if (mRunnableWhenAddingSplashScreen != null) {
diff --git a/services/tests/wmtests/src/com/android/server/wm/WindowStateTests.java b/services/tests/wmtests/src/com/android/server/wm/WindowStateTests.java
index 763d086..b731628 100644
--- a/services/tests/wmtests/src/com/android/server/wm/WindowStateTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/WindowStateTests.java
@@ -346,24 +346,28 @@
         firstWindow.mAttrs.flags |= WindowManager.LayoutParams.FLAG_TURN_SCREEN_ON;
         secondWindow.mAttrs.flags |= WindowManager.LayoutParams.FLAG_TURN_SCREEN_ON;
 
-        reset(sPowerManagerWrapper);
+        final WindowState.PowerManagerWrapper powerManagerWrapper =
+                mSystemServicesTestRule.getPowerManagerWrapper();
+        reset(powerManagerWrapper);
         firstWindow.prepareWindowToDisplayDuringRelayout(false /*wasVisible*/);
-        verify(sPowerManagerWrapper).wakeUp(anyLong(), anyInt(), anyString());
+        verify(powerManagerWrapper).wakeUp(anyLong(), anyInt(), anyString());
 
-        reset(sPowerManagerWrapper);
+        reset(powerManagerWrapper);
         secondWindow.prepareWindowToDisplayDuringRelayout(false /*wasVisible*/);
-        verify(sPowerManagerWrapper).wakeUp(anyLong(), anyInt(), anyString());
+        verify(powerManagerWrapper).wakeUp(anyLong(), anyInt(), anyString());
     }
 
     private void testPrepareWindowToDisplayDuringRelayout(WindowState appWindow,
             boolean expectedWakeupCalled, boolean expectedCurrentLaunchCanTurnScreenOn) {
-        reset(sPowerManagerWrapper);
+        final WindowState.PowerManagerWrapper powerManagerWrapper =
+                mSystemServicesTestRule.getPowerManagerWrapper();
+        reset(powerManagerWrapper);
         appWindow.prepareWindowToDisplayDuringRelayout(false /* wasVisible */);
 
         if (expectedWakeupCalled) {
-            verify(sPowerManagerWrapper).wakeUp(anyLong(), anyInt(), anyString());
+            verify(powerManagerWrapper).wakeUp(anyLong(), anyInt(), anyString());
         } else {
-            verify(sPowerManagerWrapper, never()).wakeUp(anyLong(), anyInt(), anyString());
+            verify(powerManagerWrapper, never()).wakeUp(anyLong(), anyInt(), anyString());
         }
         // If wakeup is expected to be called, the currentLaunchCanTurnScreenOn should be false
         // because the state will be consumed.
diff --git a/services/tests/wmtests/src/com/android/server/wm/WindowTestsBase.java b/services/tests/wmtests/src/com/android/server/wm/WindowTestsBase.java
index 0732066..d1cf1c3 100644
--- a/services/tests/wmtests/src/com/android/server/wm/WindowTestsBase.java
+++ b/services/tests/wmtests/src/com/android/server/wm/WindowTestsBase.java
@@ -40,7 +40,6 @@
 
 import static org.mockito.Mockito.mock;
 
-
 import android.content.Context;
 import android.content.res.Configuration;
 import android.testing.DexmakerShareClassLoaderRule;
@@ -52,15 +51,12 @@
 import android.view.WindowManager;
 
 import com.android.server.AttributeCache;
-import com.android.server.wm.utils.MockTracker;
 
 import org.junit.After;
-import org.junit.AfterClass;
 import org.junit.Before;
 import org.junit.BeforeClass;
 import org.junit.Rule;
 
-import java.io.IOException;
 import java.util.HashSet;
 import java.util.LinkedList;
 
@@ -91,8 +87,6 @@
     WindowState mChildAppWindowBelow;
     HashSet<WindowState> mCommonWindows;
 
-    private MockTracker mMockTracker;
-
     /**
      * Spied {@link Transaction} class than can be used to verify calls.
      */
@@ -104,24 +98,13 @@
     @Rule
     public final SystemServicesTestRule mSystemServicesTestRule = new SystemServicesTestRule();
 
-    static WindowState.PowerManagerWrapper sPowerManagerWrapper;
-
     @BeforeClass
     public static void setUpOnceBase() {
         AttributeCache.init(getInstrumentation().getTargetContext());
-
-        sPowerManagerWrapper = mock(WindowState.PowerManagerWrapper.class);
-    }
-
-    @AfterClass
-    public static void tearDownOnceBase() throws IOException {
-        sPowerManagerWrapper = null;
     }
 
     @Before
     public void setUpBase() {
-        mMockTracker = new MockTracker();
-
         // If @Before throws an exception, the error isn't logged. This will make sure any failures
         // in the set up are clear. This can be removed when b/37850063 is fixed.
         try {
@@ -205,11 +188,7 @@
         } catch (Exception e) {
             Log.e(TAG, "Failed to tear down test", e);
             throw e;
-        } finally {
-            mMockTracker.close();
-            mMockTracker = null;
         }
-
     }
 
     private WindowState createCommonWindow(WindowState parent, int type, String name) {
@@ -331,12 +310,13 @@
     WindowState createWindow(WindowState parent, int type, WindowToken token, String name,
             int ownerId, boolean ownerCanAddInternalSystemWindow) {
         return createWindow(parent, type, token, name, ownerId, ownerCanAddInternalSystemWindow,
-                mWm, mMockSession, mIWindow);
+                mWm, mMockSession, mIWindow, mSystemServicesTestRule.getPowerManagerWrapper());
     }
 
     static WindowState createWindow(WindowState parent, int type, WindowToken token,
             String name, int ownerId, boolean ownerCanAddInternalSystemWindow,
-            WindowManagerService service, Session session, IWindow iWindow) {
+            WindowManagerService service, Session session, IWindow iWindow,
+            WindowState.PowerManagerWrapper powerManagerWrapper) {
         synchronized (service.mGlobalLock) {
             final WindowManager.LayoutParams attrs = new WindowManager.LayoutParams(type);
             attrs.setTitle(name);
@@ -344,7 +324,7 @@
             final WindowState w = new WindowState(service, session, iWindow, token, parent,
                     OP_NONE,
                     0, attrs, VISIBLE, ownerId, ownerCanAddInternalSystemWindow,
-                    sPowerManagerWrapper);
+                    powerManagerWrapper);
             // TODO: Probably better to make this call in the WindowState ctor to avoid errors with
             // adding it to the token...
             token.addWindow(w);
diff --git a/services/tests/wmtests/src/com/android/server/wm/utils/MockTracker.java b/services/tests/wmtests/src/com/android/server/wm/utils/MockTracker.java
index a6e675a..7f09482 100644
--- a/services/tests/wmtests/src/com/android/server/wm/utils/MockTracker.java
+++ b/services/tests/wmtests/src/com/android/server/wm/utils/MockTracker.java
@@ -89,7 +89,7 @@
 
         for (final Object mock : mMocks.keySet()) {
             if (MockUtil.isMock(mock)) {
-                Mockito.reset(mock);
+                mMockitoFramework.clearInlineMock(mock);
             }
         }
         mMocks.clear();