Merge "Enable WindowContext to associate with WindowToken" into sc-dev

12 files changed, 263 insertions, 81 deletions

diff --git a/core/java/android/app/ContextImpl.java b/core/java/android/app/ContextImpl.java
index ec7d159f707b..f7ea3815d567 100644
--- a/core/java/android/app/ContextImpl.java
+++ b/core/java/android/app/ContextImpl.java
@@ -2641,9 +2641,8 @@ class ContextImpl extends Context {
         //         WindowContext's resources.
         windowTokenClient.attachContext(windowContext);
 
-        // Step 5. Register the window context's token to the server side to associate with a
-        //         window manager node.
-        windowContext.registerWithServer();
+        // Step 5. Associate the WindowContext's token to a DisplayArea.
+        windowContext.attachToDisplayArea();
 
         return windowContext;
     }
diff --git a/core/java/android/view/IWindowManager.aidl b/core/java/android/view/IWindowManager.aidl
index d0a3e4b86e2d..0f032e906991 100644
--- a/core/java/android/view/IWindowManager.aidl
+++ b/core/java/android/view/IWindowManager.aidl
@@ -786,29 +786,60 @@ interface IWindowManager
      void setDisplayHashThrottlingEnabled(boolean enable);
 
     /**
-     * Registers a listener for a {@link android.window.WindowContext} to handle configuration
-     * changes from the server side.
+     * Attaches a {@link android.window.WindowContext} to the DisplayArea specified by {@code type},
+     * {@code displayId} and {@code options}.
      * <p>
      * Note that this API should be invoked after calling
      * {@link android.window.WindowTokenClient#attachContext(Context)}
-     * </p>
+     * </p><p>
+     * Generally, this API is used for initializing a {@link android.window.WindowContext}
+     * before obtaining a valid {@link com.android.server.wm.WindowToken}. A WindowToken is usually
+     * generated when calling {@link android.view.WindowManager#addView(View, LayoutParams)}, or
+     * obtained from {@link android.view.WindowManager.LayoutParams#token}.
+     * </p><p>
+     * In some cases, the WindowToken is passed from the server side because it is managed by the
+     * system server. {@link #attachWindowContextToWindowToken(IBinder, IBinder)} could be used in
+     * this case to attach the WindowContext to the WindowToken.</p>
      *
-     * @param clientToken the window context's token
+     * @param clientToken {@link android.window.WindowContext#getWindowContextToken()
+     * the WindowContext's token}
      * @param type Window type of the window context
      * @param displayId The display associated with the window context
      * @param options A bundle used to pass window-related options and choose the right DisplayArea
      *
-     * @return {@code true} if the listener was registered successfully.
+     * @return {@code true} if the WindowContext is attached to the DisplayArea successfully.
      */
