Make virtual input device creation synchronous

To ensure that virtual device creation is synchronized with changing
other system parameters (such as changing the display of the mouse
pointer), we must make virtual input device creation synchronous.

This is required to ensure that the VirtualMouse#getCursorPosition API
returns the correct value as soon as the virtual mouse device is
created.

The system only holds on to a PointerController when there is an input
device that can control the pointer that is connected. This is so that
we don't need to hold on to cursor graphics resources when there's no
mouse or touchpad connected. This means we can only synchronize updating
the pointer display once we know such an input device is connected.

In this CL, we update the virtual input device creation to wait on the
binder thread until the system recognizes the new virtual input device.
We also clean up the virtual device creation logic to ensure the proper
cleanup is done if creation fails at any point.

Bug: 216792538
Test: atest VirtualDeviceManagerServiceTest InputControllerTest
Test: atest VirtualMouseTest
Change-Id: I3b893d23cacec988eb49c9fd3a79cebe31a59c05

5 files changed, 295 insertions, 64 deletions

diff --git a/services/companion/java/com/android/server/companion/virtual/InputController.java b/services/companion/java/com/android/server/companion/virtual/InputController.java
index 9d4b50be41fb..ab9023b4f0d2 100644
--- a/services/companion/java/com/android/server/companion/virtual/InputController.java
+++ b/services/companion/java/com/android/server/companion/virtual/InputController.java
@@ -22,6 +22,7 @@ import android.annotation.StringDef;
 import android.graphics.Point;
 import android.graphics.PointF;
 import android.hardware.display.DisplayManagerInternal;
+import android.hardware.input.InputDeviceIdentifier;
 import android.hardware.input.InputManager;
 import android.hardware.input.InputManagerInternal;
 import android.hardware.input.VirtualKeyEvent;
@@ -29,11 +30,13 @@ import android.hardware.input.VirtualMouseButtonEvent;
 import android.hardware.input.VirtualMouseRelativeEvent;
 import android.hardware.input.VirtualMouseScrollEvent;
 import android.hardware.input.VirtualTouchEvent;
+import android.os.Handler;
 import android.os.IBinder;
 import android.os.RemoteException;
 import android.util.ArrayMap;
 import android.util.Slog;
 import android.view.Display;
+import android.view.InputDevice;
 
 import com.android.internal.annotations.GuardedBy;
 import com.android.internal.annotations.VisibleForTesting;
@@ -44,7 +47,11 @@ import java.lang.annotation.Retention;
 import java.lang.annotation.RetentionPolicy;
 import java.util.Iterator;
 import java.util.Map;
