Merge "Revert "Added mockito-extended support to WmTests.""
diff --git a/services/core/java/com/android/server/policy/PhoneWindowManager.java b/services/core/java/com/android/server/policy/PhoneWindowManager.java
index dc8ee0a..282746a 100644
--- a/services/core/java/com/android/server/policy/PhoneWindowManager.java
+++ b/services/core/java/com/android/server/policy/PhoneWindowManager.java
@@ -1925,13 +1925,8 @@
 
     @Override
     public void setDefaultDisplay(DisplayContentInfo displayContentInfo) {
-        setDefaultDisplay(displayContentInfo.getDisplay(), displayContentInfo.getDisplayRotation());
-    }
-
-    @VisibleForTesting
-    void setDefaultDisplay(Display display, DisplayRotation displayRotation) {
-        mDefaultDisplay = display;
-        mDefaultDisplayRotation = displayRotation;
+        mDefaultDisplay = displayContentInfo.getDisplay();
+        mDefaultDisplayRotation = displayContentInfo.getDisplayRotation();
         mDefaultDisplayPolicy = mDefaultDisplayRotation.getDisplayPolicy();
         mDefaultOrientationListener = mDefaultDisplayRotation.getOrientationListener();
     }
@@ -2668,11 +2663,7 @@
 
     @Override
     public void onConfigurationChanged(DisplayContentInfo displayContentInfo) {
-        onConfigurationChanged(displayContentInfo.getDisplayRotation());
-    }
-
-    @VisibleForTesting
-    void onConfigurationChanged(DisplayRotation displayRotation) {
+        final DisplayRotation displayRotation = displayContentInfo.getDisplayRotation();
         // TODO(multi-display): Define policy for secondary displays.
         if (!displayRotation.isDefaultDisplay) {
             return;
diff --git a/services/core/java/com/android/server/wm/ActivityStackSupervisor.java b/services/core/java/com/android/server/wm/ActivityStackSupervisor.java
index cd708b9..082f521 100644
--- a/services/core/java/com/android/server/wm/ActivityStackSupervisor.java
+++ b/services/core/java/com/android/server/wm/ActivityStackSupervisor.java
@@ -315,7 +315,7 @@
     /** The number of distinct task ids that can be assigned to the tasks of a single user */
     private static final int MAX_TASK_IDS_PER_USER = UserHandle.PER_USER_RANGE;
 
-    final ActivityTaskManagerService mService;
+    ActivityTaskManagerService mService;
 
     /** The historial list of recent tasks including inactive tasks */
     RecentTasks mRecentTasks;
@@ -618,6 +618,11 @@
     }
 
     @VisibleForTesting
+    void setService(ActivityTaskManagerService service) {
+        mService = service;
+    }
+
+    @VisibleForTesting
     void setWindowContainerController(RootWindowContainerController controller) {
         mWindowContainerController = controller;
     }
diff --git a/services/core/java/com/android/server/wm/ActivityTaskManagerService.java b/services/core/java/com/android/server/wm/ActivityTaskManagerService.java
index 30eb500..1d00075 100644
--- a/services/core/java/com/android/server/wm/ActivityTaskManagerService.java
+++ b/services/core/java/com/android/server/wm/ActivityTaskManagerService.java
@@ -341,8 +341,7 @@
     ActivityManagerInternal mAmInternal;
     UriGrantsManagerInternal mUgmInternal;
     private PackageManagerInternal mPmInternal;
-    @VisibleForTesting
-    final ActivityTaskManagerInternal mInternal;
+    private ActivityTaskManagerInternal mInternal;
     PowerManagerInternal mPowerManagerInternal;
     private UsageStatsManagerInternal mUsageStatsInternal;
 
@@ -644,7 +643,6 @@
         mSystemThread = ActivityThread.currentActivityThread();
         mUiContext = mSystemThread.getSystemUiContext();
         mLifecycleManager = new ClientLifecycleManager();
-        mInternal = new LocalService();
         GL_ES_VERSION = SystemProperties.getInt("ro.opengles.version", GL_ES_VERSION_UNDEFINED);
     }
 
@@ -895,6 +893,7 @@
     }
 
     private void start() {
+        mInternal = new LocalService();
         LocalServices.addService(ActivityTaskManagerInternal.class, mInternal);
     }
 
@@ -6871,4 +6870,4 @@
             }
         }
     }
-}
+}
\ No newline at end of file
diff --git a/services/tests/wmtests/Android.mk b/services/tests/wmtests/Android.mk
index 81cb602..9655b3d 100644
--- a/services/tests/wmtests/Android.mk
+++ b/services/tests/wmtests/Android.mk
@@ -18,7 +18,7 @@
     services.core \
     androidx.test.runner \
     androidx.test.rules \
-    mockito-target-extended-minus-junit4 \
+    mockito-target-minus-junit4 \
     platform-test-annotations \
     truth-prebuilt \
     testables \
@@ -36,11 +36,6 @@
 
 LOCAL_CERTIFICATE := platform
 
-# These are not normally accessible from apps so they must be explicitly included.
-LOCAL_JNI_SHARED_LIBRARIES := \
-    libdexmakerjvmtiagent \
-    libstaticjvmtiagent \
-
 LOCAL_ADDITIONAL_DEPENDENCIES := $(LOCAL_PATH)/Android.mk
 
 LOCAL_JACK_FLAGS := --multi-dex native
diff --git a/services/tests/wmtests/AndroidManifest.xml b/services/tests/wmtests/AndroidManifest.xml
index ba3dc52..ff84803 100644
--- a/services/tests/wmtests/AndroidManifest.xml
+++ b/services/tests/wmtests/AndroidManifest.xml
@@ -35,7 +35,7 @@
     <uses-permission android:name="android.permission.ACCESS_SURFACE_FLINGER" />
     <uses-permission android:name="android.permission.READ_FRAME_BUFFER" />
 
-    <application android:testOnly="true" android:debuggable="true">
+    <application android:testOnly="true">
         <activity android:name="com.android.server.wm.TaskStackChangedListenerTest$ActivityA" />
         <activity android:name="com.android.server.wm.TaskStackChangedListenerTest$ActivityB" />
         <activity android:name="com.android.server.wm.TaskStackChangedListenerTest$ActivityRequestedOrientationChange" />
