(2nd try) Bind an input device via descriptor

These changes solve a problem for systems that have multiple displays
as well as 2 or more input devices with the exact same name. (e.g. 2
Xbox Series X controllers will report the exact same device name).

Currently the default behavior for overall Android is that input
device inputs go towards the last screen that the user touched.
On systems where a single computer is used for displaying content to
multiple people, this can be a problem. A user may be trying to play a
game with a controller on screen 1, while a second user tries to
browse the internet on screen 2. If the user on screen 2 touches the
screen, it'll keep sending the controller input events to screen 2
instead of the game on screen 1.

To fix this, the system can use a function to bind controllers to a
display. The system can just find an input device by device name
and make sure all inputs from that controller go to the correct
display. That ensures that no matter what, the controller inputs go to
the desired game on screen 1.

Now consider a single system with 2 screens and 2 different users
trying to play controller games. User A will want their controller
inputs to go to screen 1 and User B will want their controller inputs
to go to screen 2. This is possible with the current functionality.
That is, unless the controllers have the exact same device name. In
that case the controllers can only be locked to one screen or another.

Device names are not unique. However, descriptors are unique per
device. This would allow the system to uniquely identify each input
device and bind them to the correct screen.

Descriptors also cover certain scenarios. Such as a keyboard with a
track pad built into it. It'll be 2 different input devices with 2
different device names but the same descriptor.

The original CL [1], was reverted due to the fact that renaming
`addUniqueIdAssociation` testApi caused conflict when testing
on `next` targets. This CL reverts [2] (the revert CL).

[1] Original CL: I5c65b61652e5f4946d63d8ab7309093e49e65083
[2] Revert CL:   Idd30621950c1e5499efca71daf3e80b20c4097a6

Test: Built and run on hardware (adapted to run on Android 12L)
Test: Passed all tests under atest InputManagerServiceTests with lunch sdk_phone_x86_64-eng
Test: atest FrameworksServicesTests InputTests InputManagerServiceTests
Test: atest libsurfaceflinger_unittest
Bug: 324075859

Signed-off-by: Michael Crimando <mcriman1@ford.com>
Change-Id: Ic05feb08a667acb317d14efe7786cb103895e84f
diff --git a/core/api/test-current.txt b/core/api/test-current.txt
index 443a6c0e..13647b1 100644
--- a/core/api/test-current.txt
+++ b/core/api/test-current.txt
@@ -1769,14 +1769,16 @@
   }
 
   public final class InputManager {
-    method public void addUniqueIdAssociation(@NonNull String, @NonNull String);
+    method @RequiresPermission("android.permission.ASSOCIATE_INPUT_DEVICE_TO_DISPLAY") public void addUniqueIdAssociation(@NonNull String, @NonNull String);
+    method @FlaggedApi("com.android.input.flags.device_associations") @RequiresPermission("android.permission.ASSOCIATE_INPUT_DEVICE_TO_DISPLAY") public void addUniqueIdAssociationByDescriptor(@NonNull String, @NonNull String);
     method @RequiresPermission(android.Manifest.permission.REMAP_MODIFIER_KEYS) public void clearAllModifierKeyRemappings();
     method @NonNull public java.util.List<java.lang.String> getKeyboardLayoutDescriptors();
     method @NonNull public String getKeyboardLayoutTypeForLayoutDescriptor(@NonNull String);
     method @NonNull @RequiresPermission(android.Manifest.permission.REMAP_MODIFIER_KEYS) public java.util.Map<java.lang.Integer,java.lang.Integer> getModifierKeyRemapping();
     method public int getMousePointerSpeed();
     method @RequiresPermission(android.Manifest.permission.REMAP_MODIFIER_KEYS) public void remapModifierKey(int, int);
-    method public void removeUniqueIdAssociation(@NonNull String);
+    method @RequiresPermission("android.permission.ASSOCIATE_INPUT_DEVICE_TO_DISPLAY") public void removeUniqueIdAssociation(@NonNull String);
+    method @FlaggedApi("com.android.input.flags.device_associations") @RequiresPermission("android.permission.ASSOCIATE_INPUT_DEVICE_TO_DISPLAY") public void removeUniqueIdAssociationByDescriptor(@NonNull String);
     field public static final long BLOCK_UNTRUSTED_TOUCHES = 158002302L; // 0x96aec7eL
   }
 