-    boolean registerWindowContextListener(IBinder clientToken, int type, int displayId,
+    boolean attachWindowContextToDisplayArea(IBinder clientToken, int type, int displayId,
             in Bundle options);
 
     /**
-     * Unregisters a listener which registered with {@link #registerWindowContextListener()}.
+     * Attaches a {@link android.window.WindowContext} to a {@code WindowToken}.
+     * <p>
+     * This API is used when we hold a valid WindowToken and want to associate with the token and
+     * receive its configuration updates.
+     * </p><p>
+     * Note that this API should be invoked after calling
+     * {@link android.window.WindowTokenClient#attachContext(Context)}
+     * </p>
+     *
+     * @param clientToken {@link android.window.WindowContext#getWindowContextToken()
+     * the WindowContext's token}
+     * @param token the WindowToken to attach
+     *
+     * @throws IllegalArgumentException if the {@code clientToken} have not been attached to
+     * the server or the WindowContext's type doesn't match WindowToken {@code token}'s type.
+     *
+     * @see #attachWindowContextToDisplayArea(IBinder, int, int, Bundle)
+     */
+    void attachWindowContextToWindowToken(IBinder clientToken, IBinder token);
+
+    /**
+     * Detaches {@link android.window.WindowContext} from the window manager node it's currently
+     * attached to. It is no-op if the WindowContext is not attached to a window manager node.
      *
      * @param clientToken the window context's token
      */
-    void unregisterWindowContextListener(IBinder clientToken);
+    void detachWindowContextFromWindowContainer(IBinder clientToken);
 
     /**
      * Registers a listener, which is to be called whenever cross-window blur is enabled/disabled.
diff --git a/core/java/android/window/WindowContext.java b/core/java/android/window/WindowContext.java
index 375f4cf63c3b..901625b0732c 100644
--- a/core/java/android/window/WindowContext.java
+++ b/core/java/android/window/WindowContext.java
@@ -57,6 +57,8 @@ public class WindowContext extends ContextWrapper {
      *
      * @param base Base {@link Context} for this new instance.
      * @param type Window type to be used with this context.
+     * @param options A bundle used to pass window-related options.
+     *
      * @hide
      */
     public WindowContext(@NonNull Context base, int type, @Nullable Bundle options) {
@@ -72,11 +74,12 @@ public class WindowContext extends ContextWrapper {
     }
 
     /**
-     * Registers this {@link WindowContext} with {@link com.android.server.wm.WindowManagerService}
-     * to receive configuration changes of the associated {@link WindowManager} node.
+     * Attaches this {@link WindowContext} to the {@link com.android.server.wm.DisplayArea}
+     * specified by {@code mType}, {@link #getDisplayId() display ID} and {@code mOptions}
+     * to receive configuration changes.
      */
-    public void registerWithServer() {
-        mController.registerListener(mType, getDisplayId(), mOptions);
+    public void attachToDisplayArea() {
+        mController.attachToDisplayArea(mType, getDisplayId(), mOptions);
     }
 
     @Override
@@ -96,7 +99,7 @@ public class WindowContext extends ContextWrapper {
     /** Used for test to invoke because we can't invoke finalize directly. */
     @VisibleForTesting
     public void release() {
-        mController.unregisterListenerIfNeeded();
+        mController.detachIfNeeded();
         destroy();
     }
 
diff --git a/core/java/android/window/WindowContextController.java b/core/java/android/window/WindowContextController.java
index 61434145290f..88584f4b2571 100644
--- a/core/java/android/window/WindowContextController.java
+++ b/core/java/android/window/WindowContextController.java
@@ -29,22 +29,29 @@ import android.view.WindowManagerGlobal;
 import com.android.internal.annotations.VisibleForTesting;
 
 /**
- * The controller to manage {@link WindowContext} listener, such as registering and unregistering
- * the listener.
+ * The controller to manage {@link WindowContext}, such as attaching to a window manager node or
+ * detaching from the current attached node. The user must call
+ * {@link #attachToDisplayArea(int, int, Bundle)}, call {@link #attachToWindowToken(IBinder)}
+ * after that if necessary, and then call {@link #detachIfNeeded()} for release.
  *
  * @hide
  */
 public class WindowContextController {
     private final IWindowManager mWms;
+    /**
+     * {@code true} to indicate that the {@code mToken} is associated with a
+     * {@link com.android.server.wm.DisplayArea}. Note that {@code mToken} is able to attach a
+     * WindowToken after this flag sets to {@code true}.
+     */
     @VisibleForTesting
-    public boolean mListenerRegistered;
+    public boolean mAttachedToDisplayArea;
     @NonNull
     private final IBinder mToken;
 
     /**
      * Window Context Controller constructor
      *
-     * @param token The token to register to the window context listener. It is usually from
+     * @param token The token used to attach to a window manager node. It is usually from
      *              {@link Context#getWindowContextToken()}.
      */
     public WindowContextController(@NonNull IBinder token) {
@@ -60,19 +67,21 @@ public class WindowContextController {
     }
 
     /**
-     * Registers the {@code mToken} to the window context listener.
+     * Attaches the {@code mToken} to a {@link com.android.server.wm.DisplayArea}.
      *
      * @param type The window type of the {@link WindowContext}
      * @param displayId The {@link Context#getDisplayId() ID of display} to associate with
      * @param options The window context launched option
+     * @throws IllegalStateException if the {@code mToken} has already been attached to a
+     * DisplayArea.
      */
-    public void registerListener(@WindowType int type, int displayId,  @Nullable Bundle options) {
-        if (mListenerRegistered) {
-            throw new UnsupportedOperationException("A Window Context can only register a listener"
-                    + " once.");
+    public void attachToDisplayArea(@WindowType int type, int displayId, @Nullable Bundle options) {
+        if (mAttachedToDisplayArea) {
+            throw new IllegalStateException("A Window Context can be only attached to "
+                    + "a DisplayArea once.");
         }
         try {
-            mListenerRegistered = mWms.registerWindowContextListener(mToken, type, displayId,
+            mAttachedToDisplayArea = mWms.attachWindowContextToDisplayArea(mToken, type, displayId,
                     options);
         }  catch (RemoteException e) {
             throw e.rethrowFromSystemServer();
@@ -80,14 +89,42 @@ public class WindowContextController {
     }
 
     /**
-     * Unregisters the window context listener associated with the {@code mToken} if it has been
-     * registered.
+     * Switches to attach the window context to a window token.
+     * <p>
+     * Note that the context should have been attached to a
+     * {@link com.android.server.wm.DisplayArea} by {@link #attachToDisplayArea(int, int, Bundle)}
+     * before attaching to a window token, and the window token's type must match the window
+     * context's type.
+     * </p><p>
+     * A {@link WindowContext} can only attach to a specific window manager node, which is either a
+     * {@link com.android.server.wm.DisplayArea} by calling
+     * {@link #attachToDisplayArea(int, int, Bundle)} or the latest attached {@code windowToken}
+     * although this API is allowed to be called multiple times.
+     * </p>
+     * @throws IllegalStateException if the {@code mClientToken} has not yet attached to
+     * a {@link com.android.server.wm.DisplayArea} by
+     * {@link #attachToDisplayArea(int, int, Bundle)}.
+     *
+     * @see IWindowManager#attachWindowContextToWindowToken(IBinder, IBinder)
      */
-    public void unregisterListenerIfNeeded() {
-        if (mListenerRegistered) {
+    public void attachToWindowToken(IBinder windowToken) {
+        if (!mAttachedToDisplayArea) {
+            throw new IllegalStateException("The Window Context should have been attached"
+                    + " to a DisplayArea.");
+        }
+        try {
+            mWms.attachWindowContextToWindowToken(mToken, windowToken);
+        } catch (RemoteException e) {
+            throw e.rethrowFromSystemServer();
+        }
+    }
+
+    /** Detaches the window context from the node it's currently associated with. */
+    public void detachIfNeeded() {
+        if (mAttachedToDisplayArea) {
             try {
-                mWms.unregisterWindowContextListener(mToken);
-                mListenerRegistered = false;
+                mWms.detachWindowContextFromWindowContainer(mToken);
+                mAttachedToDisplayArea = false;
             } catch (RemoteException e) {
                 throw e.rethrowFromSystemServer();
             }
diff --git a/core/tests/coretests/src/android/window/WindowContextControllerTest.java b/core/tests/coretests/src/android/window/WindowContextControllerTest.java
index e4fc19c0e58d..020f4a06b892 100644
--- a/core/tests/coretests/src/android/window/WindowContextControllerTest.java
+++ b/core/tests/coretests/src/android/window/WindowContextControllerTest.java
@@ -60,34 +60,39 @@ public class WindowContextControllerTest {
         mMockWms = mock(IWindowManager.class);
         mController = new WindowContextController(new Binder(), mMockWms);
 
-        doReturn(true).when(mMockWms).registerWindowContextListener(
-                any(), anyInt(), anyInt(), any());
+        doReturn(true).when(mMockWms).attachWindowContextToDisplayArea(any(), anyInt(),
+                anyInt(), any());
     }
 
-    @Test(expected = UnsupportedOperationException.class)
-    public void testRegisterListenerTwiceThrowException() {
-        mController.registerListener(TYPE_APPLICATION_OVERLAY, DEFAULT_DISPLAY,
+    @Test(expected = IllegalStateException.class)
+    public void testAttachToDisplayAreaTwiceThrowException() {
+        mController.attachToDisplayArea(TYPE_APPLICATION_OVERLAY, DEFAULT_DISPLAY,
                 null /* options */);
-        mController.registerListener(TYPE_APPLICATION_OVERLAY, DEFAULT_DISPLAY,
+        mController.attachToDisplayArea(TYPE_APPLICATION_OVERLAY, DEFAULT_DISPLAY,
                 null /* options */);
     }
 
     @Test
-    public void testUnregisterListenerIfNeeded_NotRegisteredYet_DoNothing() throws Exception {
-        mController.unregisterListenerIfNeeded();
+    public void testDetachIfNeeded_NotAttachedYet_DoNothing() throws Exception {
+        mController.detachIfNeeded();
 
-        verify(mMockWms, never()).registerWindowContextListener(any(), anyInt(), anyInt(), any());
+        verify(mMockWms, never()).detachWindowContextFromWindowContainer(any());
     }
 
     @Test
-    public void testRegisterAndUnRegisterListener() {
-        mController.registerListener(TYPE_APPLICATION_OVERLAY, DEFAULT_DISPLAY,
+    public void testAttachAndDetachDisplayArea() {
+        mController.attachToDisplayArea(TYPE_APPLICATION_OVERLAY, DEFAULT_DISPLAY,
                 null /* options */);
 
-        assertThat(mController.mListenerRegistered).isTrue();
+        assertThat(mController.mAttachedToDisplayArea).isTrue();
 
-        mController.unregisterListenerIfNeeded();
+        mController.detachIfNeeded();
 
-        assertThat(mController.mListenerRegistered).isFalse();
+        assertThat(mController.mAttachedToDisplayArea).isFalse();
+    }
+
+    @Test(expected = IllegalStateException.class)
+    public void testAttachToWindowTokenBeforeAttachingToDAThrowException() {
+        mController.attachToWindowToken(new Binder());
     }
 }
diff --git a/data/etc/services.core.protolog.json b/data/etc/services.core.protolog.json
index efeaf238ac84..fec78f070311 100644
--- a/data/etc/services.core.protolog.json
+++ b/data/etc/services.core.protolog.json
@@ -2029,12 +2029,6 @@
       "group": "WM_DEBUG_WINDOW_ORGANIZER",
       "at": "com\/android\/server\/wm\/TaskOrganizerController.java"
     },
-    "236210101": {
-      "message": "registerWindowContextListener: trying to add listener to a non-existing display:%d",
-      "level": "WARN",
-      "group": "WM_ERROR",
-      "at": "com\/android\/server\/wm\/WindowManagerService.java"
-    },
     "240271590": {
       "message": "moveFocusableActivityToTop: unfocusable activity=%s",
       "level": "DEBUG",
@@ -2473,6 +2467,12 @@
       "group": "WM_DEBUG_SCREEN_ON",
       "at": "com\/android\/server\/wm\/WindowManagerService.java"
     },
+    "666937535": {
+      "message": "attachWindowContextToDisplayArea: trying to attach to a non-existing display:%d",
+      "level": "WARN",
+      "group": "WM_ERROR",
+      "at": "com\/android\/server\/wm\/WindowManagerService.java"
+    },
     "674932310": {
       "message": "Setting Intent of %s to target %s",
       "level": "VERBOSE",
@@ -3307,6 +3307,12 @@
       "group": "WM_DEBUG_RECENTS_ANIMATIONS",
       "at": "com\/android\/server\/wm\/RecentsAnimation.java"
     },
+    "1789321832": {
+      "message": "Then token:%s is invalid. It might be removed",
+      "level": "WARN",
+      "group": "WM_ERROR",
+      "at": "com\/android\/server\/wm\/WindowManagerService.java"
+    },
     "1789603530": {
       "message": "Removing activity %s hasSavedState=%b stateNotNeeded=%s finishing=%b state=%s callers=%s",
       "level": "INFO",
diff --git a/services/core/java/com/android/server/wm/ImmersiveModeConfirmation.java b/services/core/java/com/android/server/wm/ImmersiveModeConfirmation.java
index 567b6c2b964e..b94bc5b2c61d 100644
--- a/services/core/java/com/android/server/wm/ImmersiveModeConfirmation.java
+++ b/services/core/java/com/android/server/wm/ImmersiveModeConfirmation.java
@@ -390,7 +390,7 @@ public class ImmersiveModeConfirmation {
         final Bundle options = getOptionsForWindowContext(rootDisplayAreaId);
         final IWindowManager wms = WindowManagerGlobal.getWindowManagerService();
         try {
-            wms.registerWindowContextListener(mWindowContext.getWindowContextToken(),
+            wms.attachWindowContextToDisplayArea(mWindowContext.getWindowContextToken(),
                     IMMERSIVE_MODE_CONFIRMATION_WINDOW_TYPE, mContext.getDisplayId(), options);
         }  catch (RemoteException e) {
             throw e.rethrowAsRuntimeException();
diff --git a/services/core/java/com/android/server/wm/WindowContextListenerController.java b/services/core/java/com/android/server/wm/WindowContextListenerController.java
index 8b08314136d4..5e75996b0012 100644
--- a/services/core/java/com/android/server/wm/WindowContextListenerController.java
+++ b/services/core/java/com/android/server/wm/WindowContextListenerController.java
@@ -103,7 +103,11 @@ class WindowContextListenerController {
         listener.unregister();
     }
 
-    boolean assertCallerCanRemoveListener(IBinder clientToken, boolean callerCanManageAppTokens,
+    /**
+     * Verifies if the caller is allowed to do the operation to the listener specified by
+     * {@code clientToken}.
+     */
+    boolean assertCallerCanModifyListener(IBinder clientToken, boolean callerCanManageAppTokens,
             int callingUid) {
         final WindowContextListenerImpl listener = mListeners.get(clientToken);
         if (listener == null) {
diff --git a/services/core/java/com/android/server/wm/WindowManagerService.java b/services/core/java/com/android/server/wm/WindowManagerService.java
index ca8909a314a3..43dceeda3026 100644
--- a/services/core/java/com/android/server/wm/WindowManagerService.java
+++ b/services/core/java/com/android/server/wm/WindowManagerService.java
@@ -56,6 +56,7 @@ import static android.view.WindowManager.LayoutParams.FLAG_SECURE;
 import static android.view.WindowManager.LayoutParams.FLAG_SHOW_WALLPAPER;
 import static android.view.WindowManager.LayoutParams.FLAG_SHOW_WHEN_LOCKED;
 import static android.view.WindowManager.LayoutParams.INPUT_FEATURE_NO_INPUT_CHANNEL;
+import static android.view.WindowManager.LayoutParams.INVALID_WINDOW_TYPE;
 import static android.view.WindowManager.LayoutParams.LAST_APPLICATION_WINDOW;
 import static android.view.WindowManager.LayoutParams.LAST_SUB_WINDOW;
 import static android.view.WindowManager.LayoutParams.PRIVATE_FLAG_IS_ROUNDED_CORNERS_OVERLAY;
@@ -2700,30 +2701,19 @@ public class WindowManagerService extends IWindowManager.Stub
         }
     }
 
-    /**
-     * Registers a listener for a {@link android.window.WindowContext} to subscribe to configuration
-     * changes of a {@link DisplayArea}.
-     *
-     * @param clientToken the window context's token
-     * @param type Window type of the window context
-     * @param displayId The display associated with the window context
-     * @param options A bundle used to pass window-related options and choose the right DisplayArea
-     *
-     * @return {@code true} if the listener was registered successfully.
-     */
     @Override
-    public boolean registerWindowContextListener(IBinder clientToken, int type, int displayId,
+    public boolean attachWindowContextToDisplayArea(IBinder clientToken, int type, int displayId,
             Bundle options) {
         final boolean callerCanManageAppTokens = checkCallingPermission(MANAGE_APP_TOKENS,
-                "registerWindowContextListener", false /* printLog */);
+                "attachWindowContextToDisplayArea", false /* printLog */);
         final int callingUid = Binder.getCallingUid();
         final long origId = Binder.clearCallingIdentity();
         try {
             synchronized (mGlobalLock) {
                 final DisplayContent dc = mRoot.getDisplayContentOrCreate(displayId);
                 if (dc == null) {
-                    ProtoLog.w(WM_ERROR, "registerWindowContextListener: trying to add listener to"
-                            + " a non-existing display:%d", displayId);
+                    ProtoLog.w(WM_ERROR, "attachWindowContextToDisplayArea: trying to attach"
+                            + " to a non-existing display:%d", displayId);
                     return false;
                 }
                 // TODO(b/155340867): Investigate if we still need roundedCornerOverlay after
@@ -2740,14 +2730,50 @@ public class WindowManagerService extends IWindowManager.Stub
     }
 
     @Override
-    public void unregisterWindowContextListener(IBinder clientToken) {
+    public void attachWindowContextToWindowToken(IBinder clientToken, IBinder token) {
+        final boolean callerCanManageAppTokens = checkCallingPermission(MANAGE_APP_TOKENS,
+                "attachWindowContextToWindowToken", false /* printLog */);
+        final int callingUid = Binder.getCallingUid();
+        final long origId = Binder.clearCallingIdentity();
+        try {
+            synchronized (mGlobalLock) {
+                final WindowToken windowToken = mRoot.getWindowToken(token);
+                if (windowToken == null) {
+                    ProtoLog.w(WM_ERROR, "Then token:%s is invalid. It might be "
+                            + "removed", token);
+                    return;
+                }
+                final int type = mWindowContextListenerController.getWindowType(clientToken);
+                if (type == INVALID_WINDOW_TYPE) {
+                    throw new IllegalArgumentException("The clientToken:" + clientToken
+                            + " should have been attached.");
+                }
+                if (type != windowToken.windowType) {
+                    throw new IllegalArgumentException("The WindowToken's type should match"
+                            + " the created WindowContext's type. WindowToken's type is "
+                            + windowToken.windowType + ", while WindowContext's is " + type);
+                }
+                if (!mWindowContextListenerController.assertCallerCanModifyListener(clientToken,
+                        callerCanManageAppTokens, callingUid)) {
+                    return;
+                }
+                mWindowContextListenerController.registerWindowContainerListener(clientToken,
+                        windowToken, callingUid, windowToken.windowType, windowToken.mOptions);
+            }
+        } finally {
+            Binder.restoreCallingIdentity(origId);
+        }
+    }
+
+    @Override
+    public void detachWindowContextFromWindowContainer(IBinder clientToken) {
         final boolean callerCanManageAppTokens = checkCallingPermission(MANAGE_APP_TOKENS,
-                "unregisterWindowContextListener", false /* printLog */);
+                "detachWindowContextFromWindowContainer", false /* printLog */);
         final int callingUid = Binder.getCallingUid();
         final long origId = Binder.clearCallingIdentity();
         try {
             synchronized (mGlobalLock) {
-                if (!mWindowContextListenerController.assertCallerCanRemoveListener(clientToken,
+                if (!mWindowContextListenerController.assertCallerCanModifyListener(clientToken,
                         callerCanManageAppTokens, callingUid)) {
                     return;
                 }
diff --git a/services/tests/wmtests/src/com/android/server/wm/InputMethodMenuControllerTest.java b/services/tests/wmtests/src/com/android/server/wm/InputMethodMenuControllerTest.java
index 78074d6e1492..1b9308d32c8c 100644
--- a/services/tests/wmtests/src/com/android/server/wm/InputMethodMenuControllerTest.java
+++ b/services/tests/wmtests/src/com/android/server/wm/InputMethodMenuControllerTest.java
@@ -77,7 +77,7 @@ public class InputMethodMenuControllerTest extends WindowTestsBase {
                     dc.getImeContainer(), 1000 /* ownerUid */, TYPE_INPUT_METHOD_DIALOG,
                     null /* options */);
             return true;
-        }).when(wms).registerWindowContextListener(any(), eq(TYPE_INPUT_METHOD_DIALOG),
+        }).when(wms).attachWindowContextToDisplayArea(any(), eq(TYPE_INPUT_METHOD_DIALOG),
                 anyInt(), any());
 
         mSecondaryDisplay = new TestDisplayContent.Builder(mAtm, 1000, 1000).build();
diff --git a/services/tests/wmtests/src/com/android/server/wm/WindowContextListenerControllerTests.java b/services/tests/wmtests/src/com/android/server/wm/WindowContextListenerControllerTests.java
index 268969db184a..73b9173ba143 100644
--- a/services/tests/wmtests/src/com/android/server/wm/WindowContextListenerControllerTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/WindowContextListenerControllerTests.java
@@ -130,35 +130,35 @@ public class WindowContextListenerControllerTests extends WindowTestsBase {
     }
 
     @Test
-    public void testCanCallerRemoveListener_NullListener_ReturnFalse() {
-        assertFalse(mController.assertCallerCanRemoveListener(mClientToken,
+    public void testAssertCallerCanModifyListener_NullListener_ReturnFalse() {
+        assertFalse(mController.assertCallerCanModifyListener(mClientToken,
                 true /* callerCanManagerAppTokens */, TEST_UID));
     }
 
     @Test
-    public void testCanCallerRemoveListener_CanManageAppTokens_ReturnTrue() {
+    public void testAssertCallerCanModifyListener_CanManageAppTokens_ReturnTrue() {
         mController.registerWindowContainerListener(mClientToken, mContainer, TEST_UID,
                 TYPE_APPLICATION_OVERLAY, null /* options */);
 
-        assertTrue(mController.assertCallerCanRemoveListener(mClientToken,
+        assertTrue(mController.assertCallerCanModifyListener(mClientToken,
                 true /* callerCanManagerAppTokens */, ANOTHER_UID));
     }
 
     @Test
-    public void testCanCallerRemoveListener_SameUid_ReturnTrue() {
+    public void testAssertCallerCanModifyListener_SameUid_ReturnTrue() {
         mController.registerWindowContainerListener(mClientToken, mContainer, TEST_UID,
                 TYPE_APPLICATION_OVERLAY, null /* options */);
 
-        assertTrue(mController.assertCallerCanRemoveListener(mClientToken,
+        assertTrue(mController.assertCallerCanModifyListener(mClientToken,
                 false /* callerCanManagerAppTokens */, TEST_UID));
     }
 
     @Test(expected = UnsupportedOperationException.class)
-    public void testCanCallerRemoveListener_DifferentUid_ThrowException() {
+    public void testAssertCallerCanModifyListener_DifferentUid_ThrowException() {
         mController.registerWindowContainerListener(mClientToken, mContainer, TEST_UID,
                 TYPE_APPLICATION_OVERLAY, null /* options */);
 
-        mController.assertCallerCanRemoveListener(mClientToken,
+        mController.assertCallerCanModifyListener(mClientToken,
                 false /* callerCanManagerAppTokens */, ANOTHER_UID);
     }
 
diff --git a/services/tests/wmtests/src/com/android/server/wm/WindowManagerServiceTests.java b/services/tests/wmtests/src/com/android/server/wm/WindowManagerServiceTests.java
index baf072dbbc10..dd0c9e6d390e 100644
--- a/services/tests/wmtests/src/com/android/server/wm/WindowManagerServiceTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/WindowManagerServiceTests.java
@@ -21,7 +21,9 @@ import static android.app.WindowConfiguration.ACTIVITY_TYPE_STANDARD;
 import static android.app.WindowConfiguration.WINDOWING_MODE_FREEFORM;
 import static android.app.WindowConfiguration.WINDOWING_MODE_FULLSCREEN;
 import static android.view.Display.DEFAULT_DISPLAY;
+import static android.view.WindowManager.LayoutParams.INVALID_WINDOW_TYPE;
 import static android.view.WindowManager.LayoutParams.TYPE_APPLICATION;
+import static android.view.WindowManager.LayoutParams.TYPE_INPUT_METHOD;
 import static android.view.WindowManager.LayoutParams.TYPE_INPUT_METHOD_DIALOG;
 import static android.view.WindowManager.LayoutParams.TYPE_TOAST;
 import static android.window.DisplayAreaOrganizer.FEATURE_VENDOR_FIRST;
@@ -37,11 +39,15 @@ import static com.google.common.truth.Truth.assertThat;
 
 import static org.junit.Assert.assertFalse;
 import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyBoolean;
+import static org.mockito.ArgumentMatchers.anyInt;
 import static org.mockito.ArgumentMatchers.anyString;
+import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.verify;
 
 import android.content.pm.PackageManager;
+import android.os.Binder;
 import android.os.IBinder;
 import android.os.RemoteException;
 import android.platform.test.annotations.Presubmit;
@@ -179,4 +185,69 @@ public class WindowManagerServiceTests extends WindowTestsBase {
 
         assertThat(windowToken.getDisplayContent()).isEqualTo(mDefaultDisplay);
     }
+
+    @Test
+    public void testAttachWindowContextToWindowToken_InvalidToken_EarlyReturn() {
+        spyOn(mWm.mWindowContextListenerController);
+
+        mWm.attachWindowContextToWindowToken(new Binder(), new Binder());
+
+        verify(mWm.mWindowContextListenerController, never()).getWindowType(any());
+    }
+
+    @Test(expected = IllegalArgumentException.class)
+    public void testAttachWindowContextToWindowToken_InvalidWindowType_ThrowException() {
+        spyOn(mWm.mWindowContextListenerController);
+
+        final WindowToken windowToken = createTestWindowToken(TYPE_INPUT_METHOD, mDefaultDisplay);
+        doReturn(INVALID_WINDOW_TYPE).when(mWm.mWindowContextListenerController)
+                .getWindowType(any());
+
+        mWm.attachWindowContextToWindowToken(new Binder(), windowToken.token);
+    }
+
+    @Test(expected = IllegalArgumentException.class)
+    public void testAttachWindowContextToWindowToken_DifferentWindowType_ThrowException() {
+        spyOn(mWm.mWindowContextListenerController);
+
+        final WindowToken windowToken = createTestWindowToken(TYPE_INPUT_METHOD, mDefaultDisplay);
+        doReturn(TYPE_APPLICATION).when(mWm.mWindowContextListenerController)
+                .getWindowType(any());
+
+        mWm.attachWindowContextToWindowToken(new Binder(), windowToken.token);
+    }
+
+    @Test
+    public void testAttachWindowContextToWindowToken_CallerNotValid_EarlyReturn() {
+        spyOn(mWm.mWindowContextListenerController);
+
+        final WindowToken windowToken = createTestWindowToken(TYPE_INPUT_METHOD, mDefaultDisplay);
+        doReturn(TYPE_INPUT_METHOD).when(mWm.mWindowContextListenerController)
+                .getWindowType(any());
+        doReturn(false).when(mWm.mWindowContextListenerController)
+                .assertCallerCanModifyListener(any(), anyBoolean(), anyInt());
+
+        mWm.attachWindowContextToWindowToken(new Binder(), windowToken.token);
+
+        verify(mWm.mWindowContextListenerController, never()).registerWindowContainerListener(
+                any(), any(), anyInt(), anyInt(), any());
+    }
+
+    @Test
+    public void testAttachWindowContextToWindowToken_CallerValid_DoRegister() {
+        spyOn(mWm.mWindowContextListenerController);
+
+        final WindowToken windowToken = createTestWindowToken(TYPE_INPUT_METHOD, mDefaultDisplay);
+        doReturn(TYPE_INPUT_METHOD).when(mWm.mWindowContextListenerController)
+                .getWindowType(any());
+        doReturn(true).when(mWm.mWindowContextListenerController)
+                .assertCallerCanModifyListener(any(), anyBoolean(), anyInt());
+
+        final IBinder clientToken = new Binder();
+        mWm.attachWindowContextToWindowToken(clientToken, windowToken.token);
+
+        verify(mWm.mWindowContextListenerController).registerWindowContainerListener(
+                eq(clientToken), eq(windowToken), anyInt(), eq(TYPE_INPUT_METHOD),
+                eq(windowToken.mOptions));
+    }
 }