Fix flaky WindowTokenClientController

Although the WindowTokenClientController can be created for testing, it
was still using the app's real IWindowManager. When override the
IWindowManager in test setup, it can be accessed by the real app before
the test case run.

Allow mock IWindowManager getter to make sure there is no unexpected
access from the real app.

Fix: 293161156
Test: atest FrameworksCoreTests:WindowTokenClientControllerTest
Change-Id: I784f365ffa17dbddefcce73e4a1ed9530f071be4

5 files changed, 30 insertions, 39 deletions

diff --git a/core/java/android/view/WindowManagerGlobal.java b/core/java/android/view/WindowManagerGlobal.java
index 6aa85062562c..d20d95d50d61 100644
--- a/core/java/android/view/WindowManagerGlobal.java
+++ b/core/java/android/view/WindowManagerGlobal.java
@@ -34,7 +34,6 @@ import android.util.ArraySet;
 import android.util.Log;
 import android.view.inputmethod.InputMethodManager;
 
-import com.android.internal.annotations.VisibleForTesting;
 import com.android.internal.util.FastPrintWriter;
 
 import java.io.FileDescriptor;
@@ -165,6 +164,7 @@ public final class WindowManagerGlobal {
         }
     }
 
+    @Nullable
     @UnsupportedAppUsage
     public static IWindowManager getWindowManagerService() {
         synchronized (WindowManagerGlobal.class) {
@@ -172,6 +172,7 @@ public final class WindowManagerGlobal {
                 sWindowManagerService = IWindowManager.Stub.asInterface(
                         ServiceManager.getService("window"));
                 try {
+                    // Can be null if this is called before WindowManagerService is initialized.
                     if (sWindowManagerService != null) {
                         ValueAnimator.setDurationScale(
                                 sWindowManagerService.getCurrentAnimatorScale());
@@ -185,15 +186,6 @@ public final class WindowManagerGlobal {
         }
     }
 
-    /** Overrides the {@link #getWindowManagerService()} for test only. */
-    @VisibleForTesting
-    public static void overrideWindowManagerServiceForTesting(
-            @NonNull IWindowManager windowManager) {
-        synchronized (WindowManagerGlobal.class) {
-            sWindowManagerService = windowManager;
-        }
-    }
-
     @UnsupportedAppUsage
     public static IWindowSession getWindowSession() {
         synchronized (WindowManagerGlobal.class) {
diff --git a/core/java/android/window/WindowContextController.java b/core/java/android/window/WindowContextController.java
index 99e63ec7fe02..36eef5375391 100644
--- a/core/java/android/window/WindowContextController.java
+++ b/core/java/android/window/WindowContextController.java
@@ -86,6 +86,7 @@ public class WindowContextController {
      * @param token The token used to attach to a window manager node. It is usually from
      *              {@link Context#getWindowContextToken()}.
      */
+    @VisibleForTesting
     public WindowContextController(@NonNull WindowTokenClient token) {
         mToken = token;
     }
@@ -104,7 +105,7 @@ public class WindowContextController {
             throw new IllegalStateException("A Window Context can be only attached to "
                     + "a DisplayArea once.");
         }
-        mAttachedToDisplayArea = WindowTokenClientController.getInstance().attachToDisplayArea(
+        mAttachedToDisplayArea = getWindowTokenClientController().attachToDisplayArea(
                 mToken, type, displayId, options)
                 ? AttachStatus.STATUS_ATTACHED : AttachStatus.STATUS_FAILED;
         if (mAttachedToDisplayArea == AttachStatus.STATUS_FAILED) {
@@ -141,17 +142,24 @@ public class WindowContextController {
             throw new IllegalStateException("The Window Context should have been attached"
                     + " to a DisplayArea. AttachToDisplayArea:" + mAttachedToDisplayArea);
         }
-        WindowTokenClientController.getInstance().attachToWindowToken(mToken, windowToken);
+        getWindowTokenClientController().attachToWindowToken(mToken, windowToken);
     }
 
     /** Detaches the window context from the node it's currently associated with. */
     public void detachIfNeeded() {
         if (mAttachedToDisplayArea == AttachStatus.STATUS_ATTACHED) {
-            WindowTokenClientController.getInstance().detachIfNeeded(mToken);
+            getWindowTokenClientController().detachIfNeeded(mToken);
             mAttachedToDisplayArea = AttachStatus.STATUS_DETACHED;
             if (DEBUG_ATTACH) {
                 Log.d(TAG, "Detach Window Context.");
             }
         }
     }
+
+    /** Gets the {@link WindowTokenClientController}. */
+    @VisibleForTesting
+    @NonNull
+    public WindowTokenClientController getWindowTokenClientController() {
+        return WindowTokenClientController.getInstance();
+    }
 }
diff --git a/core/java/android/window/WindowTokenClientController.java b/core/java/android/window/WindowTokenClientController.java
index 44847073646d..14b9df65aa69 100644
--- a/core/java/android/window/WindowTokenClientController.java
+++ b/core/java/android/window/WindowTokenClientController.java
@@ -17,7 +17,6 @@
 package android.window;
 
 import static android.view.WindowManager.LayoutParams.WindowType;
-import static android.view.WindowManagerGlobal.getWindowManagerService;
 
 import android.annotation.NonNull;
 import android.annotation.Nullable;
@@ -33,6 +32,7 @@ import android.os.RemoteException;
 import android.util.ArrayMap;
 import android.util.Log;
 import android.view.IWindowManager;
+import android.view.WindowManagerGlobal;
 
 import com.android.internal.annotations.GuardedBy;
 import com.android.internal.annotations.VisibleForTesting;
@@ -56,6 +56,7 @@ public class WindowTokenClientController {
     private final ArrayMap<IBinder, WindowTokenClient> mWindowTokenClientMap = new ArrayMap<>();
 
     /** Gets the singleton controller. */
+    @NonNull
     public static WindowTokenClientController getInstance() {
         synchronized (WindowTokenClientController.class) {
             if (sController == null) {
@@ -75,6 +76,7 @@ public class WindowTokenClientController {
 
     /** Creates a new instance for test only. */
     @VisibleForTesting
+    @NonNull
     public static WindowTokenClientController createInstanceForTesting() {
         return new WindowTokenClientController();
     }
@@ -203,4 +205,11 @@ public class WindowTokenClientController {
         }
         return windowTokenClient;
     }
+
+    /** Gets the {@link IWindowManager}. */
+    @VisibleForTesting
+    @Nullable
+    public IWindowManager getWindowManagerService() {
+        return WindowManagerGlobal.getWindowManagerService();
+    }
 }
diff --git a/core/tests/coretests/src/android/window/WindowContextControllerTest.java b/core/tests/coretests/src/android/window/WindowContextControllerTest.java
index 940c06794b43..30c0f2b682db 100644
--- a/core/tests/coretests/src/android/window/WindowContextControllerTest.java
+++ b/core/tests/coretests/src/android/window/WindowContextControllerTest.java
@@ -28,6 +28,7 @@ import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.doNothing;
 import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.verify;
 
 import android.os.Binder;
@@ -36,7 +37,6 @@ import android.platform.test.annotations.Presubmit;
 import androidx.test.filters.SmallTest;
 import androidx.test.runner.AndroidJUnit4;
 
-import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -62,24 +62,16 @@ public class WindowContextControllerTest {
     @Mock
     private WindowTokenClient mMockToken;
 
-    private WindowTokenClientController mOriginalController;
-
     @Before
     public void setUp() throws Exception {
         MockitoAnnotations.initMocks(this);
-        mController = new WindowContextController(mMockToken);
+        mController = spy(new WindowContextController(mMockToken));
+        doReturn(mWindowTokenClientController).when(mController).getWindowTokenClientController();
         doNothing().when(mMockToken).onConfigurationChanged(any(), anyInt(), anyBoolean());
-        mOriginalController = WindowTokenClientController.getInstance();
-        WindowTokenClientController.overrideForTesting(mWindowTokenClientController);
         doReturn(true).when(mWindowTokenClientController).attachToDisplayArea(
                 eq(mMockToken), anyInt(), anyInt(), any());
     }
 
-    @After
-    public void tearDown() {
-        WindowTokenClientController.overrideForTesting(mOriginalController);
-    }
-
     @Test(expected = IllegalStateException.class)
     public void testAttachToDisplayAreaTwiceThrowException() {
         mController.attachToDisplayArea(TYPE_APPLICATION_OVERLAY, DEFAULT_DISPLAY,
diff --git a/core/tests/coretests/src/android/window/WindowTokenClientControllerTest.java b/core/tests/coretests/src/android/window/WindowTokenClientControllerTest.java
index 767dd8c20b77..c4f7b3b5cd59 100644
--- a/core/tests/coretests/src/android/window/WindowTokenClientControllerTest.java
+++ b/core/tests/coretests/src/android/window/WindowTokenClientControllerTest.java
@@ -35,11 +35,9 @@ import android.os.IBinder;
 import android.os.RemoteException;
 import android.platform.test.annotations.Presubmit;
 import android.view.IWindowManager;
-import android.view.WindowManagerGlobal;
 
 import androidx.test.filters.SmallTest;
 
-import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
 import org.mockito.Mock;
@@ -66,22 +64,14 @@ public class WindowTokenClientControllerTest {
     // Can't mock final class.
     private final Configuration mConfiguration = new Configuration();
 
-    private IWindowManager mOriginalWindowManagerService;
-
     private WindowTokenClientController mController;
 
     @Before
     public void setup() {
         MockitoAnnotations.initMocks(this);
-        mOriginalWindowManagerService = WindowManagerGlobal.getWindowManagerService();
-        WindowManagerGlobal.overrideWindowManagerServiceForTesting(mWindowManagerService);
         doReturn(mClientToken).when(mWindowTokenClient).asBinder();
         mController = spy(WindowTokenClientController.createInstanceForTesting());
-    }
-
-    @After
-    public void tearDown() {
-        WindowManagerGlobal.overrideWindowManagerServiceForTesting(mOriginalWindowManagerService);
+        doReturn(mWindowManagerService).when(mController).getWindowManagerService();
     }
 
     @Test
@@ -125,7 +115,7 @@ public class WindowTokenClientControllerTest {
                 DEFAULT_DISPLAY, null /* options */);
         mController.detachIfNeeded(mWindowTokenClient);
 
-        verify(mWindowManagerService).detachWindowContext(any());
+        verify(mWindowManagerService).detachWindowContext(mWindowTokenClient);
     }
 
     @Test
@@ -165,7 +155,7 @@ public class WindowTokenClientControllerTest {
         mController.attachToDisplayContent(mWindowTokenClient, DEFAULT_DISPLAY);
         mController.detachIfNeeded(mWindowTokenClient);
 
-        verify(mWindowManagerService).detachWindowContext(any());
+        verify(mWindowManagerService).detachWindowContext(mWindowTokenClient);
     }
 
     @Test
@@ -186,7 +176,7 @@ public class WindowTokenClientControllerTest {
         mController.attachToWindowToken(mWindowTokenClient, mWindowToken);
         mController.detachIfNeeded(mWindowTokenClient);
 
-        verify(mWindowManagerService).detachWindowContext(any());
+        verify(mWindowManagerService).detachWindowContext(mWindowTokenClient);
     }
 
     @Test