+import java.util.Objects;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicLong;
+import java.util.function.Supplier;
 
 /** Controls virtual input devices, including device lifecycle and event dispatch. */
 class InputController {
@@ -72,20 +79,27 @@ class InputController {
     @GuardedBy("mLock")
     final Map<IBinder, InputDeviceDescriptor> mInputDeviceDescriptors = new ArrayMap<>();
 
+    private final Handler mHandler;
     private final NativeWrapper mNativeWrapper;
     private final DisplayManagerInternal mDisplayManagerInternal;
     private final InputManagerInternal mInputManagerInternal;
+    private final DeviceCreationThreadVerifier mThreadVerifier;
 
-    InputController(@NonNull Object lock) {
-        this(lock, new NativeWrapper());
+    InputController(@NonNull Object lock, @NonNull Handler handler) {
+        this(lock, new NativeWrapper(), handler,
+                // Verify that virtual devices are not created on the handler thread.
+                () -> !handler.getLooper().isCurrentThread());
     }
 
     @VisibleForTesting
-    InputController(@NonNull Object lock, @NonNull NativeWrapper nativeWrapper) {
+    InputController(@NonNull Object lock, @NonNull NativeWrapper nativeWrapper,
+            @NonNull Handler handler, @NonNull DeviceCreationThreadVerifier threadVerifier) {
         mLock = lock;
+        mHandler = handler;
         mNativeWrapper = nativeWrapper;
         mDisplayManagerInternal = LocalServices.getService(DisplayManagerInternal.class);
         mInputManagerInternal = LocalServices.getService(InputManagerInternal.class);
+        mThreadVerifier = threadVerifier;
     }
 
     void close() {
@@ -108,23 +122,13 @@ class InputController {
             @NonNull IBinder deviceToken,
             int displayId) {
         final String phys = createPhys(PHYS_TYPE_KEYBOARD);
-        setUniqueIdAssociation(displayId, phys);
-        final int fd = mNativeWrapper.openUinputKeyboard(deviceName, vendorId, productId, phys);
-        if (fd < 0) {
-            throw new RuntimeException(
-                    "A native error occurred when creating keyboard: " + -fd);
-        }
-        final BinderDeathRecipient binderDeathRecipient = new BinderDeathRecipient(deviceToken);
-        synchronized (mLock) {
-            mInputDeviceDescriptors.put(deviceToken,
-                    new InputDeviceDescriptor(fd, binderDeathRecipient,
-                            InputDeviceDescriptor.TYPE_KEYBOARD, displayId, phys));
-        }
         try {
-            deviceToken.linkToDeath(binderDeathRecipient, /* flags= */ 0);
-        } catch (RemoteException e) {
-            // TODO(b/215608394): remove and close InputDeviceDescriptor
-            throw new RuntimeException("Could not create virtual keyboard", e);
+            createDeviceInternal(InputDeviceDescriptor.TYPE_KEYBOARD, deviceName, vendorId,
+                    productId, deviceToken, displayId, phys,
+                    () -> mNativeWrapper.openUinputKeyboard(deviceName, vendorId, productId, phys));
+        } catch (DeviceCreationException e) {
+            throw new RuntimeException(
+                    "Failed to create virtual keyboard device '" + deviceName + "'.", e);
         }
     }
 
@@ -134,25 +138,15 @@ class InputController {
             @NonNull IBinder deviceToken,
             int displayId) {
         final String phys = createPhys(PHYS_TYPE_MOUSE);
-        setUniqueIdAssociation(displayId, phys);
-        final int fd = mNativeWrapper.openUinputMouse(deviceName, vendorId, productId, phys);
-        if (fd < 0) {
-            throw new RuntimeException(
-                    "A native error occurred when creating mouse: " + -fd);
-        }
-        final BinderDeathRecipient binderDeathRecipient = new BinderDeathRecipient(deviceToken);
-        synchronized (mLock) {
-            mInputDeviceDescriptors.put(deviceToken,
-                    new InputDeviceDescriptor(fd, binderDeathRecipient,
-                            InputDeviceDescriptor.TYPE_MOUSE, displayId, phys));
-            mInputManagerInternal.setVirtualMousePointerDisplayId(displayId);
-        }
         try {
-            deviceToken.linkToDeath(binderDeathRecipient, /* flags= */ 0);
-        } catch (RemoteException e) {
-            // TODO(b/215608394): remove and close InputDeviceDescriptor
-            throw new RuntimeException("Could not create virtual mouse", e);
+            createDeviceInternal(InputDeviceDescriptor.TYPE_MOUSE, deviceName, vendorId, productId,
+                    deviceToken, displayId, phys,
+                    () -> mNativeWrapper.openUinputMouse(deviceName, vendorId, productId, phys));
+        } catch (DeviceCreationException e) {
+            throw new RuntimeException(
+                    "Failed to create virtual mouse device: '" + deviceName + "'.", e);
         }
+        mInputManagerInternal.setVirtualMousePointerDisplayId(displayId);
     }
 
     void createTouchscreen(@NonNull String deviceName,
@@ -162,24 +156,14 @@ class InputController {
             int displayId,
             @NonNull Point screenSize) {
         final String phys = createPhys(PHYS_TYPE_TOUCHSCREEN);
-        setUniqueIdAssociation(displayId, phys);
-        final int fd = mNativeWrapper.openUinputTouchscreen(deviceName, vendorId, productId, phys,
-                screenSize.y, screenSize.x);
-        if (fd < 0) {
-            throw new RuntimeException(
-                    "A native error occurred when creating touchscreen: " + -fd);
-        }
-        final BinderDeathRecipient binderDeathRecipient = new BinderDeathRecipient(deviceToken);
-        synchronized (mLock) {
-            mInputDeviceDescriptors.put(deviceToken,
-                    new InputDeviceDescriptor(fd, binderDeathRecipient,
-                            InputDeviceDescriptor.TYPE_TOUCHSCREEN, displayId, phys));
-        }
         try {
-            deviceToken.linkToDeath(binderDeathRecipient, /* flags= */ 0);
-        } catch (RemoteException e) {
-            // TODO(b/215608394): remove and close InputDeviceDescriptor
-            throw new RuntimeException("Could not create virtual touchscreen", e);
+            createDeviceInternal(InputDeviceDescriptor.TYPE_TOUCHSCREEN, deviceName, vendorId,
+                    productId, deviceToken, displayId, phys,
+                    () -> mNativeWrapper.openUinputTouchscreen(deviceName, vendorId, productId,
+                            phys, screenSize.y, screenSize.x));
+        } catch (DeviceCreationException e) {
+            throw new RuntimeException(
+                    "Failed to create virtual touchscreen device '" + deviceName + "'.", e);
         }
     }
 
@@ -510,4 +494,133 @@ class InputController {
             unregisterInputDevice(mDeviceToken);
         }
     }
+
+    /** A helper class used to wait for an input device to be registered. */
+    private class WaitForDevice implements  AutoCloseable {
+        private final CountDownLatch mDeviceAddedLatch = new CountDownLatch(1);
+        private final InputManager.InputDeviceListener mListener;
+
+        WaitForDevice(String deviceName, int vendorId, int productId) {
+            mListener = new InputManager.InputDeviceListener() {
+                @Override
+                public void onInputDeviceAdded(int deviceId) {
+                    final InputDevice device = InputManager.getInstance().getInputDevice(
+                            deviceId);
+                    Objects.requireNonNull(device, "Newly added input device was null.");
+                    if (!device.getName().equals(deviceName)) {
+                        return;
+                    }
+                    final InputDeviceIdentifier id = device.getIdentifier();
+                    if (id.getVendorId() != vendorId || id.getProductId() != productId) {
+                        return;
+                    }
+                    mDeviceAddedLatch.countDown();
+                }
+
+                @Override
+                public void onInputDeviceRemoved(int deviceId) {
+
+                }
+
+                @Override
+                public void onInputDeviceChanged(int deviceId) {
+
+                }
+            };
+            InputManager.getInstance().registerInputDeviceListener(mListener, mHandler);
+        }
+
+        /** Note: This must not be called from {@link #mHandler}'s thread. */
+        void waitForDeviceCreation() throws DeviceCreationException {
+            try {
+                if (!mDeviceAddedLatch.await(1, TimeUnit.MINUTES)) {
+                    throw new DeviceCreationException(
+                            "Timed out waiting for virtual device to be created.");
+                }
+            } catch (InterruptedException e) {
+                throw new DeviceCreationException(
+                        "Interrupted while waiting for virtual device to be created.", e);
+            }
+        }
+
+        @Override
+        public void close() {
+            InputManager.getInstance().unregisterInputDeviceListener(mListener);
+        }
+    }
+
+    /** An internal exception that is thrown to indicate an error when opening a virtual device. */
+    private static class DeviceCreationException extends Exception {
+        DeviceCreationException(String message) {
+            super(message);
+        }
+        DeviceCreationException(String message, Exception cause) {
+            super(message, cause);
+        }
+    }
+
+    /**
+     * Creates a virtual input device synchronously, and waits for the notification that the device
+     * was added.
+     *
+     * Note: Input device creation is expected to happen on a binder thread, and the calling thread
+     * will be blocked until the input device creation is successful. This should not be called on
+     * the handler's thread.
+     *
+     * @throws DeviceCreationException Throws this exception if anything unexpected happens in the
+     *                                 process of creating the device. This method will take care
+     *                                 to restore the state of the system in the event of any
+     *                                 unexpected behavior.
+     */
+    private void createDeviceInternal(@InputDeviceDescriptor.Type int type, String deviceName,
+            int vendorId, int productId, IBinder deviceToken, int displayId, String phys,
+            Supplier<Integer> deviceOpener)
+            throws DeviceCreationException {
+        if (!mThreadVerifier.isValidThread()) {
+            throw new IllegalStateException(
+                    "Virtual device creation should happen on an auxiliary thread (e.g. binder "
+                            + "thread) and not from the handler's thread.");
+        }
+
+        final int fd;
+        final BinderDeathRecipient binderDeathRecipient;
+
+        setUniqueIdAssociation(displayId, phys);
+        try (WaitForDevice waiter = new WaitForDevice(deviceName, vendorId, productId)) {
+            fd = deviceOpener.get();
+            if (fd < 0) {
+                throw new DeviceCreationException(
+                        "A native error occurred when creating touchscreen: " + -fd);
+            }
+            // The fd is valid from here, so ensure that all failures close the fd after this point.
+            try {
+                waiter.waitForDeviceCreation();
+
+                binderDeathRecipient = new BinderDeathRecipient(deviceToken);
+                try {
+                    deviceToken.linkToDeath(binderDeathRecipient, /* flags= */ 0);
+                } catch (RemoteException e) {
+                    throw new DeviceCreationException(
+                            "Client died before virtual device could be created.", e);
+                }
+            } catch (DeviceCreationException e) {
+                mNativeWrapper.closeUinput(fd);
+                throw e;
+            }
+        } catch (DeviceCreationException e) {
+            InputManager.getInstance().removeUniqueIdAssociation(phys);
+            throw e;
+        }
+
+        synchronized (mLock) {
+            mInputDeviceDescriptors.put(deviceToken,
+                    new InputDeviceDescriptor(fd, binderDeathRecipient, type, displayId, phys));
+        }
+    }
+
+    @VisibleForTesting
+    interface DeviceCreationThreadVerifier {
+        /** Returns true if the calling thread is a valid thread for device creation. */
+        boolean isValidThread();
+    }
 }
diff --git a/services/companion/java/com/android/server/companion/virtual/VirtualDeviceImpl.java b/services/companion/java/com/android/server/companion/virtual/VirtualDeviceImpl.java
index de14ef61a075..9802b9783da2 100644
--- a/services/companion/java/com/android/server/companion/virtual/VirtualDeviceImpl.java
+++ b/services/companion/java/com/android/server/companion/virtual/VirtualDeviceImpl.java
@@ -166,7 +166,8 @@ final class VirtualDeviceImpl extends IVirtualDevice.Stub
         mAppToken = token;
         mParams = params;
         if (inputController == null) {
-            mInputController = new InputController(mVirtualDeviceLock);
+            mInputController = new InputController(
+                    mVirtualDeviceLock, context.getMainThreadHandler());
         } else {
             mInputController = inputController;
         }
diff --git a/services/tests/servicestests/src/com/android/server/companion/virtual/InputControllerTest.java b/services/tests/servicestests/src/com/android/server/companion/virtual/InputControllerTest.java
index b4bb04d2b1b4..92e7a86876e9 100644
--- a/services/tests/servicestests/src/com/android/server/companion/virtual/InputControllerTest.java
+++ b/services/tests/servicestests/src/com/android/server/companion/virtual/InputControllerTest.java
@@ -21,20 +21,21 @@ import static org.mockito.ArgumentMatchers.anyString;
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.doNothing;
 import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 
 import android.hardware.display.DisplayManagerInternal;
 import android.hardware.input.IInputManager;
-import android.hardware.input.InputManager;
 import android.hardware.input.InputManagerInternal;
 import android.os.Binder;
+import android.os.Handler;
 import android.os.IBinder;
 import android.platform.test.annotations.Presubmit;
+import android.testing.AndroidTestingRunner;
+import android.testing.TestableLooper;
 import android.view.Display;
 import android.view.DisplayInfo;
 
-import androidx.test.runner.AndroidJUnit4;
-
 import com.android.server.LocalServices;
 
 import org.junit.Before;
@@ -44,7 +45,8 @@ import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 
 @Presubmit
-@RunWith(AndroidJUnit4.class)
+@RunWith(AndroidTestingRunner.class)
+@TestableLooper.RunWithLooper(setAsMainLooper = true)
 public class InputControllerTest {
 
     @Mock
@@ -56,11 +58,14 @@ public class InputControllerTest {
     @Mock
     private IInputManager mIInputManagerMock;
 
+    private InputManagerMockHelper mInputManagerMockHelper;
     private InputController mInputController;
 
     @Before
     public void setUp() throws Exception {
         MockitoAnnotations.initMocks(this);
+        mInputManagerMockHelper = new InputManagerMockHelper(
+                TestableLooper.get(this), mNativeWrapperMock, mIInputManagerMock);
 
         doNothing().when(mInputManagerInternalMock).setVirtualMousePointerDisplayId(anyInt());
         LocalServices.removeServiceForTest(InputManagerInternal.class);
@@ -72,10 +77,10 @@ public class InputControllerTest {
         LocalServices.removeServiceForTest(DisplayManagerInternal.class);
         LocalServices.addService(DisplayManagerInternal.class, mDisplayManagerInternalMock);
 
-        InputManager.resetInstance(mIInputManagerMock);
-        doNothing().when(mIInputManagerMock).addUniqueIdAssociation(anyString(), anyString());
-        doNothing().when(mIInputManagerMock).removeUniqueIdAssociation(anyString());
-        mInputController = new InputController(new Object(), mNativeWrapperMock);
+        // Allow virtual devices to be created on the looper thread for testing.
+        final InputController.DeviceCreationThreadVerifier threadVerifier = () -> true;
+        mInputController = new InputController(new Object(), mNativeWrapperMock,
+                new Handler(TestableLooper.get(this).getLooper()), threadVerifier);
     }
 
     @Test
@@ -83,6 +88,7 @@ public class InputControllerTest {
         final IBinder deviceToken = new Binder();
         mInputController.createMouse("name", /*vendorId= */ 1, /*productId= */ 1, deviceToken,
                 /* displayId= */ 1);
+        verify(mNativeWrapperMock).openUinputMouse(eq("name"), eq(1), eq(1), anyString());
         verify(mInputManagerInternalMock).setVirtualMousePointerDisplayId(eq(1));
         doReturn(1).when(mInputManagerInternalMock).getVirtualMousePointerDisplayId();
         mInputController.unregisterInputDevice(deviceToken);
@@ -95,10 +101,12 @@ public class InputControllerTest {
         final IBinder deviceToken = new Binder();
         mInputController.createMouse("name", /*vendorId= */ 1, /*productId= */ 1, deviceToken,
                 /* displayId= */ 1);
+        verify(mNativeWrapperMock).openUinputMouse(eq("name"), eq(1), eq(1), anyString());
         verify(mInputManagerInternalMock).setVirtualMousePointerDisplayId(eq(1));
         final IBinder deviceToken2 = new Binder();
         mInputController.createMouse("name", /*vendorId= */ 1, /*productId= */ 1, deviceToken2,
                 /* displayId= */ 2);
+        verify(mNativeWrapperMock, times(2)).openUinputMouse(eq("name"), eq(1), eq(1), anyString());
         verify(mInputManagerInternalMock).setVirtualMousePointerDisplayId(eq(2));
         mInputController.unregisterInputDevice(deviceToken);
         verify(mInputManagerInternalMock).setVirtualMousePointerDisplayId(eq(1));
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
new file mode 100644
index 000000000000..5a6d2d398f7d
--- /dev/null
+++ b/services/tests/servicestests/src/com/android/server/companion/virtual/InputManagerMockHelper.java
@@ -0,0 +1,99 @@
+/*
+ * Copyright (C) 2022 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.server.companion.virtual;
+
+import static org.mockito.ArgumentMatchers.anyInt;
+import static org.mockito.ArgumentMatchers.anyString;
+import static org.mockito.ArgumentMatchers.notNull;
+import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.doNothing;
+import static org.mockito.Mockito.when;
+
+import android.hardware.input.IInputDevicesChangedListener;
+import android.hardware.input.IInputManager;
+import android.hardware.input.InputManager;
+import android.os.RemoteException;
+import android.testing.TestableLooper;
+import android.view.InputDevice;
+
+import org.mockito.invocation.InvocationOnMock;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Objects;
+import java.util.stream.IntStream;
+
+/**
+ * A test utility class used to share the logic for setting up {@link InputManager}'s callback for
+ * when a virtual input device being added.
+ */
+class InputManagerMockHelper {
+    private final TestableLooper mTestableLooper;
+    private final InputController.NativeWrapper mNativeWrapperMock;
+    private final IInputManager mIInputManagerMock;
+    private final List<InputDevice> mDevices = new ArrayList<>();
+    private IInputDevicesChangedListener mDevicesChangedListener;
+
+    InputManagerMockHelper(TestableLooper testableLooper,
+            InputController.NativeWrapper nativeWrapperMock, IInputManager iInputManagerMock)
+            throws Exception {
+        mTestableLooper = testableLooper;
+        mNativeWrapperMock = nativeWrapperMock;
+        mIInputManagerMock = iInputManagerMock;
+
+        doAnswer(this::handleNativeOpenInputDevice).when(mNativeWrapperMock).openUinputMouse(
+                anyString(), anyInt(), anyInt(), anyString());
+        doAnswer(this::handleNativeOpenInputDevice).when(mNativeWrapperMock).openUinputKeyboard(
+                anyString(), anyInt(), anyInt(), anyString());
+        doAnswer(this::handleNativeOpenInputDevice).when(mNativeWrapperMock).openUinputTouchscreen(
+                anyString(), anyInt(), anyInt(), anyString(), anyInt(), anyInt());
+
+        doAnswer(inv -> {
+            mDevicesChangedListener = inv.getArgument(0);
+            return null;
+        }).when(mIInputManagerMock).registerInputDevicesChangedListener(notNull());
+        when(mIInputManagerMock.getInputDeviceIds()).thenReturn(new int[0]);
+        doAnswer(inv -> mDevices.get(inv.getArgument(0)))
+                .when(mIInputManagerMock).getInputDevice(anyInt());
+        doNothing().when(mIInputManagerMock).addUniqueIdAssociation(anyString(), anyString());
+        doNothing().when(mIInputManagerMock).removeUniqueIdAssociation(anyString());
+
+        // Set a new instance of InputManager for testing that uses the IInputManager mock as the
+        // interface to the server.
+        InputManager.resetInstance(mIInputManagerMock);
+    }
+
+    private Void handleNativeOpenInputDevice(InvocationOnMock inv) {
+        Objects.requireNonNull(mDevicesChangedListener,
+                "InputController did not register an InputDevicesChangedListener.");
+        // We only use a subset of the fields of InputDevice in InputController.
+        final InputDevice device = new InputDevice(mDevices.size() /*id*/, 1 /*generation*/, 0,
+                inv.getArgument(0) /*name*/, inv.getArgument(1) /*vendorId*/,
+                inv.getArgument(2) /*productId*/, inv.getArgument(3) /*descriptor*/, true, 0, 0,
+                null, false, false, false, false, false);
+        mDevices.add(device);
+        try {
+            mDevicesChangedListener.onInputDevicesChanged(
+                    mDevices.stream().flatMapToInt(
+                            d -> IntStream.of(d.getId(), d.getGeneration())).toArray());
+        } catch (RemoteException ignored) {
+        }
+        // Process the device added notification.
+        mTestableLooper.processAllMessages();
+        return null;
+    }
+}
diff --git a/services/tests/servicestests/src/com/android/server/companion/virtual/VirtualDeviceManagerServiceTest.java b/services/tests/servicestests/src/com/android/server/companion/virtual/VirtualDeviceManagerServiceTest.java
index 808f8c2cc626..22152a1953b9 100644
--- a/services/tests/servicestests/src/com/android/server/companion/virtual/VirtualDeviceManagerServiceTest.java
+++ b/services/tests/servicestests/src/com/android/server/companion/virtual/VirtualDeviceManagerServiceTest.java
@@ -54,6 +54,7 @@ import android.content.pm.ActivityInfo;
 import android.content.pm.ApplicationInfo;
 import android.graphics.Point;
 import android.hardware.display.DisplayManagerInternal;
+import android.hardware.input.IInputManager;
 import android.hardware.input.InputManagerInternal;
 import android.hardware.input.VirtualKeyEvent;
 import android.hardware.input.VirtualMouseButtonEvent;
@@ -118,6 +119,7 @@ public class VirtualDeviceManagerServiceTest {
     private static final int FLAG_CANNOT_DISPLAY_ON_REMOTE_DEVICES = 0x00000;
 
     private Context mContext;
+    private InputManagerMockHelper mInputManagerMockHelper;
     private VirtualDeviceImpl mDeviceImpl;
     private InputController mInputController;
     private AssociationInfo mAssociationInfo;
@@ -146,6 +148,8 @@ public class VirtualDeviceManagerServiceTest {
     private IAudioConfigChangedCallback mConfigChangedCallback;
     @Mock
     private ApplicationInfo mApplicationInfoMock;
+    @Mock
+    IInputManager mIInputManagerMock;
 
     private ArraySet<ComponentName> getBlockedActivities() {
         ArraySet<ComponentName> blockedActivities = new ArraySet<>();
@@ -170,7 +174,7 @@ public class VirtualDeviceManagerServiceTest {
     }
 
     @Before
-    public void setUp() {
+    public void setUp() throws Exception {
         MockitoAnnotations.initMocks(this);
 
         LocalServices.removeServiceForTest(DisplayManagerInternal.class);
@@ -199,7 +203,13 @@ public class VirtualDeviceManagerServiceTest {
                 new Handler(TestableLooper.get(this).getLooper()));
         when(mContext.getSystemService(Context.POWER_SERVICE)).thenReturn(mPowerManager);
 
-        mInputController = new InputController(new Object(), mNativeWrapperMock);
+        mInputManagerMockHelper = new InputManagerMockHelper(
+                TestableLooper.get(this), mNativeWrapperMock, mIInputManagerMock);
+        // Allow virtual devices to be created on the looper thread for testing.
+        final InputController.DeviceCreationThreadVerifier threadVerifier = () -> true;
+        mInputController = new InputController(new Object(), mNativeWrapperMock,
+                new Handler(TestableLooper.get(this).getLooper()), threadVerifier);
+
         mAssociationInfo = new AssociationInfo(1, 0, null,
                 MacAddress.BROADCAST_ADDRESS, "", null, true, false, 0, 0);