diff --git a/core/api/test-lint-baseline.txt b/core/api/test-lint-baseline.txt
index 685ea63..c4fe061 100644
--- a/core/api/test-lint-baseline.txt
+++ b/core/api/test-lint-baseline.txt
@@ -977,8 +977,12 @@
     Method 'setHdmiCecVersion' documentation mentions permissions already declared by @RequiresPermission
 RequiresPermission: android.hardware.input.InputManager#addUniqueIdAssociation(String, String):
     Method 'addUniqueIdAssociation' documentation mentions permissions without declaring @RequiresPermission
+RequiresPermission: android.hardware.input.InputManager#addUniqueIdAssociationByDescriptor(String, String):
+    Method 'addUniqueIdAssociationByDescriptor' documentation mentions permissions already declared by @RequiresPermission
 RequiresPermission: android.hardware.input.InputManager#removeUniqueIdAssociation(String):
     Method 'removeUniqueIdAssociation' documentation mentions permissions without declaring @RequiresPermission
+RequiresPermission: android.hardware.input.InputManager#removeUniqueIdAssociationByDescriptor(String):
+    Method 'removeUniqueIdAssociationByDescriptor' documentation mentions permissions already declared by @RequiresPermission
 RequiresPermission: android.hardware.location.GeofenceHardware#addGeofence(int, int, android.hardware.location.GeofenceHardwareRequest, android.hardware.location.GeofenceHardwareCallback):
     Method 'addGeofence' documentation mentions permissions without declaring @RequiresPermission
 RequiresPermission: android.hardware.location.GeofenceHardware#getMonitoringTypes():
diff --git a/core/java/android/hardware/input/IInputManager.aidl b/core/java/android/hardware/input/IInputManager.aidl
index 1c37aa2..243ae14 100644
--- a/core/java/android/hardware/input/IInputManager.aidl
+++ b/core/java/android/hardware/input/IInputManager.aidl
@@ -165,9 +165,16 @@
     // static association for the cleared input port will be restored.
     void removePortAssociation(in String inputPort);
 
-    // Add a runtime association between the input device and display.
+    // Add a runtime association between the input device and display, using device's descriptor.
+    void addUniqueIdAssociationByDescriptor(in String inputDeviceDescriptor,
+            in String displayUniqueId);
+    // Remove the runtime association between the input device and display, using device's
+    // descriptor.
+    void removeUniqueIdAssociationByDescriptor(in String inputDeviceDescriptor);
+
+    // Add a runtime association between the input device and display, using device's port.
     void addUniqueIdAssociation(in String inputPort, in String displayUniqueId);
-    // Remove the runtime association between the input device and display.
+    // Remove the runtime association between the input device and display, using device's port.
     void removeUniqueIdAssociation(in String inputPort);
 
     InputSensorInfo[] getSensorList(int deviceId);
diff --git a/core/java/android/hardware/input/InputManager.java b/core/java/android/hardware/input/InputManager.java
index f949158..dd4ea31 100644
--- a/core/java/android/hardware/input/InputManager.java
+++ b/core/java/android/hardware/input/InputManager.java
@@ -17,6 +17,7 @@
 package android.hardware.input;
 
 import static com.android.input.flags.Flags.FLAG_INPUT_DEVICE_VIEW_BEHAVIOR_API;
+import static com.android.input.flags.Flags.FLAG_DEVICE_ASSOCIATIONS;
 import static com.android.hardware.input.Flags.keyboardLayoutPreviewFlag;
 
 import android.Manifest;
@@ -1054,13 +1055,14 @@
     /**
      * Add a runtime association between the input port and the display port. This overrides any
      * static associations.
-     * @param inputPort The port of the input device.
-     * @param displayPort The physical port of the associated display.
+     * @param inputPort the port of the input device
+     * @param displayPort the physical port of the associated display
      * <p>
      * Requires {@link android.Manifest.permission#ASSOCIATE_INPUT_DEVICE_TO_DISPLAY}.
      * </p>
      * @hide
      */
