Merge "Several changes to VDM/CDM logic." into main

9 files changed, 238 insertions, 14 deletions

diff --git a/core/api/system-current.txt b/core/api/system-current.txt
index d96925d689e2..5a8209f69dcf 100644
--- a/core/api/system-current.txt
+++ b/core/api/system-current.txt
@@ -3185,6 +3185,7 @@ package android.companion.virtual {
     field public static final int LAUNCH_FAILURE_NO_ACTIVITY = 2; // 0x2
     field public static final int LAUNCH_FAILURE_PENDING_INTENT_CANCELED = 1; // 0x1
     field public static final int LAUNCH_SUCCESS = 0; // 0x0
+    field @FlaggedApi("android.companion.virtual.flags.persistent_device_id_api") public static final String PERSISTENT_DEVICE_ID_DEFAULT = "default:0";
   }
 
   public static interface VirtualDeviceManager.ActivityListener {
diff --git a/core/java/android/companion/CompanionDeviceManager.java b/core/java/android/companion/CompanionDeviceManager.java
index 9ea3dfc61ef1..e0ce917fa3b3 100644
--- a/core/java/android/companion/CompanionDeviceManager.java
+++ b/core/java/android/companion/CompanionDeviceManager.java
@@ -33,6 +33,7 @@ import android.annotation.SystemApi;
 import android.annotation.SystemService;
 import android.annotation.TestApi;
 import android.annotation.UserHandleAware;
+import android.annotation.UserIdInt;
 import android.app.Activity;
 import android.app.ActivityManager;
 import android.app.ActivityManagerInternal;
@@ -795,9 +796,19 @@ public final class CompanionDeviceManager {
     @UserHandleAware
     @RequiresPermission(android.Manifest.permission.MANAGE_COMPANION_DEVICES)
     public @NonNull List<AssociationInfo> getAllAssociations() {
+        return getAllAssociations(mContext.getUserId());
+    }
+
+    /**
+     * Per-user version of {@link #getAllAssociations()}.
+     *
+     * @hide
+     */
+    @RequiresPermission(android.Manifest.permission.MANAGE_COMPANION_DEVICES)
+    public @NonNull List<AssociationInfo> getAllAssociations(@UserIdInt int userId) {
         if (!checkFeaturePresent()) return Collections.emptyList();
         try {
-            return mService.getAllAssociationsForUser(mContext.getUserId());
+            return mService.getAllAssociationsForUser(userId);
         } catch (RemoteException e) {
             throw e.rethrowFromSystemServer();
         }
@@ -830,12 +841,25 @@ public final class CompanionDeviceManager {
     @RequiresPermission(android.Manifest.permission.MANAGE_COMPANION_DEVICES)
     public void addOnAssociationsChangedListener(
             @NonNull Executor executor, @NonNull OnAssociationsChangedListener listener) {
+        addOnAssociationsChangedListener(executor, listener, mContext.getUserId());
+    }
+
+    /**
+     * Per-user version of
+     * {@link #addOnAssociationsChangedListener(Executor, OnAssociationsChangedListener)}.
+     *
+     * @hide
+     */
+    @RequiresPermission(android.Manifest.permission.MANAGE_COMPANION_DEVICES)
+    public void addOnAssociationsChangedListener(
+            @NonNull Executor executor, @NonNull OnAssociationsChangedListener listener,
+            @UserIdInt int userId) {
         if (!checkFeaturePresent()) return;
         synchronized (mListeners) {
             final OnAssociationsChangedListenerProxy proxy = new OnAssociationsChangedListenerProxy(
                     executor, listener);
             try {
-                mService.addOnAssociationsChangedListener(proxy, mContext.getUserId());
+                mService.addOnAssociationsChangedListener(proxy, userId);
             } catch (RemoteException e) {
                 throw e.rethrowFromSystemServer();
             }
diff --git a/core/java/android/companion/virtual/VirtualDeviceManager.java b/core/java/android/companion/virtual/VirtualDeviceManager.java
index b3ea93bb8a85..60965a84f6d2 100644
--- a/core/java/android/companion/virtual/VirtualDeviceManager.java
+++ b/core/java/android/companion/virtual/VirtualDeviceManager.java
@@ -149,6 +149,19 @@ public final class VirtualDeviceManager {
     @SystemApi
     public static final int LAUNCH_FAILURE_NO_ACTIVITY = 2;
 
+    /**
+     * Persistent device identifier corresponding to the default device.
+     *
+     * @see Context#DEVICE_ID_DEFAULT
+     * @see VirtualDevice#getPersistentDeviceId()
+     *
+     * @hide
+     */
+    @SystemApi
+    @FlaggedApi(Flags.FLAG_PERSISTENT_DEVICE_ID_API)
+    public static final String PERSISTENT_DEVICE_ID_DEFAULT =
+            "default:" + Context.DEVICE_ID_DEFAULT;
+
     private final IVirtualDeviceManager mService;
     private final Context mContext;
 
diff --git a/core/java/android/companion/virtual/flags.aconfig b/core/java/android/companion/virtual/flags.aconfig
index cfab9ebec593..02066fa8a34e 100644
--- a/core/java/android/companion/virtual/flags.aconfig
+++ b/core/java/android/companion/virtual/flags.aconfig
@@ -58,6 +58,13 @@ flag {
 }
 
 flag {
+  name: "persistent_device_id_api"
+  namespace: "virtual_devices"
+  description: "Enable persistent device ID notification API"
+  bug: "295258915"
+}
+
+flag {
   name: "express_metrics"
   namespace: "virtual_devices"
   description: "Enable express metrics in VDM"
diff --git a/services/companion/java/com/android/server/companion/CompanionDeviceManagerService.java b/services/companion/java/com/android/server/companion/CompanionDeviceManagerService.java
index 2c608930b391..b9c269c91651 100644
--- a/services/companion/java/com/android/server/companion/CompanionDeviceManagerService.java
+++ b/services/companion/java/com/android/server/companion/CompanionDeviceManagerService.java
@@ -522,7 +522,8 @@ public class CompanionDeviceManagerService extends SystemService {
     private void notifyListeners(
             @UserIdInt int userId, @NonNull List<AssociationInfo> associations) {
         mListeners.broadcast((listener, callbackUserId) -> {
-            if ((int) callbackUserId == userId) {
+            int listenerUserId = (int) callbackUserId;
+            if (listenerUserId == userId || listenerUserId == UserHandle.USER_ALL) {
                 try {
                     listener.onAssociationsChanged(associations);
                 } catch (RemoteException ignored) {
@@ -660,6 +661,9 @@ public class CompanionDeviceManagerService extends SystemService {
 
             enforceCallerIsSystemOrCanInteractWithUserId(getContext(), userId);
 
+            if (userId == UserHandle.USER_ALL) {
+                return List.copyOf(mAssociationStore.getAssociations());
+            }
             return mAssociationStore.getAssociationsForUser(userId);
         }
 
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 ae8fddfd35ef..562fe3b17f9f 100644
--- a/services/companion/java/com/android/server/companion/virtual/VirtualDeviceImpl.java
+++ b/services/companion/java/com/android/server/companion/virtual/VirtualDeviceImpl.java
@@ -304,7 +304,7 @@ final class VirtualDeviceImpl extends IVirtualDevice.Stub
         UserHandle ownerUserHandle = UserHandle.getUserHandleForUid(attributionSource.getUid());
         mContext = context.createContextAsUser(ownerUserHandle, 0);
         mAssociationInfo = associationInfo;
-        mPersistentDeviceId = PERSISTENT_ID_PREFIX_CDM_ASSOCIATION + associationInfo.getId();
+        mPersistentDeviceId = createPersistentDeviceId(associationInfo.getId());
         mService = service;
         mPendingTrampolineCallback = pendingTrampolineCallback;
         mActivityListener = activityListener;
@@ -380,6 +380,10 @@ final class VirtualDeviceImpl extends IVirtualDevice.Stub
         return mSensorController;
     }
 
+    static String createPersistentDeviceId(int associationId) {
+        return PERSISTENT_ID_PREFIX_CDM_ASSOCIATION + associationId;
+    }
+
     /**
      * Returns the flags that should be added to any virtual displays created on this virtual
      * device.
diff --git a/services/companion/java/com/android/server/companion/virtual/VirtualDeviceManagerService.java b/services/companion/java/com/android/server/companion/virtual/VirtualDeviceManagerService.java
index 85b3c9a2b33c..215970eedff1 100644
--- a/services/companion/java/com/android/server/companion/virtual/VirtualDeviceManagerService.java
+++ b/services/companion/java/com/android/server/companion/virtual/VirtualDeviceManagerService.java
@@ -27,6 +27,7 @@ import android.annotation.RequiresPermission;
 import android.annotation.SuppressLint;
 import android.app.ActivityOptions;
 import android.companion.AssociationInfo;
+import android.companion.AssociationRequest;
 import android.companion.CompanionDeviceManager;
 import android.companion.virtual.IVirtualDevice;
 import android.companion.virtual.IVirtualDeviceActivityListener;
@@ -94,6 +95,11 @@ public class VirtualDeviceManagerService extends SystemService {
 
     private static final String VIRTUAL_DEVICE_NATIVE_SERVICE = "virtualdevice_native";
 
+    private static final List<String> VIRTUAL_DEVICE_COMPANION_DEVICE_PROFILES = Arrays.asList(
+            AssociationRequest.DEVICE_PROFILE_AUTOMOTIVE_PROJECTION,
+            AssociationRequest.DEVICE_PROFILE_APP_STREAMING,
+            AssociationRequest.DEVICE_PROFILE_NEARBY_DEVICE_STREAMING);
+
     private final Object mVirtualDeviceManagerLock = new Object();
     private final VirtualDeviceManagerImpl mImpl;
     private final VirtualDeviceManagerNativeImpl mNativeImpl;
@@ -105,6 +111,9 @@ public class VirtualDeviceManagerService extends SystemService {
     private static AtomicInteger sNextUniqueIndex = new AtomicInteger(
             Context.DEVICE_ID_DEFAULT + 1);
 
+    @GuardedBy("mVirtualDeviceManagerLock")
+    private List<AssociationInfo> mActiveAssociations = new ArrayList<>();
+
     private final CompanionDeviceManager.OnAssociationsChangedListener mCdmAssociationListener =
             new CompanionDeviceManager.OnAssociationsChangedListener() {
                 @Override
@@ -161,6 +170,7 @@ public class VirtualDeviceManagerService extends SystemService {
             };
 
     @Override
+    @RequiresPermission(android.Manifest.permission.MANAGE_COMPANION_DEVICES)
     public void onStart() {
         publishBinderService(Context.VIRTUAL_DEVICE_SERVICE, mImpl);
         if (Flags.enableNativeVdm()) {
@@ -172,6 +182,21 @@ public class VirtualDeviceManagerService extends SystemService {
         activityTaskManagerInternal.registerActivityStartInterceptor(
                 VIRTUAL_DEVICE_SERVICE_ORDERED_ID,
                 mActivityInterceptorCallback);
+
+        if (Flags.persistentDeviceIdApi()) {
+            CompanionDeviceManager cdm =
+                    getContext().getSystemService(CompanionDeviceManager.class);
+            if (cdm != null) {
+                synchronized (mVirtualDeviceManagerLock) {
+                    mActiveAssociations = cdm.getAllAssociations(UserHandle.USER_ALL);
+                }
+                cdm.addOnAssociationsChangedListener(getContext().getMainExecutor(),
+                        this::onCdmAssociationsChanged, UserHandle.USER_ALL);
+            } else {
+                Slog.e(TAG, "Failed to find CompanionDeviceManager. No CDM association info "
+                        + " will be available.");
+            }
+        }
     }
 
     void onCameraAccessBlocked(int appUid) {
@@ -264,9 +289,11 @@ public class VirtualDeviceManagerService extends SystemService {
         try {
             getContext().sendBroadcastAsUser(i, UserHandle.ALL);
 
-            synchronized (mVirtualDeviceManagerLock) {
-                if (mVirtualDevices.size() == 0) {
-                    unregisterCdmAssociationListener();
+            if (!Flags.persistentDeviceIdApi()) {
+                synchronized (mVirtualDeviceManagerLock) {
+                    if (mVirtualDevices.size() == 0) {
+                        unregisterCdmAssociationListener();
+                    }
                 }
             }
         } finally {
@@ -316,6 +343,45 @@ public class VirtualDeviceManagerService extends SystemService {
         cdm.removeOnAssociationsChangedListener(mCdmAssociationListener);
     }
 
+    @RequiresPermission(android.Manifest.permission.CREATE_VIRTUAL_DEVICE)
+    void onCdmAssociationsChanged(List<AssociationInfo> associations) {
+        Set<VirtualDeviceImpl> virtualDevicesToRemove = new HashSet<>();
+        Set<String> removedPersistentDeviceIds = new HashSet<>();
+        synchronized (mVirtualDeviceManagerLock) {
+            Set<Integer> activeAssociationIds = new HashSet<>(associations.size());
+            for (int i = 0; i < associations.size(); ++i) {
+                activeAssociationIds.add(associations.get(i).getId());
+            }
+
+            for (int i = 0; i < mActiveAssociations.size(); ++i) {
+                AssociationInfo associationInfo = mActiveAssociations.get(i);
+                if (!activeAssociationIds.contains(associationInfo.getId())
+                        && VIRTUAL_DEVICE_COMPANION_DEVICE_PROFILES.contains(
+                                associationInfo.getDeviceProfile())) {
+                    removedPersistentDeviceIds.add(
+                            VirtualDeviceImpl.createPersistentDeviceId(associationInfo.getId()));
+                }
+            }
+
+            for (int i = 0; i < mVirtualDevices.size(); i++) {
+                VirtualDeviceImpl virtualDevice = mVirtualDevices.valueAt(i);
+                if (!activeAssociationIds.contains(virtualDevice.getAssociationId())) {
+                    virtualDevicesToRemove.add(virtualDevice);
+                }
+            }
+
+            mActiveAssociations = associations;
+        }
+
+        for (VirtualDeviceImpl virtualDevice : virtualDevicesToRemove) {
+            virtualDevice.close();
+        }
+
+        if (!removedPersistentDeviceIds.isEmpty()) {
+            mLocalService.onPersistentDeviceIdsRemoved(removedPersistentDeviceIds);
+        }
+    }
+
     private ArrayList<VirtualDeviceImpl> getVirtualDevicesSnapshot() {
         synchronized (mVirtualDeviceManagerLock) {
             ArrayList<VirtualDeviceImpl> virtualDevices = new ArrayList<>(mVirtualDevices.size());
@@ -393,7 +459,7 @@ public class VirtualDeviceManagerService extends SystemService {
             }
 
             synchronized (mVirtualDeviceManagerLock) {
-                if (mVirtualDevices.size() == 0) {
+                if (!Flags.persistentDeviceIdApi() && mVirtualDevices.size() == 0) {
                     final long callindId = Binder.clearCallingIdentity();
                     try {
                         registerCdmAssociationListener();
@@ -623,8 +689,12 @@ public class VirtualDeviceManagerService extends SystemService {
 
     private final class LocalService extends VirtualDeviceManagerInternal {
         @GuardedBy("mVirtualDeviceManagerLock")
-        private final ArrayList<AppsOnVirtualDeviceListener>
-                mAppsOnVirtualDeviceListeners = new ArrayList<>();
+        private final ArrayList<AppsOnVirtualDeviceListener> mAppsOnVirtualDeviceListeners =
+                new ArrayList<>();
+        @GuardedBy("mVirtualDeviceManagerLock")
+        private final ArrayList<Consumer<String>> mPersistentDeviceIdRemovedListeners =
+                new ArrayList<>();
+
         @GuardedBy("mVirtualDeviceManagerLock")
         private final ArraySet<Integer> mAllUidsOnVirtualDevice = new ArraySet<>();
 
@@ -700,6 +770,22 @@ public class VirtualDeviceManagerService extends SystemService {
         }
 
         @Override
+        public void onPersistentDeviceIdsRemoved(Set<String> removedPersistentDeviceIds) {
+            final List<Consumer<String>> persistentDeviceIdRemovedListeners;
+            synchronized (mVirtualDeviceManagerLock) {
+                persistentDeviceIdRemovedListeners = List.copyOf(
+                        mPersistentDeviceIdRemovedListeners);
+            }
+            mHandler.post(() -> {
+                for (String persistentDeviceId : removedPersistentDeviceIds) {
+                    for (Consumer<String> listener : persistentDeviceIdRemovedListeners) {
+                        listener.accept(persistentDeviceId);
+                    }
+                }
+            });
+        }
+
+        @Override
         public void onAuthenticationPrompt(int uid) {
             synchronized (mVirtualDeviceManagerLock) {
                 for (int i = 0; i < mVirtualDevices.size(); i++) {
@@ -766,6 +852,10 @@ public class VirtualDeviceManagerService extends SystemService {
 
         @Override
         public @Nullable String getPersistentIdForDevice(int deviceId) {
+            if (deviceId == Context.DEVICE_ID_DEFAULT) {
+                return VirtualDeviceManager.PERSISTENT_DEVICE_ID_DEFAULT;
+            }
+
             VirtualDeviceImpl virtualDevice;
             synchronized (mVirtualDeviceManagerLock) {
                 virtualDevice = mVirtualDevices.get(deviceId);
@@ -788,6 +878,22 @@ public class VirtualDeviceManagerService extends SystemService {
                 mAppsOnVirtualDeviceListeners.remove(listener);
             }
         }
+
+        @Override
+        public void registerPersistentDeviceIdRemovedListener(
+                @NonNull Consumer<String> persistentDeviceIdRemovedListener) {
+            synchronized (mVirtualDeviceManagerLock) {
+                mPersistentDeviceIdRemovedListeners.add(persistentDeviceIdRemovedListener);
+            }
+        }
+
+        @Override
+        public void unregisterPersistentDeviceIdRemovedListener(
+                @NonNull Consumer<String> persistentDeviceIdRemovedListener) {
+            synchronized (mVirtualDeviceManagerLock) {
+                mPersistentDeviceIdRemovedListeners.remove(persistentDeviceIdRemovedListener);
+            }
+        }
     }
 
     private static final class PendingTrampolineMap {
diff --git a/services/core/java/com/android/server/companion/virtual/VirtualDeviceManagerInternal.java b/services/core/java/com/android/server/companion/virtual/VirtualDeviceManagerInternal.java
index 0d7f778bb326..c629b2b91603 100644
--- a/services/core/java/com/android/server/companion/virtual/VirtualDeviceManagerInternal.java
+++ b/services/core/java/com/android/server/companion/virtual/VirtualDeviceManagerInternal.java
@@ -25,6 +25,7 @@ import android.os.LocaleList;
 import android.util.ArraySet;
 
 import java.util.Set;
+import java.util.function.Consumer;
 
 /**
  * Virtual device manager local service interface.
@@ -46,6 +47,14 @@ public abstract class VirtualDeviceManagerInternal {
     public abstract void unregisterAppsOnVirtualDeviceListener(
             @NonNull AppsOnVirtualDeviceListener listener);
 
+    /** Register a listener for removal of persistent device IDs. */
+    public abstract void registerPersistentDeviceIdRemovedListener(
+            @NonNull Consumer<String> persistentDeviceIdRemovedListener);
+
+    /** Unregister a listener for the removal of persistent device IDs. */
+    public abstract void unregisterPersistentDeviceIdRemovedListener(
+            @NonNull Consumer<String> persistentDeviceIdRemovedListener);
+
     /**
      * Notifies that the set of apps running on virtual devices has changed.
      * This method only notifies the listeners when the union of running UIDs on all virtual devices
@@ -59,6 +68,11 @@ public abstract class VirtualDeviceManagerInternal {
     public abstract void onAuthenticationPrompt(int uid);
 
     /**
+     * Notifies the given persistent device IDs have been removed.
+     */
+    public abstract void onPersistentDeviceIdsRemoved(Set<String> removedPersistentDeviceIds);
+
+    /**
      * Gets the owner uid for a deviceId.
      *
      * @param deviceId which device we're asking about
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 cdff62320918..f978990f0fac 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
@@ -52,9 +52,11 @@ import android.Manifest;
 import android.app.WindowConfiguration;
 import android.app.admin.DevicePolicyManager;
 import android.companion.AssociationInfo;
+import android.companion.AssociationRequest;
 import android.companion.virtual.IVirtualDeviceActivityListener;
 import android.companion.virtual.IVirtualDeviceIntentInterceptor;
 import android.companion.virtual.IVirtualDeviceSoundEffectListener;
+import android.companion.virtual.VirtualDeviceManager;
 import android.companion.virtual.VirtualDeviceParams;
 import android.companion.virtual.audio.IAudioConfigChangedCallback;
 import android.companion.virtual.audio.IAudioRoutingCallback;
@@ -134,6 +136,8 @@ import org.mockito.Mockito;
 import org.mockito.MockitoAnnotations;
 
 import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
 import java.util.Set;
 import java.util.function.Consumer;
 
@@ -261,6 +265,8 @@ public class VirtualDeviceManagerServiceTest {
     @Mock
     private VirtualDeviceManagerInternal.AppsOnVirtualDeviceListener mAppsOnVirtualDeviceListener;
     @Mock
+    private Consumer<String> mPersistentDeviceIdRemovedListener;
+    @Mock
     IPowerManager mIPowerManagerMock;
     @Mock
     IThermalService mIThermalServiceMock;
@@ -372,9 +378,8 @@ public class VirtualDeviceManagerServiceTest {
         mCameraAccessController =
                 new CameraAccessController(mContext, mLocalService, mCameraAccessBlockedCallback);
 
-        mAssociationInfo = new AssociationInfo(/* associationId= */ 1, 0, null,
-                null, MacAddress.BROADCAST_ADDRESS, "", null, null, true, false, false,
-                0, 0, -1);
+        mAssociationInfo = createAssociationInfo(
+                /* associationId= */ 1, AssociationRequest.DEVICE_PROFILE_APP_STREAMING);
 
         mVdms = new VirtualDeviceManagerService(mContext);
         mLocalService = mVdms.getLocalServiceInstance();
@@ -722,6 +727,39 @@ public class VirtualDeviceManagerServiceTest {
     }
 
     @Test
+    public void onPersistentDeviceIdsRemoved_listenersNotified() {
+        mLocalService.registerPersistentDeviceIdRemovedListener(mPersistentDeviceIdRemovedListener);
+        mLocalService.onPersistentDeviceIdsRemoved(Set.of(mDeviceImpl.getPersistentDeviceId()));
+        TestableLooper.get(this).processAllMessages();
+
+        verify(mPersistentDeviceIdRemovedListener).accept(mDeviceImpl.getPersistentDeviceId());
+    }
+
+    @Test
+    public void onCdmAssociationsChanged_persistentDeviceIdRemovedListenersNotified() {
+        mLocalService.registerPersistentDeviceIdRemovedListener(mPersistentDeviceIdRemovedListener);
+        mVdms.onCdmAssociationsChanged(List.of(mAssociationInfo));
+        TestableLooper.get(this).processAllMessages();
+
+        mVdms.onCdmAssociationsChanged(List.of(
+                createAssociationInfo(2, AssociationRequest.DEVICE_PROFILE_APP_STREAMING),
+                createAssociationInfo(3, AssociationRequest.DEVICE_PROFILE_AUTOMOTIVE_PROJECTION),
+                createAssociationInfo(4, AssociationRequest.DEVICE_PROFILE_WATCH)));
+        TestableLooper.get(this).processAllMessages();
+
+        verify(mPersistentDeviceIdRemovedListener).accept(mDeviceImpl.getPersistentDeviceId());
+
+        mVdms.onCdmAssociationsChanged(Collections.emptyList());
+        TestableLooper.get(this).processAllMessages();
+
+        verify(mPersistentDeviceIdRemovedListener)
+                .accept(VirtualDeviceImpl.createPersistentDeviceId(2));
+        verify(mPersistentDeviceIdRemovedListener)
+                .accept(VirtualDeviceImpl.createPersistentDeviceId(3));
+        verifyNoMoreInteractions(mPersistentDeviceIdRemovedListener);
+    }
+
+    @Test
     public void onAppsOnVirtualDeviceChanged_singleVirtualDevice_listenersNotified() {
         ArraySet<Integer> uids = new ArraySet<>(Arrays.asList(UID_1, UID_2));
         mLocalService.registerAppsOnVirtualDeviceListener(mAppsOnVirtualDeviceListener);
@@ -1860,11 +1898,16 @@ public class VirtualDeviceManagerServiceTest {
     @Test
     public void getPersistentIdForDevice_invalidDeviceId_returnsNull() {
         assertThat(mLocalService.getPersistentIdForDevice(DEVICE_ID_INVALID)).isNull();
-        assertThat(mLocalService.getPersistentIdForDevice(DEVICE_ID_DEFAULT)).isNull();
         assertThat(mLocalService.getPersistentIdForDevice(VIRTUAL_DEVICE_ID_2)).isNull();
     }
 
     @Test
+    public void getPersistentIdForDevice_defaultDeviceId() {
+        assertThat(mLocalService.getPersistentIdForDevice(DEVICE_ID_DEFAULT)).isEqualTo(
+                VirtualDeviceManager.PERSISTENT_DEVICE_ID_DEFAULT);
+    }
+
+    @Test
     public void getPersistentIdForDevice_returnsCorrectId() {
         assertThat(mLocalService.getPersistentIdForDevice(VIRTUAL_DEVICE_ID_1))
                 .isEqualTo(mDeviceImpl.getPersistentDeviceId());
@@ -1919,6 +1962,14 @@ public class VirtualDeviceManagerServiceTest {
         return intent.resolveActivity(packageManager);
     }
 
+    private AssociationInfo createAssociationInfo(int associationId, String deviceProfile) {
+        return new AssociationInfo(associationId, /* userId= */ 0, /* packageName=*/ null,
+                /* tag= */ null, MacAddress.BROADCAST_ADDRESS, /* displayName= */ "", deviceProfile,
+                /* associatedDevice= */ null, /* selfManaged= */ true,
+                /* notifyOnDeviceNearby= */ false, /* revoked= */false,  /* timeApprovedMs= */0,
+                /* lastTimeConnectedMs= */0, /* systemDataSyncFlags= */ -1);
+    }
+
     /** Helper class to drop permissions temporarily and restore them at the end of a test. */
     static final class DropShellPermissionsTemporarily implements AutoCloseable {
         DropShellPermissionsTemporarily() {