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 03b2c8c067049..59c9d8c625b52 100644
--- a/services/companion/java/com/android/server/companion/virtual/VirtualDeviceImpl.java
+++ b/services/companion/java/com/android/server/companion/virtual/VirtualDeviceImpl.java
@@ -53,8 +53,7 @@ import com.android.internal.annotations.VisibleForTesting;
 
 import java.io.FileDescriptor;
 import java.io.PrintWriter;
-import java.util.ArrayList;
-import java.util.List;
+import java.util.Set;
 
 
 final class VirtualDeviceImpl extends IVirtualDevice.Stub
@@ -69,7 +68,7 @@ final class VirtualDeviceImpl extends IVirtualDevice.Stub
     private final int mOwnerUid;
     private final InputController mInputController;
     @VisibleForTesting
-    final List<Integer> mVirtualDisplayIds = new ArrayList<>();
+    final Set<Integer> mVirtualDisplayIds = new ArraySet<>();
     private final OnDeviceCloseListener mListener;
     private final IBinder mAppToken;
     private final VirtualDeviceParams mParams;
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 0a0f7d7a0ae0b..a6b4aecf1cb60 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
@@ -23,9 +23,11 @@ import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.doCallRealMethod;
 import static org.mockito.Mockito.doNothing;
 import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
 import static org.testng.Assert.assertThrows;
 
 import android.Manifest;
+import android.app.admin.DevicePolicyManager;
 import android.companion.virtual.VirtualDeviceParams;
 import android.content.Context;
 import android.content.ContextWrapper;
@@ -73,6 +75,8 @@ public class VirtualDeviceManagerServiceTest {
     private DisplayManagerInternal mDisplayManagerInternalMock;
     @Mock
     private VirtualDeviceImpl.PendingTrampolineCallback mPendingTrampolineCallback;
+    @Mock
+    private DevicePolicyManager mDevicePolicyManagerMock;
 
     @Before
     public void setUp() {
@@ -84,6 +88,9 @@ public class VirtualDeviceManagerServiceTest {
         mContext = Mockito.spy(new ContextWrapper(InstrumentationRegistry.getTargetContext()));
         doNothing().when(mContext).enforceCallingOrSelfPermission(
                 eq(Manifest.permission.CREATE_VIRTUAL_DEVICE), anyString());
+        when(mContext.getSystemService(Context.DEVICE_POLICY_SERVICE)).thenReturn(
+                mDevicePolicyManagerMock);
+
         mInputController = new InputController(new Object(), mNativeWrapperMock);
         mDeviceImpl = new VirtualDeviceImpl(mContext,
                 /* association info */ null, new Binder(), /* uid */ 0, mInputController,
@@ -91,6 +98,14 @@ public class VirtualDeviceManagerServiceTest {
                 new VirtualDeviceParams.Builder().build());
     }
 
+    @Test
+    public void onVirtualDisplayRemovedLocked_doesNotThrowException() {
+        final int displayId = 2;
+        mDeviceImpl.onVirtualDisplayCreatedLocked(displayId);
+        // This call should not throw any exceptions.
+        mDeviceImpl.onVirtualDisplayRemovedLocked(displayId);
+    }
+
     @Test
     public void createVirtualKeyboard_noDisplay_failsSecurityException() {
         assertThrows(