Modify onProfileActionStatusChange and add tests.

Modified onProfileActionStatusChange in UserManagerState to make it more
generalised for all kinds of profiles that may come up in the future.
Added tests for getUserIdToLabelMap and getUserIdToBadgeMap.

Bug: 321894374
Test: atest DocumentsUIGoogleTests:UserManagerStateTest
Change-Id: Id2e9f668b3746cee3475d1faf824d40cef1e974b
diff --git a/src/com/android/documentsui/BaseActivity.java b/src/com/android/documentsui/BaseActivity.java
index 4228ead..1e9c5fd 100644
--- a/src/com/android/documentsui/BaseActivity.java
+++ b/src/com/android/documentsui/BaseActivity.java
@@ -140,7 +140,7 @@
 
     @VisibleForTesting
     protected void initConfigStore() {
-        mConfigStore = DocumentsApplication.getConfigStore(this);
+        mConfigStore = DocumentsApplication.getConfigStore();
     }
 
     @VisibleForTesting
@@ -287,6 +287,12 @@
 
         mUserIdManager = DocumentsApplication.getUserIdManager(this);
         mUserManagerState = DocumentsApplication.getUserManagerState(this);
+        // If private space feature flag is enabled, we should store the intent that launched docsUi
+        // so that we can use this intent to get CrossProfileResolveInfo when ever we want to,
+        // for example when ACTION_PROFILE_AVAILABLE intent is received
+        if (mUserManagerState != null) {
+            mUserManagerState.setCurrentStateIntent(intent);
+        }
         mSearchManager = new SearchViewManager(searchListener, queryInterceptor,
                 chipGroup, savedInstanceState);
         // initialize the chip sets by accept mime types
