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

7 files changed, 265 insertions, 57 deletions

diff --git a/src/com/android/documentsui/BaseActivity.java b/src/com/android/documentsui/BaseActivity.java
index 4228ead20..1e9c5fde0 100644
--- a/src/com/android/documentsui/BaseActivity.java
+++ b/src/com/android/documentsui/BaseActivity.java
@@ -140,7 +140,7 @@ public abstract class BaseActivity
 
     @VisibleForTesting
     protected void initConfigStore() {
-        mConfigStore = DocumentsApplication.getConfigStore(this);
+        mConfigStore = DocumentsApplication.getConfigStore();
     }
 
     @VisibleForTesting
@@ -287,6 +287,12 @@ public abstract class BaseActivity
 
         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 92aa3f75e..60fbdaf7c 100644
--- a/src/com/android/documentsui/DocumentsApplication.java
+++ b/src/com/android/documentsui/DocumentsApplication.java
@@ -85,7 +85,7 @@ public class DocumentsApplication extends Application {
     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 class DocumentsApplication extends Application {
     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 @@ public class DocumentsApplication extends Application {
     /**
      * 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 @@ public class DocumentsApplication extends Application {
     }
 
     /**
-     * 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 @@ public class DocumentsApplication extends Application {
             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 @@ public class DocumentsApplication extends Application {
             } 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 28a83e88f..eccc6b00c 100644
--- a/src/com/android/documentsui/UserManagerState.java
+++ b/src/com/android/documentsui/UserManagerState.java
@@ -95,6 +95,11 @@ public interface UserManagerState {
     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 @@ public interface UserManagerState {
         @GuardedBy("mCanFrowardToProfileIdMap")
         private final Map<UserId, Boolean> mCanFrowardToProfileIdMap = new HashMap<>();
 
+        private Intent mCurrentStateIntent;
 
         private final BroadcastReceiver mIntentReceiver = new BroadcastReceiver() {
             @Override
@@ -159,7 +165,7 @@ public interface UserManagerState {
         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 interface UserManagerState {
         public List<UserId> getUserIds() {
             synchronized (mUserIds) {
                 if (mUserIds.isEmpty()) {
-                    Log.d("profileAction", "user ids empty");
                     mUserIds.addAll(getUserIdsInternal());
                 }
                 return mUserIds;
@@ -234,15 +239,6 @@ public interface UserManagerState {
                 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 @@ public interface UserManagerState {
                     }
                 }
                 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 @@ public interface UserManagerState {
             /*
              * 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 053e5dbac..709938174 100644
--- a/src/com/android/documentsui/files/FilesActivity.java
+++ b/src/com/android/documentsui/files/FilesActivity.java
@@ -300,7 +300,7 @@ public class FilesActivity extends BaseActivity implements AbstractActionHandler
     }
 
     @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 2061a8fee..8ba48dbdd 100644
--- a/src/com/android/documentsui/picker/PickActivity.java
+++ b/src/com/android/documentsui/picker/PickActivity.java
@@ -300,6 +300,11 @@ public class PickActivity extends BaseActivity implements ActionHandler.Addons {
     }
 
     @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 8dbe62a2a..5f41ce3e0 100644
--- a/tests/common/com/android/documentsui/TestUserManagerState.java
+++ b/tests/common/com/android/documentsui/TestUserManagerState.java
@@ -18,6 +18,7 @@ package com.android.documentsui;
 
 import android.content.Intent;
 import android.graphics.drawable.Drawable;
+import android.util.Log;
 
 import com.android.documentsui.base.UserId;
 
@@ -28,12 +29,14 @@ import java.util.Map;
 
 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 class TestUserManagerState implements UserManagerState {
     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 25c760727..afe8f4564 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 @@ import java.util.Map;
 @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 @@ public class UserManagerStateTest {
     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 @@ public class UserManagerStateTest {
             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 @@ public class UserManagerStateTest {
         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 @@ public class UserManagerStateTest {
     }
 
     @Test
-    public void testOnProfileStatusChange_anyIntentActionOnManagedProfile() {
+    public void testOnProfileStatusChange_anyIntentActionForManagedProfile() {
         if (!SdkLevel.isAtLeastV()) return;
         UserId currentUser = UserId.of(mSystemUser);
         initializeUserManagerState(currentUser,
@@ -584,7 +606,7 @@ public class UserManagerStateTest {
     }
 
     @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 @@ public class UserManagerStateTest {
                 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,56 @@ public class UserManagerStateTest {
                 .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(
+        assertWithMessage("CanForwardToLabelMap should be same before and after receiving intent: "
+                + action)
+                .that(mUserManagerState.getCanForwardToProfileIdMap(mMockIntent)).isEqualTo(
                         canForwardToProfileIdMapBeforeIntent);
+    }
+
+    @Test
+    public void testOnProfileStatusChange_actionProfileAvailable_profileInitialised() {
+        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);
+        initializeUserManagerState(currentUser,
+                Lists.newArrayList(mSystemUser, mManagedUser, mPrivateUser));
+
+        // 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(
-                "Unexpected changes to canForwardToProfileIdMap on receiving intent: " + action)
+                "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(
-                        expectedCanForwardToProfileIdMapAfterIntent);
+                        canForwardToProfileIdMapBeforeIntent);
     }
 
     @Test
-    public void testOnProfileStatusChange_actionProfileAvailableOnPrivateProfile() {
+    public void testOnProfileStatusChange_actionProfileAvailable_profileNotInitialised() {
         if (!SdkLevel.isAtLeastV()) return;
         UserId currentUser = UserId.of(mSystemUser);
         UserId managedUser = UserId.of(mManagedUser);
@@ -636,10 +692,13 @@ public class UserManagerStateTest {
         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));
+                Lists.newArrayList(mSystemUser, mManagedUser, mPrivateUser));
 
-        // UserManagerState#mUserId and UserManagerState#mCanForwardToProfileIdMap will empty
+        // 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 @@ public class UserManagerStateTest {
                         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 @@ public class UserManagerStateTest {
         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;
+    }
 }