diff --git a/services/tests/wmtests/src/com/android/server/policy/PhoneWindowManagerTestBase.java b/services/tests/wmtests/src/com/android/server/policy/PhoneWindowManagerTestBase.java
index 6521cf5..fc8fe23 100644
--- a/services/tests/wmtests/src/com/android/server/policy/PhoneWindowManagerTestBase.java
+++ b/services/tests/wmtests/src/com/android/server/policy/PhoneWindowManagerTestBase.java
@@ -28,11 +28,10 @@
 import static android.view.WindowManager.LayoutParams.TYPE_NAVIGATION_BAR;
 import static android.view.WindowManager.LayoutParams.TYPE_STATUS_BAR;
 
-import static com.android.dx.mockito.inline.extended.ExtendedMockito.mock;
-import static com.android.server.wm.WindowTestUtils.createTestDisplayRotation;
 import static com.android.server.wm.utils.CoordinateTransforms.transformPhysicalToLogicalCoordinates;
 
 import static org.junit.Assert.assertEquals;
+import static org.mockito.Mockito.mock;
 
 import android.content.Context;
 import android.content.ContextWrapper;
@@ -60,7 +59,6 @@
 
 import com.android.server.policy.keyguard.KeyguardServiceDelegate;
 import com.android.server.wm.DisplayFrames;
-import com.android.server.wm.DisplayRotation;
 import com.android.server.wm.WindowTestUtils.TestDisplayContent;
 import com.android.server.wm.utils.WmDisplayCutout;
 
@@ -97,8 +95,7 @@
         mContext.getResourceMocker().addOverride(
                 com.android.internal.R.dimen.navigation_bar_width, NAV_BAR_HEIGHT);
 
-        InstrumentationRegistry.getInstrumentation().runOnMainSync(
-                () -> mPolicy = new TestablePhoneWindowManager(mContext));
+        mPolicy = TestablePhoneWindowManager.create(mContext);
 
         updateDisplayFrames();
     }
@@ -235,21 +232,9 @@
         }
     }
 
-    class TestablePhoneWindowManager extends PhoneWindowManager {
+    static class TestablePhoneWindowManager extends PhoneWindowManager {
 
-        TestablePhoneWindowManager(Context context) {
-            mContext = context;
-            mKeyguardDelegate = mock(KeyguardServiceDelegate.class);
-            mAccessibilityManager = new AccessibilityManager(context,
-                    mock(IAccessibilityManager.class), UserHandle.USER_CURRENT);
-            mSystemGestures = mock(SystemGesturesPointerEventListener.class);
-
-            final TestDisplayContent displayContent = mock(TestDisplayContent.class);
-            final DisplayRotation displayRotation =
-                    createTestDisplayRotation(mContext, displayContent);
-
-            setDefaultDisplay(displayContent.getDisplay(), displayRotation);
-            onConfigurationChanged(displayRotation);
+        TestablePhoneWindowManager() {
         }
 
         @Override
@@ -271,5 +256,22 @@
             adjustWindowParamsLw(state, state.getAttrs(), true /* hasStatusBarPermission */);
             assertEquals(WindowManagerGlobal.ADD_OKAY, prepareAddWindowLw(state, state.getAttrs()));
         }
+
+        public static TestablePhoneWindowManager create(Context context) {
+            TestablePhoneWindowManager[] policy = new TestablePhoneWindowManager[1];
+            InstrumentationRegistry.getInstrumentation().runOnMainSync(() -> {
+                policy[0] = new TestablePhoneWindowManager();
+                policy[0].mContext = context;
+                policy[0].mKeyguardDelegate = mock(KeyguardServiceDelegate.class);
+                policy[0].mAccessibilityManager = new AccessibilityManager(context,
+                        mock(IAccessibilityManager.class), UserHandle.USER_CURRENT);
+                policy[0].mSystemGestures = mock(SystemGesturesPointerEventListener.class);
+
+                final TestDisplayContent displayContent = TestDisplayContent.create(context);
+                policy[0].setDefaultDisplay(displayContent);
+                policy[0].onConfigurationChanged(displayContent);
+            });
+            return policy[0];
+        }
     }
 }
diff --git a/services/tests/wmtests/src/com/android/server/wm/ActivityDisplayTests.java b/services/tests/wmtests/src/com/android/server/wm/ActivityDisplayTests.java
index 41ed3b5..7a9c8dc 100644
--- a/services/tests/wmtests/src/com/android/server/wm/ActivityDisplayTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/ActivityDisplayTests.java
@@ -23,19 +23,20 @@
 import static android.content.pm.ActivityInfo.FLAG_ALWAYS_FOCUSABLE;
 import static android.content.pm.ActivityInfo.FLAG_SHOW_WHEN_LOCKED;
 
-import static com.android.dx.mockito.inline.extended.ExtendedMockito.doReturn;
-import static com.android.dx.mockito.inline.extended.ExtendedMockito.never;
-import static com.android.dx.mockito.inline.extended.ExtendedMockito.reset;
-import static com.android.dx.mockito.inline.extended.ExtendedMockito.spy;
-import static com.android.dx.mockito.inline.extended.ExtendedMockito.verify;
 import static com.android.server.wm.ActivityStackSupervisor.ON_TOP;
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertTrue;
 import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.reset;
+import static org.mockito.Mockito.spy;
+import static org.mockito.Mockito.verify;
 
 import android.platform.test.annotations.Presubmit;
+
 import androidx.test.filters.SmallTest;
 
 import org.junit.Before;
diff --git a/services/tests/wmtests/src/com/android/server/wm/ActivityMetricsLaunchObserverTests.java b/services/tests/wmtests/src/com/android/server/wm/ActivityMetricsLaunchObserverTests.java
index c7f0521..215c51d 100644
--- a/services/tests/wmtests/src/com/android/server/wm/ActivityMetricsLaunchObserverTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/ActivityMetricsLaunchObserverTests.java
@@ -21,12 +21,11 @@
 import static android.app.WindowConfiguration.ACTIVITY_TYPE_STANDARD;
 import static android.app.WindowConfiguration.WINDOWING_MODE_FULLSCREEN;
 
-import static com.android.dx.mockito.inline.extended.ExtendedMockito.mock;
-import static com.android.dx.mockito.inline.extended.ExtendedMockito.verify;
-import static com.android.dx.mockito.inline.extended.ExtendedMockito.verifyNoMoreInteractions;
-
 import static org.mockito.ArgumentMatchers.anyInt;
 import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.verifyNoMoreInteractions;
 
 import android.content.Intent;
 import android.os.SystemClock;
@@ -35,6 +34,7 @@
 
 import androidx.test.filters.FlakyTest;
 import androidx.test.filters.SmallTest;
+import androidx.test.runner.AndroidJUnit4;
 
 import org.junit.Before;
 import org.junit.Test;