+    @RequiresPermission(android.Manifest.permission.ASSOCIATE_INPUT_DEVICE_TO_DISPLAY)
     public void addPortAssociation(@NonNull String inputPort, int displayPort) {
         try {
             mIm.addPortAssociation(inputPort, displayPort);
@@ -1072,12 +1074,13 @@
     /**
      * Remove the runtime association between the input port and the display port. Any existing
      * static association for the cleared input port will be restored.
-     * @param inputPort The port of the input device to be cleared.
+     * @param inputPort the port of the input device to be cleared
      * <p>
      * Requires {@link android.Manifest.permission#ASSOCIATE_INPUT_DEVICE_TO_DISPLAY}.
      * </p>
      * @hide
      */
+    @RequiresPermission(android.Manifest.permission.ASSOCIATE_INPUT_DEVICE_TO_DISPLAY)
     public void removePortAssociation(@NonNull String inputPort) {
         try {
             mIm.removePortAssociation(inputPort);
@@ -1089,14 +1092,16 @@
     /**
      * Add a runtime association between the input port and display, by unique id. Input ports are
      * expected to be unique.
-     * @param inputPort The port of the input device.
-     * @param displayUniqueId The unique id of the associated display.
+     * @param inputPort the port of the input device
+     * @param displayUniqueId the unique id of the associated display
      * <p>
      * Requires {@link android.Manifest.permission#ASSOCIATE_INPUT_DEVICE_TO_DISPLAY}.
      * </p>
      * @hide
      */
+    @RequiresPermission(android.Manifest.permission.ASSOCIATE_INPUT_DEVICE_TO_DISPLAY)
     @TestApi
+    // TODO(b/324075859): Rename to addUniqueIdAssociationByPort
     public void addUniqueIdAssociation(@NonNull String inputPort,
             @NonNull String displayUniqueId) {
         mGlobal.addUniqueIdAssociation(inputPort, displayUniqueId);
@@ -1104,18 +1109,60 @@
 
     /**
      * Removes a runtime association between the input device and display.
-     * @param inputPort The port of the input device.
+     * @param inputPort the port of the input device
      * <p>
      * Requires {@link android.Manifest.permission#ASSOCIATE_INPUT_DEVICE_TO_DISPLAY}.
      * </p>
      * @hide
      */
+    @RequiresPermission(android.Manifest.permission.ASSOCIATE_INPUT_DEVICE_TO_DISPLAY)
     @TestApi
+    // TODO(b/324075859): Rename to removeUniqueIdAssociationByPort
     public void removeUniqueIdAssociation(@NonNull String inputPort) {
         mGlobal.removeUniqueIdAssociation(inputPort);
     }
 
     /**
+     * Add a runtime association between the input device name and display, by descriptor. Input
+     * device descriptors are expected to be unique per physical device, though one physical
+     * device can have multiple virtual input devices that possess the same descriptor.
+     * E.g. a keyboard with built in trackpad will be 2 different input devices with the same
+     * descriptor.
+     * @param inputDeviceDescriptor the descriptor of the input device
+     * @param displayUniqueId the unique id of the associated display
+     * <p>
+     * Requires {@link android.Manifest.permissions.ASSOCIATE_INPUT_DEVICE_TO_DISPLAY}.
+     * </p>
+     * @hide
+     */
+    @FlaggedApi(FLAG_DEVICE_ASSOCIATIONS)
+    @RequiresPermission(android.Manifest.permission.ASSOCIATE_INPUT_DEVICE_TO_DISPLAY)
+    @TestApi
+    public void addUniqueIdAssociationByDescriptor(@NonNull String inputDeviceDescriptor,
+                                                   @NonNull String displayUniqueId) {
+        mGlobal.addUniqueIdAssociationByDescriptor(inputDeviceDescriptor, displayUniqueId);
+    }
+
+    /**
+     * Removes a runtime association between the input device and display.
+    }
+
+    /**
+     * Removes a runtime association between the input device and display.
+     * @param inputDeviceDescriptor the descriptor of the input device
+     * <p>
+     * Requires {@link android.Manifest.permissions.ASSOCIATE_INPUT_DEVICE_TO_DISPLAY}.
+     * </p>
+     * @hide
+     */
+    @FlaggedApi(FLAG_DEVICE_ASSOCIATIONS)
+    @RequiresPermission(android.Manifest.permission.ASSOCIATE_INPUT_DEVICE_TO_DISPLAY)
+    @TestApi
+    public void removeUniqueIdAssociationByDescriptor(@NonNull String inputDeviceDescriptor) {
+        mGlobal.removeUniqueIdAssociationByDescriptor(inputDeviceDescriptor);
+    }
+
+    /**
      * Reports the version of the Universal Stylus Initiative (USI) protocol supported by the given
      * display, if any.
      *
diff --git a/core/java/android/hardware/input/InputManagerGlobal.java b/core/java/android/hardware/input/InputManagerGlobal.java
index 7b29666..a9c97b1 100644
--- a/core/java/android/hardware/input/InputManagerGlobal.java
+++ b/core/java/android/hardware/input/InputManagerGlobal.java
@@ -1489,6 +1489,29 @@
     }
 
     /**
+     * @see InputManager#addUniqueIdAssociationByDescriptor(String, String)
+     */
+    public void addUniqueIdAssociationByDescriptor(@NonNull String inputDeviceDescriptor,
+                                                   @NonNull String displayUniqueId) {
+        try {
+            mIm.addUniqueIdAssociationByDescriptor(inputDeviceDescriptor, displayUniqueId);
+        } catch (RemoteException e) {
+            throw e.rethrowFromSystemServer();
+        }
+    }
+
+    /**
+     * @see InputManager#removeUniqueIdAssociationByDescriptor(String)
+     */
+    public void removeUniqueIdAssociationByDescriptor(@NonNull String inputDeviceDescriptor) {
+        try {
+            mIm.removeUniqueIdAssociationByDescriptor(inputDeviceDescriptor);
+        } catch (RemoteException e) {
+            throw e.rethrowFromSystemServer();
+        }
+    }
+
+    /**
      * @see InputManager#getInputDeviceBluetoothAddress(int)
      */
     @RequiresPermission(Manifest.permission.BLUETOOTH)
diff --git a/services/core/java/com/android/server/input/InputManagerService.java b/services/core/java/com/android/server/input/InputManagerService.java
index f32c11d..8686e9a 100644
--- a/services/core/java/com/android/server/input/InputManagerService.java
+++ b/services/core/java/com/android/server/input/InputManagerService.java
@@ -250,8 +250,19 @@
     private final Object mAssociationsLock = new Object();
     @GuardedBy("mAssociationsLock")
     private final Map<String, Integer> mRuntimeAssociations = new ArrayMap<>();
+
+    // The associations of input devices to displays by port. Maps from {InputDevice#mName} (String)
+    // to {DisplayInfo#uniqueId} (String) so that events from the Input Device go to a
+    // specific display.
     @GuardedBy("mAssociationsLock")
     private final Map<String, String> mUniqueIdAssociations = new ArrayMap<>();
+
+    // The associations of input devices to displays by descriptor. Maps from
+    // {InputDevice#mDescriptor} to {DisplayInfo#uniqueId} (String) so that events from the
+    // input device go to a specific display.
+    @GuardedBy("mAssociationsLock")
+    private final Map<String, String> mUniqueIdAssociationsByDescriptor = new ArrayMap<>();
+
     // The map from input port (String) to the keyboard layout identifiers (comma separated string
     // containing language tag and layout type) associated with the corresponding keyboard device.
     // Currently only accessed by InputReader.
@@ -1741,8 +1752,8 @@
     /**
      * Add a runtime association between the input port and the display port. This overrides any
      * static associations.
-     * @param inputPort The port of the input device.
-     * @param displayPort The physical port of the associated display.
+     * @param inputPort the port of the input device
+     * @param displayPort the physical port of the associated display
      */
     @Override // Binder call
     public void addPortAssociation(@NonNull String inputPort, int displayPort) {
@@ -1763,7 +1774,7 @@
     /**
      * Remove the runtime association between the input port and the display port. Any existing
      * static association for the cleared input port will be restored.
-     * @param inputPort The port of the input device to be cleared.
+     * @param inputPort the port of the input device to be cleared
      */
     @Override // Binder call
     public void removePortAssociation(@NonNull String inputPort) {
@@ -1813,6 +1824,49 @@
         mNative.changeUniqueIdAssociation();
     }
 
+    /**
+     * Adds a runtime association between the input device descriptor and the display unique id.
+     * @param inputDeviceDescriptor the descriptor of the input device
+     * @param displayUniqueId the unique ID of the display
+     */
+    @Override // Binder call
+    public void addUniqueIdAssociationByDescriptor(@NonNull String inputDeviceDescriptor,
+                                                   @NonNull String displayUniqueId) {
+        if (!checkCallingPermission(
+                android.Manifest.permission.ASSOCIATE_INPUT_DEVICE_TO_DISPLAY,
+                "addUniqueIdAssociationByDescriptor()")) {
+            throw new SecurityException(
+                    "Requires ASSOCIATE_INPUT_DEVICE_TO_DISPLAY permission");
+        }
+
+        Objects.requireNonNull(inputDeviceDescriptor);
+        Objects.requireNonNull(displayUniqueId);
+        synchronized (mAssociationsLock) {
+            mUniqueIdAssociationsByDescriptor.put(inputDeviceDescriptor, displayUniqueId);
+        }
+        mNative.changeUniqueIdAssociation();
+    }
+
+    /**
+     * Removes the runtime association between the input device and the display.
+     * @param inputDeviceDescriptor the descriptor of the input device
+     */
+    @Override // Binder call
+    public void removeUniqueIdAssociationByDescriptor(@NonNull String inputDeviceDescriptor) {
+        if (!checkCallingPermission(
+                android.Manifest.permission.ASSOCIATE_INPUT_DEVICE_TO_DISPLAY,
+                "removeUniqueIdAssociationByDescriptor()")) {
+            throw new SecurityException(
+                    "Requires ASSOCIATE_INPUT_DEVICE_TO_DISPLAY permission");
+        }
+
+        Objects.requireNonNull(inputDeviceDescriptor);
+        synchronized (mAssociationsLock) {
+            mUniqueIdAssociationsByDescriptor.remove(inputDeviceDescriptor);
+        }
+        mNative.changeUniqueIdAssociation();
+    }
+
     void setTypeAssociationInternal(@NonNull String inputPort, @NonNull String type) {
         Objects.requireNonNull(inputPort);
         Objects.requireNonNull(type);
@@ -2190,6 +2244,13 @@
                     pw.println("  uniqueId: " + v);
                 });
             }
+            if (!mUniqueIdAssociationsByDescriptor.isEmpty()) {
+                pw.println("Unique Id Associations:");
+                mUniqueIdAssociationsByDescriptor.forEach((k, v) -> {
+                    pw.print("  descriptor: " + k);
+                    pw.println("  uniqueId: " + v);
+                });
+            }
             if (!mDeviceTypeAssociations.isEmpty()) {
                 pw.println("Type Associations:");
                 mDeviceTypeAssociations.forEach((k, v) -> {
@@ -2633,6 +2694,17 @@
 
     // Native callback
     @SuppressWarnings("unused")
+    private String[] getInputUniqueIdAssociationsByDescriptor() {
+        final Map<String, String> associations;
+        synchronized (mAssociationsLock) {
+            associations = new HashMap<>(mUniqueIdAssociationsByDescriptor);
+        }
+
+        return flatten(associations);
+    }
+
+    // Native callback
+    @SuppressWarnings("unused")
     @VisibleForTesting
     String[] getDeviceTypeAssociations() {
         final Map<String, String> associations;
diff --git a/services/core/jni/com_android_server_input_InputManagerService.cpp b/services/core/jni/com_android_server_input_InputManagerService.cpp
index 88c47f3..553f721 100644
--- a/services/core/jni/com_android_server_input_InputManagerService.cpp
+++ b/services/core/jni/com_android_server_input_InputManagerService.cpp
@@ -129,6 +129,7 @@
     jmethodID getExcludedDeviceNames;
     jmethodID getInputPortAssociations;
     jmethodID getInputUniqueIdAssociations;
+    jmethodID getInputUniqueIdAssociationsByDescriptor;
     jmethodID getDeviceTypeAssociations;
     jmethodID getKeyboardLayoutAssociations;
     jmethodID getHoverTapTimeout;
@@ -634,11 +635,15 @@
         env->DeleteLocalRef(portAssociations);
     }
 
-    outConfig->uniqueIdAssociations =
+    outConfig->uniqueIdAssociationsByPort =
             readMapFromInterleavedJavaArray<std::string>(gServiceClassInfo
                                                                  .getInputUniqueIdAssociations,
                                                          "getInputUniqueIdAssociations");
 
+    outConfig->uniqueIdAssociationsByDescriptor = readMapFromInterleavedJavaArray<
+            std::string>(gServiceClassInfo.getInputUniqueIdAssociationsByDescriptor,
+                         "getInputUniqueIdAssociationsByDescriptor");
+
     outConfig->deviceTypeAssociations =
             readMapFromInterleavedJavaArray<std::string>(gServiceClassInfo
                                                                  .getDeviceTypeAssociations,
@@ -3093,6 +3098,9 @@
     GET_METHOD_ID(gServiceClassInfo.getInputUniqueIdAssociations, clazz,
                   "getInputUniqueIdAssociations", "()[Ljava/lang/String;");
 
+    GET_METHOD_ID(gServiceClassInfo.getInputUniqueIdAssociationsByDescriptor, clazz,
+                  "getInputUniqueIdAssociationsByDescriptor", "()[Ljava/lang/String;");
+
     GET_METHOD_ID(gServiceClassInfo.getDeviceTypeAssociations, clazz, "getDeviceTypeAssociations",
                   "()[Ljava/lang/String;");
 
diff --git a/services/tests/servicestests/src/com/android/server/companion/virtual/InputManagerMockHelper.java b/services/tests/servicestests/src/com/android/server/companion/virtual/InputManagerMockHelper.java
index 74e854e4..b33a8aa 100644
--- a/services/tests/servicestests/src/com/android/server/companion/virtual/InputManagerMockHelper.java
+++ b/services/tests/servicestests/src/com/android/server/companion/virtual/InputManagerMockHelper.java
@@ -79,8 +79,9 @@
         when(mIInputManagerMock.getInputDeviceIds()).thenReturn(new int[0]);
         doAnswer(inv -> mDevices.get(inv.getArgument(0)))
                 .when(mIInputManagerMock).getInputDevice(anyInt());
-        doAnswer(inv -> mUniqueIdAssociation.put(inv.getArgument(0), inv.getArgument(1))).when(
-                mIInputManagerMock).addUniqueIdAssociation(anyString(), anyString());
+        doAnswer(inv -> mUniqueIdAssociation.put(inv.getArgument(0),
+                inv.getArgument(1))).when(mIInputManagerMock).addUniqueIdAssociation(
+                        anyString(), anyString());
         doAnswer(inv -> mUniqueIdAssociation.remove(inv.getArgument(0))).when(
                 mIInputManagerMock).removeUniqueIdAssociation(anyString());
 
diff --git a/tests/Input/src/com/android/server/input/InputManagerServiceTests.kt b/tests/Input/src/com/android/server/input/InputManagerServiceTests.kt
index f6f766a..c9c6574 100644
--- a/tests/Input/src/com/android/server/input/InputManagerServiceTests.kt
+++ b/tests/Input/src/com/android/server/input/InputManagerServiceTests.kt
@@ -19,17 +19,26 @@
 
 import android.content.Context
 import android.content.ContextWrapper
+import android.hardware.display.DisplayManager
 import android.hardware.display.DisplayViewport
+import android.hardware.display.VirtualDisplay
 import android.hardware.input.InputManager
 import android.hardware.input.InputManagerGlobal
+import android.os.InputEventInjectionSync
+import android.os.SystemClock
 import android.os.test.TestLooper
 import android.platform.test.annotations.Presubmit
 import android.platform.test.annotations.RequiresFlagsDisabled
 import android.platform.test.flag.junit.DeviceFlagsValueProvider
 import android.provider.Settings
-import android.test.mock.MockContentResolver
+import android.view.View.OnKeyListener
 import android.view.Display
+import android.view.InputDevice
+import android.view.KeyEvent
 import android.view.PointerIcon
+import android.view.SurfaceHolder
+import android.view.SurfaceView
+import android.test.mock.MockContentResolver
 import androidx.test.platform.app.InstrumentationRegistry
 import com.android.internal.util.test.FakeSettingsProvider
 import com.google.common.truth.Truth.assertThat
@@ -48,6 +57,7 @@
 import org.mockito.Mockito.`when`
 import org.mockito.Mockito.clearInvocations
 import org.mockito.Mockito.doAnswer
+import org.mockito.Mockito.mock
 import org.mockito.Mockito.never
 import org.mockito.Mockito.spy
 import org.mockito.Mockito.times
@@ -412,6 +422,175 @@
         verify(wmCallbacks).notifyPointerDisplayIdChanged(overrideDisplayId, 0f, 0f)
         thread.join(100 /*millis*/)
     }
+
+    private fun createVirtualDisplays(count: Int): List<VirtualDisplay> {
+        val displayManager: DisplayManager = context.getSystemService(
+                DisplayManager::class.java
+        ) as DisplayManager
+        val virtualDisplays = mutableListOf<VirtualDisplay>()
+        for (i in 0 until count) {
+            virtualDisplays.add(displayManager.createVirtualDisplay(
+                    /* displayName= */ "testVirtualDisplay$i",
+                    /* width= */ 100,
+                    /* height= */ 100,
+                    /* densityDpi= */ 100,
+                    /* surface= */ null,
+                    /* flags= */ 0
+            ))
+        }
+        return virtualDisplays
+    }
+
+    // Helper function that creates a KeyEvent with Keycode A with the given action
+    private fun createKeycodeAEvent(inputDevice: InputDevice, action: Int): KeyEvent {
+        val eventTime = SystemClock.uptimeMillis()
+        return KeyEvent(
+                /* downTime= */ eventTime,
+                /* eventTime= */ eventTime,
+                /* action= */ action,
+                /* code= */ KeyEvent.KEYCODE_A,
+                /* repeat= */ 0,
+                /* metaState= */ 0,
+                /* deviceId= */ inputDevice.id,
+                /* scanCode= */ 0,
+                /* flags= */ KeyEvent.FLAG_FROM_SYSTEM,
+                /* source= */ InputDevice.SOURCE_KEYBOARD
+        )
+    }
+
+    private fun createInputDevice(): InputDevice {
+        return InputDevice.Builder()
+                .setId(123)
+                .setName("abc")
+                .setDescriptor("def")
+                .setSources(InputDevice.SOURCE_KEYBOARD)
+                .build()
+    }
+
+    @Test
+    fun addUniqueIdAssociationByDescriptor_verifyAssociations() {
+        // Overall goal is to have 2 displays and verify that events from the InputDevice are
+        // sent only to the view that is on the associated display.
+        // So, associate the InputDevice with display 1, then send and verify KeyEvents.
+        // Then remove associations, then associate the InputDevice with display 2, then send
+        // and verify commands.
+
+        // Make 2 virtual displays with some mock SurfaceViews
+        val mockSurfaceView1 = mock(SurfaceView::class.java)
+        val mockSurfaceView2 = mock(SurfaceView::class.java)
+        val mockSurfaceHolder1 = mock(SurfaceHolder::class.java)
+        `when`(mockSurfaceView1.holder).thenReturn(mockSurfaceHolder1)
+        val mockSurfaceHolder2 = mock(SurfaceHolder::class.java)
+        `when`(mockSurfaceView2.holder).thenReturn(mockSurfaceHolder2)
+
+        val virtualDisplays = createVirtualDisplays(2)
+
+        // Simulate an InputDevice
+        val inputDevice = createInputDevice()
+
+        // Associate input device with display
+        service.addUniqueIdAssociationByDescriptor(
+                inputDevice.descriptor,
+                virtualDisplays[0].display.displayId.toString()
+        )
+
+        // Simulate 2 different KeyEvents
+        val downEvent = createKeycodeAEvent(inputDevice, KeyEvent.ACTION_DOWN)
+        val upEvent = createKeycodeAEvent(inputDevice, KeyEvent.ACTION_UP)
+
+        // Create a mock OnKeyListener object
+        val mockOnKeyListener = mock(OnKeyListener::class.java)
+
+        // Verify that the event went to Display 1 not Display 2
+        service.injectInputEvent(downEvent, InputEventInjectionSync.NONE)
+
+        // Call the onKey method on the mock OnKeyListener object
+        mockOnKeyListener.onKey(mockSurfaceView1, /* keyCode= */ KeyEvent.KEYCODE_A, downEvent)
+        mockOnKeyListener.onKey(mockSurfaceView2, /* keyCode= */ KeyEvent.KEYCODE_A, upEvent)
+
+        // Verify that the onKey method was called with the expected arguments
+        verify(mockOnKeyListener).onKey(mockSurfaceView1, KeyEvent.KEYCODE_A, downEvent)
+        verify(mockOnKeyListener, never()).onKey(mockSurfaceView2, KeyEvent.KEYCODE_A, downEvent)
+
+        // Remove association
+        service.removeUniqueIdAssociationByDescriptor(inputDevice.descriptor)
+
+        // Associate with Display 2
+        service.addUniqueIdAssociationByDescriptor(
+                inputDevice.descriptor,
+                virtualDisplays[1].display.displayId.toString()
+        )
+
+        // Simulate a KeyEvent
+        service.injectInputEvent(upEvent, InputEventInjectionSync.NONE)
+
+        // Verify that the event went to Display 2 not Display 1
+        verify(mockOnKeyListener).onKey(mockSurfaceView2, KeyEvent.KEYCODE_A, upEvent)
+        verify(mockOnKeyListener, never()).onKey(mockSurfaceView1, KeyEvent.KEYCODE_A, upEvent)
+    }
+
+    // TODO(b/324075859): Rename this method to addUniqueIdAssociationByPort_verifyAssociations
+    @Test
+    fun addUniqueIdAssociation_verifyAssociations() {
+        // Overall goal is to have 2 displays and verify that events from the InputDevice are
+        // sent only to the view that is on the associated display.
+        // So, associate the InputDevice with display 1, then send and verify KeyEvents.
+        // Then remove associations, then associate the InputDevice with display 2, then send
+        // and verify commands.
+
+        // Make 2 virtual displays with some mock SurfaceViews
+        val mockSurfaceView1 = mock(SurfaceView::class.java)
+        val mockSurfaceView2 = mock(SurfaceView::class.java)
+        val mockSurfaceHolder1 = mock(SurfaceHolder::class.java)
+        `when`(mockSurfaceView1.holder).thenReturn(mockSurfaceHolder1)
+        val mockSurfaceHolder2 = mock(SurfaceHolder::class.java)
+        `when`(mockSurfaceView2.holder).thenReturn(mockSurfaceHolder2)
+
+        val virtualDisplays = createVirtualDisplays(2)
+
+        // Simulate an InputDevice
+        val inputDevice = createInputDevice()
+
+        // Associate input device with display
+        service.addUniqueIdAssociation(
+                inputDevice.name,
+                virtualDisplays[0].display.displayId.toString()
+        )
+
+        // Simulate 2 different KeyEvents
+        val downEvent = createKeycodeAEvent(inputDevice, KeyEvent.ACTION_DOWN)
+        val upEvent = createKeycodeAEvent(inputDevice, KeyEvent.ACTION_UP)
+
+        // Create a mock OnKeyListener object
+        val mockOnKeyListener = mock(OnKeyListener::class.java)
+
+        // Verify that the event went to Display 1 not Display 2
+        service.injectInputEvent(downEvent, InputEventInjectionSync.NONE)
+
+        // Call the onKey method on the mock OnKeyListener object
+        mockOnKeyListener.onKey(mockSurfaceView1, /* keyCode= */ KeyEvent.KEYCODE_A, downEvent)
+        mockOnKeyListener.onKey(mockSurfaceView2, /* keyCode= */ KeyEvent.KEYCODE_A, upEvent)
+
+        // Verify that the onKey method was called with the expected arguments
+        verify(mockOnKeyListener).onKey(mockSurfaceView1, KeyEvent.KEYCODE_A, downEvent)
+        verify(mockOnKeyListener, never()).onKey(mockSurfaceView2, KeyEvent.KEYCODE_A, downEvent)
+
+        // Remove association
+        service.removeUniqueIdAssociation(inputDevice.name)
+
+        // Associate with Display 2
+        service.addUniqueIdAssociation(
+                inputDevice.name,
+                virtualDisplays[1].display.displayId.toString()
+        )
+
+        // Simulate a KeyEvent
+        service.injectInputEvent(upEvent, InputEventInjectionSync.NONE)
+
+        // Verify that the event went to Display 2 not Display 1
+        verify(mockOnKeyListener).onKey(mockSurfaceView2, KeyEvent.KEYCODE_A, upEvent)
+        verify(mockOnKeyListener, never()).onKey(mockSurfaceView1, KeyEvent.KEYCODE_A, upEvent)
+    }
 }
 
 private fun <T> whenever(methodCall: T): OngoingStubbing<T> = `when`(methodCall)