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;
+    }
 }