diff --git a/services/tests/wmtests/src/com/android/server/wm/ActivityRecordTests.java b/services/tests/wmtests/src/com/android/server/wm/ActivityRecordTests.java
index 0e2974e..b865772 100644
--- a/services/tests/wmtests/src/com/android/server/wm/ActivityRecordTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/ActivityRecordTests.java
@@ -32,9 +32,9 @@
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertTrue;
-import static com.android.dx.mockito.inline.extended.ExtendedMockito.any;
-import static com.android.dx.mockito.inline.extended.ExtendedMockito.doAnswer;
-import static com.android.dx.mockito.inline.extended.ExtendedMockito.when;
+import static org.mockito.Mockito.any;
+import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.when;
 
 import android.app.ActivityOptions;
 import android.app.servertransaction.ClientTransaction;
diff --git a/services/tests/wmtests/src/com/android/server/wm/ActivityStackSupervisorTests.java b/services/tests/wmtests/src/com/android/server/wm/ActivityStackSupervisorTests.java
index f755035..16dd92f 100644
--- a/services/tests/wmtests/src/com/android/server/wm/ActivityStackSupervisorTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/ActivityStackSupervisorTests.java
@@ -44,14 +44,14 @@
 import static org.mockito.ArgumentMatchers.anyInt;
 import static org.mockito.ArgumentMatchers.contains;
 import static org.mockito.ArgumentMatchers.eq;
-import static com.android.dx.mockito.inline.extended.ExtendedMockito.doAnswer;
-import static com.android.dx.mockito.inline.extended.ExtendedMockito.doReturn;
-import static com.android.dx.mockito.inline.extended.ExtendedMockito.mock;
-import static com.android.dx.mockito.inline.extended.ExtendedMockito.never;
-import static com.android.dx.mockito.inline.extended.ExtendedMockito.reset;
-import static com.android.dx.mockito.inline.extended.ExtendedMockito.spy;
-import static com.android.dx.mockito.inline.extended.ExtendedMockito.times;
-import static com.android.dx.mockito.inline.extended.ExtendedMockito.verify;
+import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.reset;
+import static org.mockito.Mockito.spy;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
 
 import android.app.ActivityOptions;
 import android.app.WaitResult;
diff --git a/services/tests/wmtests/src/com/android/server/wm/ActivityStackTests.java b/services/tests/wmtests/src/com/android/server/wm/ActivityStackTests.java
index b4d9c16..62767e3 100644
--- a/services/tests/wmtests/src/com/android/server/wm/ActivityStackTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/ActivityStackTests.java
@@ -44,10 +44,10 @@
 import static org.mockito.ArgumentMatchers.anyBoolean;
 import static org.mockito.ArgumentMatchers.anyInt;
 import static org.mockito.ArgumentMatchers.eq;
-import static com.android.dx.mockito.inline.extended.ExtendedMockito.doReturn;
-import static com.android.dx.mockito.inline.extended.ExtendedMockito.mock;
-import static com.android.dx.mockito.inline.extended.ExtendedMockito.spy;
-import static com.android.dx.mockito.inline.extended.ExtendedMockito.verify;
+import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.spy;
+import static org.mockito.Mockito.verify;
 
 import android.content.pm.ActivityInfo;
 import android.os.UserHandle;
diff --git a/services/tests/wmtests/src/com/android/server/wm/ActivityStartControllerTests.java b/services/tests/wmtests/src/com/android/server/wm/ActivityStartControllerTests.java
index 9d93c85..e8de05c 100644
--- a/services/tests/wmtests/src/com/android/server/wm/ActivityStartControllerTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/ActivityStartControllerTests.java
@@ -19,12 +19,12 @@
 import static android.app.WindowConfiguration.ACTIVITY_TYPE_STANDARD;
 import static android.app.WindowConfiguration.WINDOWING_MODE_FULLSCREEN;
 
-import static com.android.dx.mockito.inline.extended.ExtendedMockito.doReturn;
-import static com.android.dx.mockito.inline.extended.ExtendedMockito.eq;
-import static com.android.dx.mockito.inline.extended.ExtendedMockito.mock;
-import static com.android.dx.mockito.inline.extended.ExtendedMockito.spy;
-import static com.android.dx.mockito.inline.extended.ExtendedMockito.times;
-import static com.android.dx.mockito.inline.extended.ExtendedMockito.verify;
+import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.eq;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.spy;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
 
 import android.app.IApplicationThread;
 import android.content.Intent;
diff --git a/services/tests/wmtests/src/com/android/server/wm/ActivityStartInterceptorTest.java b/services/tests/wmtests/src/com/android/server/wm/ActivityStartInterceptorTest.java
index 2bad802..dda077e 100644
--- a/services/tests/wmtests/src/com/android/server/wm/ActivityStartInterceptorTest.java
+++ b/services/tests/wmtests/src/com/android/server/wm/ActivityStartInterceptorTest.java
@@ -25,7 +25,7 @@
 import static org.junit.Assert.assertTrue;
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.ArgumentMatchers.nullable;
-import static com.android.dx.mockito.inline.extended.ExtendedMockito.when;
+import static org.mockito.Mockito.when;
 
 import android.app.ActivityManagerInternal;
 import android.app.KeyguardManager;
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 80f30e3..88479ec 100644
--- a/services/tests/wmtests/src/com/android/server/wm/ActivityStarterTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/ActivityStarterTests.java
@@ -48,15 +48,15 @@
 import static org.mockito.ArgumentMatchers.anyInt;
 import static org.mockito.ArgumentMatchers.anyObject;
 import static org.mockito.ArgumentMatchers.eq;
-import static com.android.dx.mockito.inline.extended.ExtendedMockito.clearInvocations;
-import static com.android.dx.mockito.inline.extended.ExtendedMockito.doAnswer;
-import static com.android.dx.mockito.inline.extended.ExtendedMockito.doNothing;
-import static com.android.dx.mockito.inline.extended.ExtendedMockito.doReturn;
-import static com.android.dx.mockito.inline.extended.ExtendedMockito.mock;
-import static com.android.dx.mockito.inline.extended.ExtendedMockito.never;
-import static com.android.dx.mockito.inline.extended.ExtendedMockito.spy;
-import static com.android.dx.mockito.inline.extended.ExtendedMockito.times;
-import static com.android.dx.mockito.inline.extended.ExtendedMockito.verify;
+import static org.mockito.Mockito.clearInvocations;
+import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.doNothing;
+import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.spy;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
 
 import android.app.ActivityOptions;
 import android.app.IApplicationThread;
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 cbb8ecb..26286e2 100644
--- a/services/tests/wmtests/src/com/android/server/wm/ActivityTestsBase.java
+++ b/services/tests/wmtests/src/com/android/server/wm/ActivityTestsBase.java
@@ -29,16 +29,15 @@
 import static com.android.server.wm.ActivityStack.REMOVE_TASK_MODE_DESTROYING;
 import static com.android.server.wm.ActivityStackSupervisor.ON_TOP;
 
