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)