diff --git a/src/com/android/documentsui/DocumentsApplication.java b/src/com/android/documentsui/DocumentsApplication.java
index 92aa3f7..60fbdaf 100644
--- a/src/com/android/documentsui/DocumentsApplication.java
+++ b/src/com/android/documentsui/DocumentsApplication.java
@@ -85,7 +85,7 @@
     public static ProvidersCache getProvidersCache(Context context) {
         ProvidersCache providers =
                 ((DocumentsApplication) context.getApplicationContext()).mProviders;
-        final ConfigStore configStore = getConfigStore(context);
+        final ConfigStore configStore = getConfigStore();
         // When private space in DocsUI is enabled then ProvidersCache should use UserManagerState
         // else it should use UserIdManager. The following if-check will ensure the construction of
         // a new ProvidersCache instance whenever there is a mismatch in this.
@@ -140,7 +140,7 @@
     public static UserManagerState getUserManagerState(Context context) {
         UserManagerState userManagerState =
                 ((DocumentsApplication) context.getApplicationContext()).mUserManagerState;
-        if (userManagerState == null && getConfigStore(context).isPrivateSpaceInDocsUIEnabled()) {
+        if (userManagerState == null && getConfigStore().isPrivateSpaceInDocsUIEnabled()) {
             userManagerState = UserManagerState.create(context);
             ((DocumentsApplication) context.getApplicationContext()).mUserManagerState =
                     userManagerState;
@@ -159,7 +159,7 @@
     /**
      * Retrieve {@link ConfigStore} instance to access feature flags in production code.
      */
-    public static synchronized ConfigStore getConfigStore(Context context) {
+    public static synchronized ConfigStore getConfigStore() {
         if (sConfigStore == null) {
             sConfigStore = new ConfigStore.ConfigStoreImpl();
         }
@@ -167,14 +167,6 @@
     }
 
     /**
-     * Set {@link #mProviders} as null onDestroy of BaseActivity so that new session uses new
-     * instance of {@link #mProviders}
-     */
-    public static void invalidateProvidersCache(Context context) {
-        ((DocumentsApplication) context.getApplicationContext()).mProviders = null;
-    }
-
-    /**
      * Set {@link #mUserManagerState} as null onDestroy of BaseActivity so that new session uses new
      * instance of {@link #mUserManagerState}
      */
@@ -217,17 +209,17 @@
             Log.w(TAG, "Can't obtain OverlayManager from System Service!");
         }
 
-        if (getConfigStore(this).isPrivateSpaceInDocsUIEnabled()) {
+        if (getConfigStore().isPrivateSpaceInDocsUIEnabled()) {
             mUserManagerState = UserManagerState.create(this);
             mUserIdManager = null;
             synchronized (DocumentsApplication.class) {
-                mProviders = new ProvidersCache(this, mUserManagerState, getConfigStore(this));
+                mProviders = new ProvidersCache(this, mUserManagerState, getConfigStore());
             }
         } else {
             mUserManagerState = null;
             mUserIdManager = UserIdManager.create(this);
             synchronized (DocumentsApplication.class) {
-                mProviders = new ProvidersCache(this, mUserIdManager, getConfigStore(this));
+                mProviders = new ProvidersCache(this, mUserIdManager, getConfigStore());
             }
         }
 
@@ -290,7 +282,7 @@
             } else if (PROFILE_FILTER_ACTIONS.contains(action)) {
                 // After we have reloaded roots. Resend the broadcast locally so the other
                 // components can reload properly after roots are updated.
-                if (getConfigStore(context).isPrivateSpaceInDocsUIEnabled()) {
+                if (getConfigStore().isPrivateSpaceInDocsUIEnabled()) {
                     UserHandle userHandle = intent.getParcelableExtra(Intent.EXTRA_USER);
                     UserId userId = UserId.of(userHandle);
                     getUserManagerState(context).onProfileActionStatusChange(action, userId);
diff --git a/src/com/android/documentsui/UserManagerState.java b/src/com/android/documentsui/UserManagerState.java
index 28a83e8..eccc6b0 100644
--- a/src/com/android/documentsui/UserManagerState.java
+++ b/src/com/android/documentsui/UserManagerState.java
@@ -95,6 +95,11 @@
     void onProfileActionStatusChange(String action, UserId userId);
 
     /**
+     * Sets the intent that triggered the launch of the DocsUI
+     */
+    void setCurrentStateIntent(Intent intent);
+
+    /**
      * Creates an implementation of {@link UserManagerState}.
      */
     // TODO: b/314746383 Make this class a singleton
@@ -136,6 +141,7 @@
         @GuardedBy("mCanFrowardToProfileIdMap")
         private final Map<UserId, Boolean> mCanFrowardToProfileIdMap = new HashMap<>();
 
+        private Intent mCurrentStateIntent;
 
         private final BroadcastReceiver mIntentReceiver = new BroadcastReceiver() {
             @Override
@@ -159,7 +165,7 @@
         private RuntimeUserManagerState(Context context) {
             this(context, UserId.CURRENT_USER,
                     Features.CROSS_PROFILE_TABS && isDeviceSupported(context),
-                    DocumentsApplication.getConfigStore(context));
+                    DocumentsApplication.getConfigStore());
         }
 
         @VisibleForTesting
@@ -185,7 +191,6 @@
         public List<UserId> getUserIds() {
             synchronized (mUserIds) {
                 if (mUserIds.isEmpty()) {
-                    Log.d("profileAction", "user ids empty");
                     mUserIds.addAll(getUserIdsInternal());
                 }
                 return mUserIds;
@@ -234,15 +239,6 @@
                 synchronized (mUserIds) {
                     mUserIds.remove(userId);
                 }
-                synchronized (mUserIdToLabelMap) {
-                    mUserIdToLabelMap.remove(userId);
-                }
-                synchronized (mUserIdToBadgeMap) {
-                    mUserIdToBadgeMap.remove(userId);
-                }
-                synchronized (mCanFrowardToProfileIdMap) {
-                    mCanFrowardToProfileIdMap.remove(userId);
-                }
             } else if (Intent.ACTION_PROFILE_AVAILABLE.equals(action)) {
                 synchronized (mUserIds) {
                     if (!mUserIds.contains(userId)) {
@@ -250,19 +246,40 @@
                     }
                 }
                 synchronized (mUserIdToLabelMap) {
-                    mUserIdToLabelMap.put(userId, getProfileLabel(userId));
+                    if (!mUserIdToLabelMap.containsKey(userId)) {
+                        mUserIdToLabelMap.put(userId, getProfileLabel(userId));
+                    }
                 }
                 synchronized (mUserIdToBadgeMap) {
-                    mUserIdToBadgeMap.put(userId, getProfileBadge(userId));
+                    if (!mUserIdToBadgeMap.containsKey(userId)) {
+                        mUserIdToBadgeMap.put(userId, getProfileBadge(userId));
+                    }
                 }
                 synchronized (mCanFrowardToProfileIdMap) {
-                    mCanFrowardToProfileIdMap.put(userId, true);
+                    if (!mCanFrowardToProfileIdMap.containsKey(userId)) {
+                        if (userId.getIdentifier() == ActivityManager.getCurrentUser()
+                                || isCrossProfileContentSharingStrategyDelegatedFromParent(
+                                UserHandle.of(userId.getIdentifier()))
+                                || CrossProfileUtils.getCrossProfileResolveInfo(mCurrentUser,
+                                mContext.getPackageManager(), mCurrentStateIntent, mContext,
+                                mConfigStore.isPrivateSpaceInDocsUIEnabled()) != null) {
+                            mCanFrowardToProfileIdMap.put(userId, true);
+                        } else {
+                            mCanFrowardToProfileIdMap.put(userId, false);
+                        }
+
+                    }
                 }
             } else {
                 Log.e(TAG, "Unexpected action received: " + action);
             }
         }
 
+        @Override
+        public void setCurrentStateIntent(Intent intent) {
+            mCurrentStateIntent = intent;
+        }
+
         private List<UserId> getUserIdsInternal() {
             final List<UserId> result = new ArrayList<>();
 
@@ -538,8 +555,7 @@
             /*
              * Cross profile resolve info need to be checked in the following 2 cases:
              * 1. current user is either parent or delegates check to parent and the target user
-             * does
-             *    not delegate to parent
+             *    does not delegate to parent
              * 2. current user does not delegate check to the parent and the target user is the
              *    parent profile
              */
diff --git a/src/com/android/documentsui/files/FilesActivity.java b/src/com/android/documentsui/files/FilesActivity.java
index 053e5db..7099381 100644
--- a/src/com/android/documentsui/files/FilesActivity.java
+++ b/src/com/android/documentsui/files/FilesActivity.java
@@ -300,7 +300,7 @@
     }
 
     @Override
-    public void onDestroy() {
+    protected void onDestroy() {
         super.onDestroy();
     }
 
diff --git a/src/com/android/documentsui/picker/PickActivity.java b/src/com/android/documentsui/picker/PickActivity.java
index 2061a8f..8ba48db 100644
--- a/src/com/android/documentsui/picker/PickActivity.java
+++ b/src/com/android/documentsui/picker/PickActivity.java
@@ -300,6 +300,11 @@
     }
 
     @Override
+    protected void onDestroy() {
+        super.onDestroy();
+    }
+
+    @Override
     public String getDrawerTitle() {
         String title;
         try {
diff --git a/tests/common/com/android/documentsui/TestUserManagerState.java b/tests/common/com/android/documentsui/TestUserManagerState.java
index 8dbe62a..5f41ce3 100644
--- a/tests/common/com/android/documentsui/TestUserManagerState.java
+++ b/tests/common/com/android/documentsui/TestUserManagerState.java
@@ -18,6 +18,7 @@
 
 import android.content.Intent;
 import android.graphics.drawable.Drawable;
+import android.util.Log;
 
 import com.android.documentsui.base.UserId;
 
@@ -28,12 +29,14 @@
 
 public class TestUserManagerState implements UserManagerState {
 
+    private static final String TAG = "TestUserManagerState";
     public List<UserId> userIds = new ArrayList<>();
     public Map<UserId, String> userIdToLabelMap = new HashMap<>();
     public Map<UserId, Boolean> canFrowardToProfileIdMap = new HashMap<>();
     public Map<UserId, Drawable> userIdToBadgeMap = new HashMap<>();
     public String profileLabel = "Test";
     public Drawable profileBadge = null;
+    public Intent intent = new Intent();
 
     @Override
     public List<UserId> getUserIds() {
@@ -59,16 +62,26 @@
     public void onProfileActionStatusChange(String action, UserId userId) {
         if (Intent.ACTION_PROFILE_UNAVAILABLE.equals(action)) {
             userIds.remove(userId);
-            userIdToLabelMap.remove(userId);
-            userIdToBadgeMap.remove(userId);
-            canFrowardToProfileIdMap.put(userId, false);
             return;
+        } else if (Intent.ACTION_PROFILE_AVAILABLE.equals(action)) {
+            if (!userIds.contains(userId)) {
+                userIds.add(userId);
+            }
+            if (!userIdToLabelMap.containsKey(userId)) {
+                userIdToLabelMap.put(userId, profileLabel);
+            }
+            if (!userIdToBadgeMap.containsKey(userId)) {
+                userIdToBadgeMap.put(userId, profileBadge);
+            }
+            if (!canFrowardToProfileIdMap.containsKey(userId)) {
+                canFrowardToProfileIdMap.put(userId, true);
+            }
+        } else {
+            Log.e(TAG, "Unexpected action received: " + action);
         }
-        if (!userIds.contains(userId)) {
-            userIds.add(userId);
-        }
-        userIdToLabelMap.put(userId, profileLabel);
-        userIdToBadgeMap.put(userId, profileBadge);
-        canFrowardToProfileIdMap.put(userId, true);
+    }
+
+    @Override
+    public void setCurrentStateIntent(Intent intent) {
     }
 }
diff --git a/tests/unit/com/android/documentsui/UserManagerStateTest.java b/tests/unit/com/android/documentsui/UserManagerStateTest.java
index 25c7607..afe8f45 100644
--- a/tests/unit/com/android/documentsui/UserManagerStateTest.java
+++ b/tests/unit/com/android/documentsui/UserManagerStateTest.java
@@ -16,22 +16,33 @@
 
 package com.android.documentsui;
 
+import static com.android.documentsui.DevicePolicyResources.Drawables.Style.SOLID_COLORED;
+import static com.android.documentsui.DevicePolicyResources.Drawables.WORK_PROFILE_ICON;
+import static com.android.documentsui.DevicePolicyResources.Strings.PERSONAL_TAB;
+import static com.android.documentsui.DevicePolicyResources.Strings.WORK_TAB;
+
 import static com.google.common.truth.Truth.assertWithMessage;
 
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
 
 import android.Manifest;
+import android.app.admin.DevicePolicyManager;
+import android.app.admin.DevicePolicyResourcesManager;
 import android.content.Context;
 import android.content.Intent;
 import android.content.pm.PackageManager;
 import android.content.pm.ResolveInfo;
 import android.content.pm.UserProperties;
+import android.graphics.drawable.Drawable;
 import android.os.UserHandle;
 import android.os.UserManager;
 
 import androidx.test.filters.SdkSuppress;
 import androidx.test.filters.SmallTest;
+import androidx.test.platform.app.InstrumentationRegistry;
 
 import com.android.documentsui.base.UserId;
 import com.android.documentsui.testing.UserManagers;
@@ -52,6 +63,10 @@
 @SdkSuppress(minSdkVersion = 31, codeName = "S")
 public class UserManagerStateTest {
 
+    private static final String PERSONAL = "Personal";
+    private static final String WORK = "Work";
+    private static final String PRIVATE = "Private";
+
     private final UserHandle mSystemUser = UserHandle.SYSTEM;
     private final UserHandle mManagedUser = UserHandle.of(100);
     private final UserHandle mPrivateUser = UserHandle.of(101);
@@ -66,6 +81,7 @@
     private final Intent mMockIntent = mock(Intent.class);
     private final UserManager mMockUserManager = UserManagers.create();
     private final PackageManager mMockPackageManager = mock(PackageManager.class);
+    private final DevicePolicyManager mDevicePolicyManager = mock(DevicePolicyManager.class);
     private UserManagerState mUserManagerState;
 
     @Before
@@ -114,7 +130,7 @@
             when(mMockUserManager.getUserProperties(mNormalUser)).thenReturn(normalUserProperties);
         }
 
-        when(mMockUserManager.getProfileParent(mSystemUser)).thenReturn(mSystemUser);
+        when(mMockUserManager.getProfileParent(mSystemUser)).thenReturn(null);
         when(mMockUserManager.getProfileParent(mManagedUser)).thenReturn(mSystemUser);
         when(mMockUserManager.getProfileParent(mPrivateUser)).thenReturn(mSystemUser);
         when(mMockUserManager.getProfileParent(mOtherUser)).thenReturn(mSystemUser);
@@ -129,6 +145,12 @@
         when(mMockContext.getPackageManager()).thenReturn(mMockPackageManager);
         when(mMockContext.getSystemServiceName(UserManager.class)).thenReturn("mMockUserManager");
         when(mMockContext.getSystemService(UserManager.class)).thenReturn(mMockUserManager);
+        when(mMockContext.getSystemServiceName(DevicePolicyManager.class))
+                .thenReturn(Context.DEVICE_POLICY_SERVICE);
+        when(mMockContext.getSystemService(Context.DEVICE_POLICY_SERVICE))
+                .thenReturn(mDevicePolicyManager);
+        when(mMockContext.getResources()).thenReturn(
+                InstrumentationRegistry.getInstrumentation().getTargetContext().getResources());
     }
 
     @Test
@@ -560,7 +582,7 @@
     }
 
     @Test
-    public void testOnProfileStatusChange_anyIntentActionOnManagedProfile() {
+    public void testOnProfileStatusChange_anyIntentActionForManagedProfile() {
         if (!SdkLevel.isAtLeastV()) return;
         UserId currentUser = UserId.of(mSystemUser);
         initializeUserManagerState(currentUser,
@@ -584,7 +606,7 @@
     }
 
     @Test
-    public void testOnProfileStatusChange_actionProfileUnavailableOnPrivateProfile() {
+    public void testOnProfileStatusChange_actionProfileUnavailableForPrivateProfile() {
         if (!SdkLevel.isAtLeastV()) return;
         UserId currentUser = UserId.of(mSystemUser);
         UserId managedUser = UserId.of(mManagedUser);
@@ -603,9 +625,6 @@
                 mUserManagerState.getCanForwardToProfileIdMap(mMockIntent));
 
         List<UserId> expectedUserIdsAfterIntent = Lists.newArrayList(currentUser, managedUser);
-        Map<UserId, Boolean> expectedCanForwardToProfileIdMapAfterIntent = new HashMap<>();
-        expectedCanForwardToProfileIdMapAfterIntent.put(currentUser, true);
-        expectedCanForwardToProfileIdMapAfterIntent.put(managedUser, true);
 
         String action = Intent.ACTION_PROFILE_UNAVAILABLE;
         mUserManagerState.onProfileActionStatusChange(action, privateUser);
@@ -615,19 +634,14 @@
                 .that(mUserManagerState.getUserIds()).isNotEqualTo(userIdsBeforeIntent);
         assertWithMessage("Unexpected changes to user id list on receiving intent: " + action)
                 .that(mUserManagerState.getUserIds()).isEqualTo(expectedUserIdsAfterIntent);
-        assertWithMessage(
-                "CanForwardToLabelMap should not be same before and after receiving intent: "
-                        + action)
-                .that(mUserManagerState.getCanForwardToProfileIdMap(mMockIntent)).isNotEqualTo(
-                        canForwardToProfileIdMapBeforeIntent);
-        assertWithMessage(
-                "Unexpected changes to canForwardToProfileIdMap on receiving intent: " + action)
+        assertWithMessage("CanForwardToLabelMap should be same before and after receiving intent: "
+                + action)
                 .that(mUserManagerState.getCanForwardToProfileIdMap(mMockIntent)).isEqualTo(
-                        expectedCanForwardToProfileIdMapAfterIntent);
+                        canForwardToProfileIdMapBeforeIntent);
     }
 
     @Test
-    public void testOnProfileStatusChange_actionProfileAvailableOnPrivateProfile() {
+    public void testOnProfileStatusChange_actionProfileAvailable_profileInitialised() {
         if (!SdkLevel.isAtLeastV()) return;
         UserId currentUser = UserId.of(mSystemUser);
         UserId managedUser = UserId.of(mManagedUser);
@@ -637,9 +651,54 @@
                 PackageManager.MATCH_DEFAULT_ONLY, mSystemUser)).thenReturn(
                 mMockResolveInfoList);
         initializeUserManagerState(currentUser,
-                Lists.newArrayList(mSystemUser, mManagedUser));
+                Lists.newArrayList(mSystemUser, mManagedUser, mPrivateUser));
 
-        // UserManagerState#mUserId and UserManagerState#mCanForwardToProfileIdMap will empty
+        // initialising the userIds list and canForwardToProfileIdMap
+        mUserManagerState.getUserIds();
+        mUserManagerState.getCanForwardToProfileIdMap(mMockIntent);
+
+        // Making the private profile unavailable after it has been initialised
+        mUserManagerState.onProfileActionStatusChange(Intent.ACTION_PROFILE_UNAVAILABLE,
+                privateUser);
+
+        List<UserId> userIdsBeforeIntent = new ArrayList<>(mUserManagerState.getUserIds());
+        Map<UserId, Boolean> canForwardToProfileIdMapBeforeIntent = new HashMap<>(
+                mUserManagerState.getCanForwardToProfileIdMap(mMockIntent));
+
+        List<UserId> expectedUserIdsAfterIntent = Lists.newArrayList(currentUser, managedUser,
+                privateUser);
+
+        String action = Intent.ACTION_PROFILE_AVAILABLE;
+        mUserManagerState.onProfileActionStatusChange(action, privateUser);
+
+        assertWithMessage(
+                "UserIds list should not be same before and after receiving intent: " + action)
+                .that(mUserManagerState.getUserIds()).isNotEqualTo(userIdsBeforeIntent);
+        assertWithMessage("Unexpected changes to user id list on receiving intent: " + action)
+                .that(mUserManagerState.getUserIds()).isEqualTo(expectedUserIdsAfterIntent);
+        assertWithMessage("CanForwardToLabelMap should be same before and after receiving intent: "
+                + action)
+                .that(mUserManagerState.getCanForwardToProfileIdMap(mMockIntent)).isEqualTo(
+                        canForwardToProfileIdMapBeforeIntent);
+    }
+
+    @Test
+    public void testOnProfileStatusChange_actionProfileAvailable_profileNotInitialised() {
+        if (!SdkLevel.isAtLeastV()) return;
+        UserId currentUser = UserId.of(mSystemUser);
+        UserId managedUser = UserId.of(mManagedUser);
+        UserId privateUser = UserId.of(mPrivateUser);
+        final List<ResolveInfo> mMockResolveInfoList = Lists.newArrayList(mMockInfo1, mMockInfo2);
+        when(mMockPackageManager.queryIntentActivitiesAsUser(mMockIntent,
+                PackageManager.MATCH_DEFAULT_ONLY, mSystemUser)).thenReturn(
+                mMockResolveInfoList);
+
+        // Private user will not be initialised if it is in quiet mode
+        when(mMockUserManager.isQuietModeEnabled(mPrivateUser)).thenReturn(true);
+        initializeUserManagerState(currentUser,
+                Lists.newArrayList(mSystemUser, mManagedUser, mPrivateUser));
+
+        // UserManagerState#mUserId and UserManagerState#mCanForwardToProfileIdMap will be empty
         // by default if the getters of these member variables have not been called
         List<UserId> userIdsBeforeIntent = new ArrayList<>(mUserManagerState.getUserIds());
         Map<UserId, Boolean> canForwardToProfileIdMapBeforeIntent = new HashMap<>(
@@ -671,6 +730,103 @@
                         expectedCanForwardToProfileIdMapAfterIntent);
     }
 
+    @Test
+    public void testGetUserIdToLabelMap_systemUserAndManagedUser_PreV() {
+        if (SdkLevel.isAtLeastV()) return;
+        UserId currentUser = UserId.of(mSystemUser);
+        initializeUserManagerState(currentUser,
+                Lists.newArrayList(mSystemUser, mManagedUser));
+        if (SdkLevel.isAtLeastT()) {
+            DevicePolicyResourcesManager devicePolicyResourcesManager = mock(
+                    DevicePolicyResourcesManager.class);
+            when(mDevicePolicyManager.getResources()).thenReturn(devicePolicyResourcesManager);
+            when(devicePolicyResourcesManager.getString(eq(PERSONAL_TAB), any())).thenReturn(
+                    PERSONAL);
+            when(devicePolicyResourcesManager.getString(eq(WORK_TAB), any())).thenReturn(WORK);
+        }
+
+        Map<UserId, String> userIdToLabelMap = mUserManagerState.getUserIdToLabelMap();
+
+        assertWithMessage("Incorrect label returned for user id " + mSystemUser)
+                .that(userIdToLabelMap.get(UserId.of(mSystemUser))).isEqualTo(PERSONAL);
+        assertWithMessage("Incorrect label returned for user id " + mManagedUser)
+                .that(userIdToLabelMap.get(UserId.of(mManagedUser))).isEqualTo(WORK);
+    }
+
+    @Test
+    public void testGetUserIdToLabelMap_systemUserManagedUserPrivateUser_PostV() {
+        if (!SdkLevel.isAtLeastV()) return;
+        UserId currentUser = UserId.of(mSystemUser);
+        initializeUserManagerState(currentUser,
+                Lists.newArrayList(mSystemUser, mManagedUser, mPrivateUser));
+        if (SdkLevel.isAtLeastT()) {
+            DevicePolicyResourcesManager devicePolicyResourcesManager = mock(
+                    DevicePolicyResourcesManager.class);
+            when(mDevicePolicyManager.getResources()).thenReturn(devicePolicyResourcesManager);
+            when(devicePolicyResourcesManager.getString(eq(PERSONAL_TAB), any())).thenReturn(
+                    PERSONAL);
+        }
+        UserManager managedUserManager = getUserManagerForManagedUser();
+        UserManager privateUserManager = getUserManagerForPrivateUser();
+        when(managedUserManager.getProfileLabel()).thenReturn(WORK);
+        when(privateUserManager.getProfileLabel()).thenReturn(PRIVATE);
+
+        Map<UserId, String> userIdToLabelMap = mUserManagerState.getUserIdToLabelMap();
+
+        assertWithMessage("Incorrect label returned for user id " + mSystemUser)
+                .that(userIdToLabelMap.get(UserId.of(mSystemUser))).isEqualTo(PERSONAL);
+        assertWithMessage("Incorrect label returned for user id " + mManagedUser)
+                .that(userIdToLabelMap.get(UserId.of(mManagedUser))).isEqualTo(WORK);
+        assertWithMessage("Incorrect label returned for user id " + mPrivateUser)
+                .that(userIdToLabelMap.get(UserId.of(mPrivateUser))).isEqualTo(PRIVATE);
+    }
+
+    @Test
+    public void testGetUserIdToBadgeMap_systemUserManagedUser_PreV() {
+        if (SdkLevel.isAtLeastV()) return;
+        UserId currentUser = UserId.of(mSystemUser);
+        initializeUserManagerState(currentUser,
+                Lists.newArrayList(mSystemUser, mManagedUser));
+        Drawable workBadge = mMockContext.getDrawable(R.drawable.ic_briefcase);
+        if (SdkLevel.isAtLeastT()) {
+            DevicePolicyResourcesManager devicePolicyResourcesManager = mock(
+                    DevicePolicyResourcesManager.class);
+            when(mDevicePolicyManager.getResources()).thenReturn(devicePolicyResourcesManager);
+            when(devicePolicyResourcesManager.getDrawable(eq(WORK_PROFILE_ICON), eq(SOLID_COLORED),
+                    any())).thenReturn(workBadge);
+        }
+
+        Map<UserId, Drawable> userIdToBadgeMap = mUserManagerState.getUserIdToBadgeMap();
+
+        assertWithMessage("There should be no badge present for personal user")
+                .that(userIdToBadgeMap.containsKey(UserId.of(mSystemUser))).isFalse();
+        assertWithMessage("Incorrect badge returned for user id " + mManagedUser)
+                .that(userIdToBadgeMap.get(UserId.of(mManagedUser))).isEqualTo(workBadge);
+    }
+
+    @Test
+    public void testGetUserIdToBadgeMap_systemUserManagedUserPrivateUser_PostV() {
+        if (!SdkLevel.isAtLeastV()) return;
+        UserId currentUser = UserId.of(mSystemUser);
+        initializeUserManagerState(currentUser,
+                Lists.newArrayList(mSystemUser, mManagedUser, mPrivateUser));
+        Drawable workBadge = mock(Drawable.class);
+        Drawable privateBadge = mock(Drawable.class);
+        UserManager managedUserManager = getUserManagerForManagedUser();
+        UserManager privateUserManager = getUserManagerForPrivateUser();
+        when(managedUserManager.getUserBadge()).thenReturn(workBadge);
+        when(privateUserManager.getUserBadge()).thenReturn(privateBadge);
+
+        Map<UserId, Drawable> userIdToBadgeMap = mUserManagerState.getUserIdToBadgeMap();
+
+        assertWithMessage("There should be no badge present for personal user")
+                .that(userIdToBadgeMap.get(UserId.of(mSystemUser))).isNull();
+        assertWithMessage("Incorrect badge returned for user id " + mManagedUser)
+                .that(userIdToBadgeMap.get(UserId.of(mManagedUser))).isEqualTo(workBadge);
+        assertWithMessage("Incorrect badge returned for user id " + mPrivateUser)
+                .that(userIdToBadgeMap.get(UserId.of(mPrivateUser))).isEqualTo(privateBadge);
+    }
+
     private void initializeUserManagerState(UserId current, List<UserHandle> usersOnDevice) {
         when(mMockUserManager.getUserProfiles()).thenReturn(usersOnDevice);
         TestConfigStore testConfigStore = new TestConfigStore();
@@ -678,4 +834,24 @@
         mUserManagerState = new UserManagerState.RuntimeUserManagerState(mMockContext, current,
                 true, testConfigStore);
     }
+
+    private UserManager getUserManagerForManagedUser() {
+        Context managedUserContext = mock(Context.class);
+        when(mMockContext.createContextAsUser(mManagedUser, 0)).thenReturn(managedUserContext);
+        UserManager managedUserManager = mock(UserManager.class);
+        when(managedUserContext.getSystemServiceName(UserManager.class))
+                .thenReturn("managedUserManager");
+        when(managedUserContext.getSystemService(UserManager.class)).thenReturn(managedUserManager);
+        return managedUserManager;
+    }
+
+    private UserManager getUserManagerForPrivateUser() {
+        Context privateUserContext = mock(Context.class);
+        when(mMockContext.createContextAsUser(mPrivateUser, 0)).thenReturn(privateUserContext);
+        UserManager privateUserManager = mock(UserManager.class);
+        when(privateUserContext.getSystemServiceName(UserManager.class))
+                .thenReturn("privateUserManager");
+        when(privateUserContext.getSystemService(UserManager.class)).thenReturn(privateUserManager);
+        return privateUserManager;
+    }
 }