-import static com.android.dx.mockito.inline.extended.ExtendedMockito.any;
-import static com.android.dx.mockito.inline.extended.ExtendedMockito.anyBoolean;
-import static com.android.dx.mockito.inline.extended.ExtendedMockito.anyInt;
-import static com.android.dx.mockito.inline.extended.ExtendedMockito.anyString;
-import static com.android.dx.mockito.inline.extended.ExtendedMockito.doAnswer;
-import static com.android.dx.mockito.inline.extended.ExtendedMockito.doNothing;
-import static com.android.dx.mockito.inline.extended.ExtendedMockito.doReturn;
-import static com.android.dx.mockito.inline.extended.ExtendedMockito.mock;
-import static com.android.dx.mockito.inline.extended.ExtendedMockito.spy;
-import static com.android.dx.mockito.inline.extended.ExtendedMockito.spyOn;
+import static org.mockito.Mockito.any;
+import static org.mockito.Mockito.anyBoolean;
+import static org.mockito.Mockito.anyInt;
+import static org.mockito.Mockito.anyString;
+import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.doNothing;
+import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.spy;
 
 import android.app.ActivityManagerInternal;
 import android.app.ActivityOptions;
@@ -68,8 +67,6 @@
 import com.android.server.AttributeCache;
 import com.android.server.ServiceThread;
 import com.android.server.am.ActivityManagerService;
-import com.android.server.am.PendingIntentController;
-import com.android.server.firewall.IntentFirewall;
 import com.android.server.uri.UriGrantsManagerInternal;
 
 import org.junit.After;
@@ -119,13 +116,27 @@
     }
 
     ActivityTaskManagerService createActivityTaskManagerService() {
-        mService = new TestActivityTaskManagerService(mContext);
-        mSupervisor = mService.mStackSupervisor;
-        return mService;
+        final TestActivityTaskManagerService atm =
+                spy(new TestActivityTaskManagerService(mContext));
+        setupActivityManagerService(atm);
+        return atm;
     }
 
     void setupActivityTaskManagerService() {
-        createActivityTaskManagerService();
+        mService = createActivityTaskManagerService();
+        mSupervisor = mService.mStackSupervisor;
+    }
+
+    ActivityManagerService createActivityManagerService() {
+        final TestActivityTaskManagerService atm =
+                spy(new TestActivityTaskManagerService(mContext));
+        return setupActivityManagerService(atm);
+    }
+
+    ActivityManagerService setupActivityManagerService(TestActivityTaskManagerService atm) {
+        final TestActivityManagerService am = spy(new TestActivityManagerService(mTestInjector));
+        setupActivityManagerService(am, atm);
+        return am;
     }
 
     /** Creates a {@link TestActivityDisplay}. */
@@ -144,6 +155,32 @@
         return display;
     }
 
+    void setupActivityManagerService(
+            TestActivityManagerService am, TestActivityTaskManagerService atm) {
+        atm.setActivityManagerService(am.mIntentFirewall, am.mPendingIntentController);
+        atm.mAmInternal = am.getLocalService();
+        am.mAtmInternal = atm.getLocalService();
+        // Makes sure the supervisor is using with the spy object.
+        atm.mStackSupervisor.setService(atm);
+        doReturn(mock(IPackageManager.class)).when(am).getPackageManager();
+        doReturn(mock(IPackageManager.class)).when(atm).getPackageManager();
+        PackageManagerInternal mockPackageManager = mock(PackageManagerInternal.class);
+        doReturn(mockPackageManager).when(am).getPackageManagerInternalLocked();
+        doReturn(null).when(mockPackageManager).getDefaultHomeActivity(anyInt());
+        doNothing().when(am).grantEphemeralAccessLocked(anyInt(), any(), anyInt(), anyInt());
+        am.mActivityTaskManager = atm;
+        am.mWindowManager = prepareMockWindowManager();
+        atm.setWindowManager(am.mWindowManager);
+
+        // Put a home stack on the default display, so that we'll always have something focusable.
+        final TestActivityStackSupervisor supervisor =
+                (TestActivityStackSupervisor) atm.mStackSupervisor;
+        supervisor.mDisplay.createStack(WINDOWING_MODE_FULLSCREEN, ACTIVITY_TYPE_HOME, ON_TOP);
+        final TaskRecord task = new TaskBuilder(atm.mStackSupervisor)
+                .setStack(supervisor.getDefaultDisplay().getHomeStack()).build();
+        new ActivityBuilder(atm).setTask(task).build();
+    }
+
     /**
      * Builder for creating new activities.
      */
@@ -369,52 +406,23 @@
         }
     }
 
-    protected class TestActivityTaskManagerService extends ActivityTaskManagerService {
+    protected static class TestActivityTaskManagerService extends ActivityTaskManagerService {
+        private LockTaskController mLockTaskController;
+        private ActivityTaskManagerInternal mInternal;
         private PackageManagerInternal mPmInternal;
 
         // ActivityStackSupervisor may be created more than once while setting up AMS and ATMS.
         // We keep the reference in order to prevent creating it twice.
-        ActivityStackSupervisor mTestStackSupervisor;
+        private ActivityStackSupervisor mTestStackSupervisor;
 
         TestActivityTaskManagerService(Context context) {
             super(context);
-
-            mUgmInternal = mock(UriGrantsManagerInternal.class);
-
             mSupportsMultiWindow = true;
             mSupportsMultiDisplay = true;
             mSupportsSplitScreenMultiWindow = true;
             mSupportsFreeformWindowManagement = true;
             mSupportsPictureInPicture = true;
-
-            final TestActivityManagerService am =
-                    new TestActivityManagerService(mTestInjector, this);
-            // TODO: Get some odd native errors when I put the 2 lines below in
-            // TestActivityManagerService which should be okay...but it isn't.
-            spyOn(am);
-            am.initMocks();
-
-            // Put a home stack on the default display, so that we'll always have something
-            // focusable.
-            final TestActivityStackSupervisor supervisor =
-                    (TestActivityStackSupervisor) mStackSupervisor;
-            supervisor.mDisplay.createStack(WINDOWING_MODE_FULLSCREEN, ACTIVITY_TYPE_HOME, ON_TOP);
-            final TaskRecord task = new TaskBuilder(mStackSupervisor)
-                    .setStack(supervisor.getDefaultDisplay().getHomeStack()).build();
-            new ActivityBuilder(this).setTask(task).build();
-
-            spyOn(this);
-            spyOn(getLifecycleManager());
-            spyOn(getLockTaskController());
-            doReturn(mock(IPackageManager.class)).when(this).getPackageManager();
-        }
-
-        void setActivityManagerService(IntentFirewall intentFirewall,
-                PendingIntentController intentController, ActivityManagerInternal amInternal,
-                WindowManagerService wm) {
-            mAmInternal = amInternal;
-            setActivityManagerService(intentFirewall, intentController);
-            setWindowManager(wm);
+            mUgmInternal = mock(UriGrantsManagerInternal.class);
         }
 
         @Override
@@ -423,17 +431,54 @@
         }
 
         @Override
+        public LockTaskController getLockTaskController() {
+            if (mLockTaskController == null) {
+                mLockTaskController = spy(super.getLockTaskController());
+            }
+
+            return mLockTaskController;
+        }
+
+        @Override
         void updateUsageStats(ActivityRecord component, boolean resumed) {
         }
 
         @Override
-        protected ActivityStackSupervisor createStackSupervisor() {
+        protected final ActivityStackSupervisor createStackSupervisor() {
             if (mTestStackSupervisor == null) {
-                mTestStackSupervisor = new TestActivityStackSupervisor(this, mH.getLooper());
+                final ActivityStackSupervisor supervisor = spy(createTestSupervisor());
+                final KeyguardController keyguardController = mock(KeyguardController.class);
+
+                // Invoked during {@link ActivityStack} creation.
+                doNothing().when(supervisor).updateUIDsPresentOnDisplay();
+                // Always keep things awake.
+                doReturn(true).when(supervisor).hasAwakeDisplay();
+                // Called when moving activity to pinned stack.
+                doNothing().when(supervisor).ensureActivitiesVisibleLocked(any(), anyInt(),
+                        anyBoolean());
+                // Do not schedule idle timeouts
+                doNothing().when(supervisor).scheduleIdleTimeoutLocked(any());
+                // unit test version does not handle launch wake lock
+                doNothing().when(supervisor).acquireLaunchWakelock();
+                doReturn(keyguardController).when(supervisor).getKeyguardController();
+
+                supervisor.initialize();
+                mTestStackSupervisor = supervisor;
             }
             return mTestStackSupervisor;
         }
 
+        protected ActivityStackSupervisor createTestSupervisor() {
+            return new TestActivityStackSupervisor(this, mH.getLooper());
+        }
+
+        ActivityTaskManagerInternal getLocalService() {
+            if (mInternal == null) {
+                mInternal = new ActivityTaskManagerService.LocalService();
+            }
+            return mInternal;
+        }
+
         @Override
         PackageManagerInternal getPackageManagerInternalLocked() {
             if (mPmInternal == null) {
@@ -484,27 +529,20 @@
      * An {@link ActivityManagerService} subclass which provides a test
      * {@link ActivityStackSupervisor}.
      */
-    class TestActivityManagerService extends ActivityManagerService {
+    static class TestActivityManagerService extends ActivityManagerService {
 
-        TestActivityManagerService(TestInjector testInjector, TestActivityTaskManagerService atm) {
+        private ActivityManagerInternal mInternal;
+
+        TestActivityManagerService(TestInjector testInjector) {
             super(testInjector, testInjector.mHandlerThread);
-
-            mWindowManager = prepareMockWindowManager();
             mUgmInternal = mock(UriGrantsManagerInternal.class);
-
-            atm.setActivityManagerService(mIntentFirewall, mPendingIntentController,
-                    new LocalService(), mWindowManager);
-
-            mActivityTaskManager = atm;
-            mAtmInternal = atm.mInternal;
         }
 
-        void initMocks() {
-            doReturn(mock(IPackageManager.class)).when(this).getPackageManager();
-            PackageManagerInternal mockPackageManager = mock(PackageManagerInternal.class);
-            doReturn(mockPackageManager).when(this).getPackageManagerInternalLocked();
-            doReturn(null).when(mockPackageManager).getDefaultHomeActivity(anyInt());
-            doNothing().when(this).grantEphemeralAccessLocked(anyInt(), any(), anyInt(), anyInt());
+        ActivityManagerInternal getLocalService() {
+            if (mInternal == null) {
+                mInternal = new LocalService();
+            }
+            return mInternal;
         }
     }
 
@@ -512,40 +550,23 @@
      * An {@link ActivityStackSupervisor} which stubs out certain methods that depend on
      * setup not available in the test environment. Also specifies an injector for
      */
-    protected class TestActivityStackSupervisor extends ActivityStackSupervisor {
+    protected static class TestActivityStackSupervisor extends ActivityStackSupervisor {
         private ActivityDisplay mDisplay;
         private KeyguardController mKeyguardController;
 
-        TestActivityStackSupervisor(ActivityTaskManagerService service, Looper looper) {
+        public TestActivityStackSupervisor(ActivityTaskManagerService service, Looper looper) {
             super(service, looper);
-            spyOn(this);
             mDisplayManager =
                     (DisplayManager) mService.mContext.getSystemService(Context.DISPLAY_SERVICE);
             mWindowManager = prepareMockWindowManager();
             mKeyguardController = mock(KeyguardController.class);
             setWindowContainerController(mock(RootWindowContainerController.class));
-
-            // Invoked during {@link ActivityStack} creation.
-            doNothing().when(this).updateUIDsPresentOnDisplay();
-            // Always keep things awake.
-            doReturn(true).when(this).hasAwakeDisplay();
-            // Called when moving activity to pinned stack.
-            doNothing().when(this).ensureActivitiesVisibleLocked(any(), anyInt(),
-                    anyBoolean());
-            // Do not schedule idle timeouts
-            doNothing().when(this).scheduleIdleTimeoutLocked(any());
-            // unit test version does not handle launch wake lock
-            doNothing().when(this).acquireLaunchWakelock();
-            doReturn(mKeyguardController).when(this).getKeyguardController();
-
-            initialize();
         }
 
         @Override
         public void initialize() {
             super.initialize();
-            mDisplay = TestActivityDisplay.create(this, DEFAULT_DISPLAY);
-            spyOn(mDisplay);
+            mDisplay = spy(TestActivityDisplay.create(this, DEFAULT_DISPLAY));
             addChild(mDisplay, ActivityDisplay.POSITION_TOP);
         }
 
diff --git a/services/tests/wmtests/src/com/android/server/wm/AssistDataRequesterTest.java b/services/tests/wmtests/src/com/android/server/wm/AssistDataRequesterTest.java
index 9e1d88b..c6c1c52 100644
--- a/services/tests/wmtests/src/com/android/server/wm/AssistDataRequesterTest.java
+++ b/services/tests/wmtests/src/com/android/server/wm/AssistDataRequesterTest.java
@@ -32,9 +32,9 @@
 import static org.mockito.ArgumentMatchers.anyInt;
 import static org.mockito.ArgumentMatchers.anyString;
 import static org.mockito.ArgumentMatchers.eq;
-import static com.android.dx.mockito.inline.extended.ExtendedMockito.doAnswer;
-import static com.android.dx.mockito.inline.extended.ExtendedMockito.doReturn;
-import static com.android.dx.mockito.inline.extended.ExtendedMockito.mock;
+import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.mock;
 
 import android.app.AppOpsManager;
 import android.app.IActivityManager;
diff --git a/services/tests/wmtests/src/com/android/server/wm/ClientLifecycleManagerTests.java b/services/tests/wmtests/src/com/android/server/wm/ClientLifecycleManagerTests.java
index 28dd458..f1d840d 100644
--- a/services/tests/wmtests/src/com/android/server/wm/ClientLifecycleManagerTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/ClientLifecycleManagerTests.java
@@ -16,10 +16,10 @@
 
 package com.android.server.wm;
 
-import static com.android.dx.mockito.inline.extended.ExtendedMockito.mock;
-import static com.android.dx.mockito.inline.extended.ExtendedMockito.spy;
-import static com.android.dx.mockito.inline.extended.ExtendedMockito.times;
-import static com.android.dx.mockito.inline.extended.ExtendedMockito.verify;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.spy;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
 
 import android.app.IApplicationThread;
 import android.app.servertransaction.ClientTransaction;
diff --git a/services/tests/wmtests/src/com/android/server/wm/LaunchParamsControllerTests.java b/services/tests/wmtests/src/com/android/server/wm/LaunchParamsControllerTests.java
index 62693f3..f8d64e9 100644
--- a/services/tests/wmtests/src/com/android/server/wm/LaunchParamsControllerTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/LaunchParamsControllerTests.java
@@ -28,15 +28,15 @@
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNotEquals;
-import static com.android.dx.mockito.inline.extended.ExtendedMockito.any;
-import static com.android.dx.mockito.inline.extended.ExtendedMockito.anyInt;
-import static com.android.dx.mockito.inline.extended.ExtendedMockito.doNothing;
-import static com.android.dx.mockito.inline.extended.ExtendedMockito.eq;
-import static com.android.dx.mockito.inline.extended.ExtendedMockito.mock;
-import static com.android.dx.mockito.inline.extended.ExtendedMockito.never;
-import static com.android.dx.mockito.inline.extended.ExtendedMockito.spy;
-import static com.android.dx.mockito.inline.extended.ExtendedMockito.times;
-import static com.android.dx.mockito.inline.extended.ExtendedMockito.verify;
+import static org.mockito.Mockito.any;
+import static org.mockito.Mockito.anyInt;
+import static org.mockito.Mockito.doNothing;
+import static org.mockito.Mockito.eq;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.spy;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
 
 import android.app.ActivityOptions;
 import android.content.ComponentName;
diff --git a/services/tests/wmtests/src/com/android/server/wm/LaunchParamsPersisterTests.java b/services/tests/wmtests/src/com/android/server/wm/LaunchParamsPersisterTests.java
index 42bce75b..59e9ce3 100644
--- a/services/tests/wmtests/src/com/android/server/wm/LaunchParamsPersisterTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/LaunchParamsPersisterTests.java
@@ -26,9 +26,9 @@
 import static org.junit.Assert.assertTrue;
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Matchers.any;
-import static com.android.dx.mockito.inline.extended.ExtendedMockito.mock;
-import static com.android.dx.mockito.inline.extended.ExtendedMockito.verify;
-import static com.android.dx.mockito.inline.extended.ExtendedMockito.when;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
 
 import android.content.ComponentName;
 import android.content.pm.PackageList;
diff --git a/services/tests/wmtests/src/com/android/server/wm/LockTaskControllerTest.java b/services/tests/wmtests/src/com/android/server/wm/LockTaskControllerTest.java
index 73eebb6..6b613ed 100644
--- a/services/tests/wmtests/src/com/android/server/wm/LockTaskControllerTest.java
+++ b/services/tests/wmtests/src/com/android/server/wm/LockTaskControllerTest.java
@@ -44,13 +44,13 @@
 import static org.mockito.ArgumentMatchers.anyInt;
 import static org.mockito.ArgumentMatchers.anyString;
 import static org.mockito.ArgumentMatchers.eq;
-import static com.android.dx.mockito.inline.extended.ExtendedMockito.doAnswer;
-import static com.android.dx.mockito.inline.extended.ExtendedMockito.mock;
-import static com.android.dx.mockito.inline.extended.ExtendedMockito.never;
-import static com.android.dx.mockito.inline.extended.ExtendedMockito.reset;
-import static com.android.dx.mockito.inline.extended.ExtendedMockito.times;
-import static com.android.dx.mockito.inline.extended.ExtendedMockito.verify;
-import static com.android.dx.mockito.inline.extended.ExtendedMockito.when;
+import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.reset;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
 
 import android.app.StatusBarManager;
 import android.app.admin.DevicePolicyManager;
diff --git a/services/tests/wmtests/src/com/android/server/wm/RecentTasksTest.java b/services/tests/wmtests/src/com/android/server/wm/RecentTasksTest.java
index c24b493..26241d2 100644
--- a/services/tests/wmtests/src/com/android/server/wm/RecentTasksTest.java
+++ b/services/tests/wmtests/src/com/android/server/wm/RecentTasksTest.java
@@ -38,8 +38,8 @@
 import static org.junit.Assert.fail;
 import static org.mockito.ArgumentMatchers.anyInt;
 import static org.mockito.ArgumentMatchers.anyString;
-import static com.android.dx.mockito.inline.extended.ExtendedMockito.doReturn;
-import static com.android.dx.mockito.inline.extended.ExtendedMockito.spy;
+import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.spy;
 
 import static java.lang.Integer.MAX_VALUE;
 
@@ -109,7 +109,9 @@
     @Before
     public void setUp() throws Exception {
         mTaskPersister = new TestTaskPersister(mContext.getFilesDir());
-        mTestService = new MyTestActivityTaskManagerService(mContext);
+        mTestService = spy(new MyTestActivityTaskManagerService(mContext));
+        final TestActivityManagerService am = spy(new MyTestActivityManagerService());
+        setupActivityManagerService(am, mTestService);
         mRecentTasks = (TestRecentTasks) mTestService.getRecentTasks();
         mRecentTasks.loadParametersFromResources(mContext.getResources());
         mHomeStack = mTestService.mStackSupervisor.getDefaultDisplay().getOrCreateStack(
@@ -865,11 +867,20 @@
         }
 
         @Override
-        protected ActivityStackSupervisor createStackSupervisor() {
-            if (mTestStackSupervisor == null) {
-                mTestStackSupervisor = new MyTestActivityStackSupervisor(this, mH.getLooper());
-            }
-            return mTestStackSupervisor;
+        protected ActivityStackSupervisor createTestSupervisor() {
+            return new MyTestActivityStackSupervisor(this, mH.getLooper());
+        }
+
+    }
+
+    private class MyTestActivityManagerService extends TestActivityManagerService {
+        MyTestActivityManagerService() {
+            super(mTestInjector);
+        }
+
+        @Override
+        public boolean isUserRunning(int userId, int flags) {
+            return true;
         }
     }
 
diff --git a/services/tests/wmtests/src/com/android/server/wm/RecentsAnimationTest.java b/services/tests/wmtests/src/com/android/server/wm/RecentsAnimationTest.java
index 93152f1..a01a3d9 100644
--- a/services/tests/wmtests/src/com/android/server/wm/RecentsAnimationTest.java
+++ b/services/tests/wmtests/src/com/android/server/wm/RecentsAnimationTest.java
@@ -22,14 +22,13 @@
 
 import static com.android.server.wm.RecentsAnimationController.REORDER_KEEP_IN_PLACE;
 
-import static com.android.dx.mockito.inline.extended.ExtendedMockito.any;
-import static com.android.dx.mockito.inline.extended.ExtendedMockito.doReturn;
-import static com.android.dx.mockito.inline.extended.ExtendedMockito.eq;
-import static com.android.dx.mockito.inline.extended.ExtendedMockito.mock;
-import static com.android.dx.mockito.inline.extended.ExtendedMockito.spy;
-import static com.android.dx.mockito.inline.extended.ExtendedMockito.spyOn;
-import static com.android.dx.mockito.inline.extended.ExtendedMockito.times;
-import static com.android.dx.mockito.inline.extended.ExtendedMockito.verify;
+import static org.mockito.Mockito.any;
+import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.eq;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.spy;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
 
 import android.content.ComponentName;
 import android.content.Context;
@@ -52,50 +51,61 @@
 public class RecentsAnimationTest extends ActivityTestsBase {
 
     private Context mContext = InstrumentationRegistry.getContext();
+    private TestActivityTaskManagerService mTestService;
     private ComponentName mRecentsComponent;
 
     @Before
     public void setUp() throws Exception {
         mRecentsComponent = new ComponentName(mContext.getPackageName(), "RecentsActivity");
-        mService = new TestActivityTaskManagerService(mContext);
-
-        final RecentTasks recentTasks = mService.getRecentTasks();
-        spyOn(recentTasks);
-        mRecentsComponent = new ComponentName(mContext.getPackageName(), "RecentsActivity");
-        doReturn(mRecentsComponent).when(recentTasks).getRecentsComponent();
+        mTestService = spy(new MyTestActivityTaskManagerService(mContext));
+        setupActivityManagerService(mTestService);
     }
 
     @Test
     public void testCancelAnimationOnStackOrderChange() {
         ActivityStack fullscreenStack =
-                mService.mStackSupervisor.getDefaultDisplay().createStack(
+                mTestService.mStackSupervisor.getDefaultDisplay().createStack(
                         WINDOWING_MODE_FULLSCREEN, ACTIVITY_TYPE_STANDARD, true /* onTop */);
-        ActivityStack recentsStack = mService.mStackSupervisor.getDefaultDisplay().createStack(
+        ActivityStack recentsStack = mTestService.mStackSupervisor.getDefaultDisplay().createStack(
                 WINDOWING_MODE_FULLSCREEN, ACTIVITY_TYPE_RECENTS, true /* onTop */);
-        ActivityRecord recentsActivity = new ActivityBuilder(mService)
+        ActivityRecord recentsActivity = new ActivityBuilder(mTestService)
                 .setComponent(mRecentsComponent)
                 .setCreateTask(true)
                 .setStack(recentsStack)
                 .build();
         ActivityStack fullscreenStack2 =
-                mService.mStackSupervisor.getDefaultDisplay().createStack(
+                mTestService.mStackSupervisor.getDefaultDisplay().createStack(
                         WINDOWING_MODE_FULLSCREEN, ACTIVITY_TYPE_STANDARD, true /* onTop */);
-        ActivityRecord fsActivity = new ActivityBuilder(mService)
+        ActivityRecord fsActivity = new ActivityBuilder(mTestService)
                 .setComponent(new ComponentName(mContext.getPackageName(), "App1"))
                 .setCreateTask(true)
                 .setStack(fullscreenStack2)
                 .build();
-        doReturn(true).when(mService.mWindowManager).canStartRecentsAnimation();
+        doReturn(true).when(mTestService.mWindowManager).canStartRecentsAnimation();
 
         // Start the recents animation
         Intent recentsIntent = new Intent();
         recentsIntent.setComponent(mRecentsComponent);
-        mService.startRecentsActivity(recentsIntent, null, mock(IRecentsAnimationRunner.class));
+        mTestService.startRecentsActivity(recentsIntent, null, mock(IRecentsAnimationRunner.class));
 
         fullscreenStack.moveToFront("Activity start");
 
         // Ensure that the recents animation was canceled
-        verify(mService.mWindowManager, times(1)).cancelRecentsAnimationSynchronously(
+        verify(mTestService.mWindowManager, times(1)).cancelRecentsAnimationSynchronously(
                 eq(REORDER_KEEP_IN_PLACE), any());
     }
+
+    private class MyTestActivityTaskManagerService extends TestActivityTaskManagerService {
+        MyTestActivityTaskManagerService(Context context) {
+            super(context);
+        }
+
+        @Override
+        protected RecentTasks createRecentTasks() {
+            RecentTasks recents = mock(RecentTasks.class);
+            doReturn(mRecentsComponent).when(recents).getRecentsComponent();
+            System.out.println(mRecentsComponent);
+            return recents;
+        }
+    }
 }
diff --git a/services/tests/wmtests/src/com/android/server/wm/TaskLaunchParamsModifierTests.java b/services/tests/wmtests/src/com/android/server/wm/TaskLaunchParamsModifierTests.java
index 63258c7..2168fab 100644
--- a/services/tests/wmtests/src/com/android/server/wm/TaskLaunchParamsModifierTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/TaskLaunchParamsModifierTests.java
@@ -32,7 +32,7 @@
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertTrue;
 import static org.mockito.ArgumentMatchers.any;
-import static com.android.dx.mockito.inline.extended.ExtendedMockito.when;
+import static org.mockito.Mockito.when;
 
 import android.app.ActivityOptions;
 import android.content.pm.ActivityInfo;
diff --git a/services/tests/wmtests/src/com/android/server/wm/WindowAnimationSpecTest.java b/services/tests/wmtests/src/com/android/server/wm/WindowAnimationSpecTest.java
index 0447e8a..21e5d99 100644
--- a/services/tests/wmtests/src/com/android/server/wm/WindowAnimationSpecTest.java
+++ b/services/tests/wmtests/src/com/android/server/wm/WindowAnimationSpecTest.java
@@ -22,8 +22,8 @@
 
 import static org.mockito.ArgumentMatchers.argThat;
 import static org.mockito.ArgumentMatchers.eq;
-import static com.android.dx.mockito.inline.extended.ExtendedMockito.mock;
-import static com.android.dx.mockito.inline.extended.ExtendedMockito.verify;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
 
 import android.graphics.Point;
 import android.graphics.Rect;
diff --git a/services/tests/wmtests/src/com/android/server/wm/WindowTestUtils.java b/services/tests/wmtests/src/com/android/server/wm/WindowTestUtils.java
index ab9a35d..3b4ab38 100644
--- a/services/tests/wmtests/src/com/android/server/wm/WindowTestUtils.java
+++ b/services/tests/wmtests/src/com/android/server/wm/WindowTestUtils.java
@@ -16,14 +16,14 @@
 
 package com.android.server.wm;
 
-import static com.android.dx.mockito.inline.extended.ExtendedMockito.any;
-import static com.android.dx.mockito.inline.extended.ExtendedMockito.anyBoolean;
-import static com.android.dx.mockito.inline.extended.ExtendedMockito.anyFloat;
-import static com.android.dx.mockito.inline.extended.ExtendedMockito.anyInt;
-import static com.android.dx.mockito.inline.extended.ExtendedMockito.doAnswer;
-import static com.android.dx.mockito.inline.extended.ExtendedMockito.doReturn;
-import static com.android.dx.mockito.inline.extended.ExtendedMockito.mock;
-import static com.android.dx.mockito.inline.extended.ExtendedMockito.when;
+import static org.mockito.Mockito.any;
+import static org.mockito.Mockito.anyBoolean;
+import static org.mockito.Mockito.anyFloat;
+import static org.mockito.Mockito.anyInt;
+import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
 
 import android.content.Context;
 import android.content.res.Configuration;
@@ -46,29 +46,34 @@
                 WallpaperController wallpaperController, DisplayWindowController controller) {
             super(display, service, wallpaperController, controller);
         }
+
+        /** Create a mocked default {@link DisplayContent}. */
+        public static TestDisplayContent create(Context context) {
+            final TestDisplayContent displayContent = mock(TestDisplayContent.class);
+            displayContent.isDefaultDisplay = true;
+
+            final DisplayPolicy displayPolicy = mock(DisplayPolicy.class);
+            when(displayPolicy.navigationBarCanMove()).thenReturn(true);
+            when(displayPolicy.hasNavigationBar()).thenReturn(true);
+
+            final DisplayRotation displayRotation = new DisplayRotation(
+                    mock(WindowManagerService.class), displayContent, displayPolicy,
+                    context, new Object());
+            displayRotation.mPortraitRotation = Surface.ROTATION_0;
+            displayRotation.mLandscapeRotation = Surface.ROTATION_90;
+            displayRotation.mUpsideDownRotation = Surface.ROTATION_180;
+            displayRotation.mSeascapeRotation = Surface.ROTATION_270;
+
+            when(displayContent.getDisplayRotation()).thenReturn(displayRotation);
+
+            return displayContent;
+        }
     }
 
-    public static DisplayRotation createTestDisplayRotation(
-            Context context, TestDisplayContent displayContent) {
-        displayContent.isDefaultDisplay = true;
-
-        final DisplayPolicy displayPolicy = mock(DisplayPolicy.class);
-        doReturn(true).when(displayPolicy).navigationBarCanMove();
-        doReturn(true).when(displayPolicy).hasNavigationBar();
-
-        final DisplayRotation displayRotation = new DisplayRotation(
-                mock(WindowManagerService.class), displayContent, displayPolicy,
-                context, new Object());
-        displayRotation.mPortraitRotation = Surface.ROTATION_0;
-        displayRotation.mLandscapeRotation = Surface.ROTATION_90;
-        displayRotation.mUpsideDownRotation = Surface.ROTATION_180;
-        displayRotation.mSeascapeRotation = Surface.ROTATION_270;
-
-        return displayRotation;
-    }
-
-    /** Creates a mock instance of {@link StackWindowController}. */
-    static StackWindowController createMockStackWindowContainerController() {
+    /**
+     * Creates a mock instance of {@link StackWindowController}.
+     */
+    public static StackWindowController createMockStackWindowContainerController() {
         StackWindowController controller = mock(StackWindowController.class);
         controller.mContainer = mock(TestTaskStack.class);
 
@@ -89,7 +94,7 @@
      * An extension of {@link TestTaskStack}, which overrides package scoped methods that would not
      * normally be mocked out.
      */
-    static class TestTaskStack extends TaskStack {
+    public static class TestTaskStack extends TaskStack {
         TestTaskStack(WindowManagerService service, int stackId) {
             super(service, stackId, null);
         }
diff --git a/services/tests/wmtests/src/com/android/server/wm/WindowTracingTest.java b/services/tests/wmtests/src/com/android/server/wm/WindowTracingTest.java
index 274c560..0445ea0 100644
--- a/services/tests/wmtests/src/com/android/server/wm/WindowTracingTest.java
+++ b/services/tests/wmtests/src/com/android/server/wm/WindowTracingTest.java
@@ -24,10 +24,10 @@
 import static org.junit.Assert.assertTrue;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.eq;
-import static com.android.dx.mockito.inline.extended.ExtendedMockito.doAnswer;
-import static com.android.dx.mockito.inline.extended.ExtendedMockito.mock;
-import static com.android.dx.mockito.inline.extended.ExtendedMockito.verify;
-import static com.android.dx.mockito.inline.extended.ExtendedMockito.verifyZeroInteractions;
+import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.verifyZeroInteractions;
 
 import android.content.Context;
 import android.platform.test.annotations.Presubmit;