Load App icons from correct user and badge when cross profile is supported.

Bug: 148271447
Test: atest DocumentsUIGoogleTests
Test: manual
Change-Id: I281682a277ab1e064e5867e0aeae76a45eb83984

21 files changed, 219 insertions, 94 deletions

diff --git a/src/com/android/documentsui/IconUtils.java b/src/com/android/documentsui/IconUtils.java
index 46d032124..5a829ff90 100644
--- a/src/com/android/documentsui/IconUtils.java
+++ b/src/com/android/documentsui/IconUtils.java
@@ -22,21 +22,29 @@ import android.content.pm.ProviderInfo;
 import android.graphics.drawable.Drawable;
 import android.util.TypedValue;
 
-import com.android.documentsui.base.MimeTypes;
+import com.android.documentsui.base.UserId;
 
 public class IconUtils {
-    public static Drawable loadPackageIcon(Context context, String authority, int icon) {
+    public static Drawable loadPackageIcon(Context context, UserId userId, String authority,
+            int icon, boolean maybeShowBadge) {
         if (icon != 0) {
+            final PackageManager pm = userId.getPackageManager(context);
+            Drawable packageIcon = null;
             if (authority != null) {
-                final PackageManager pm = context.getPackageManager();
                 final ProviderInfo info = pm.resolveContentProvider(authority, 0);
                 if (info != null) {
-                    return pm.getDrawable(info.packageName, icon, info.applicationInfo);
+                    packageIcon = pm.getDrawable(info.packageName, icon, info.applicationInfo);
                 }
             } else {
-                return context.getDrawable(icon);
+                packageIcon = userId.getDrawable(context, icon);
+            }
+            if (maybeShowBadge) {
+                return userId.getUserBadgedIcon(context, packageIcon);
+            } else {
+                return packageIcon;
             }
         }
+
         return null;
     }
 
diff --git a/src/com/android/documentsui/base/RootInfo.java b/src/com/android/documentsui/base/RootInfo.java
index 8e5a197db..53fedf2ff 100644
--- a/src/com/android/documentsui/base/RootInfo.java
+++ b/src/com/android/documentsui/base/RootInfo.java
@@ -423,24 +423,25 @@ public class RootInfo implements Durable, Parcelable, Comparable<RootInfo> {
         }
     }
 
-    public Drawable loadIcon(Context context) {
+    public Drawable loadIcon(Context context, boolean maybeShowBadge) {
         if (derivedIcon == LOAD_FROM_CONTENT_RESOLVER) {
             return loadMimeTypeIcon(context);
         } else if (derivedIcon != 0) {
+            // derivedIcon is set with the resources of the current user.
             return context.getDrawable(derivedIcon);
         } else {
-            return IconUtils.loadPackageIcon(context, authority, icon);
+            return IconUtils.loadPackageIcon(context, userId, authority, icon, maybeShowBadge);
         }
     }
 
-    public Drawable loadDrawerIcon(Context context) {
+    public Drawable loadDrawerIcon(Context context, boolean maybeShowBadge) {
         if (derivedIcon == LOAD_FROM_CONTENT_RESOLVER) {
             return IconUtils.applyTintColor(context, loadMimeTypeIcon(context),
                     R.color.item_root_icon);
         } else if (derivedIcon != 0) {
             return IconUtils.applyTintColor(context, derivedIcon, R.color.item_root_icon);
         } else {
-            return IconUtils.loadPackageIcon(context, authority, icon);
+            return IconUtils.loadPackageIcon(context, userId, authority, icon, maybeShowBadge);
         }
     }
 
diff --git a/src/com/android/documentsui/base/UserId.java b/src/com/android/documentsui/base/UserId.java
index 1b6b3182f..b65300c32 100644
--- a/src/com/android/documentsui/base/UserId.java
+++ b/src/com/android/documentsui/base/UserId.java
@@ -22,6 +22,7 @@ import android.content.ContentResolver;
 import android.content.Context;
 import android.content.Intent;
 import android.content.pm.PackageManager;
+import android.graphics.drawable.Drawable;
 import android.net.Uri;
 import android.os.Process;
 import android.os.UserHandle;
@@ -90,6 +91,7 @@ public final class UserId {
         } catch (PackageManager.NameNotFoundException e) {
             throw new IllegalStateException("android package not found.");
         }
+
     }
 
     /**
@@ -107,6 +109,21 @@ public final class UserId {
     }
 
     /**
+     * Returns a drawable object associated with a particular resource ID in this user.
+     */
+    public Drawable getDrawable(Context context, int resId) {
+        return asContext(context).getDrawable(resId);
+    }
+
+    /**
+     * If this target user is a managed profile, then this returns a badged copy of the given icon
+     * to be able to distinguish it from the original icon.
+     */
+    public Drawable getUserBadgedIcon(Context context, Drawable drawable) {
+        return getPackageManager(context).getUserBadgedIcon(drawable, mUserHandle);
+    }
+
+    /**
      * Returns true if this user refers to the system user; false otherwise.
      */
     public boolean isSystem() {
diff --git a/src/com/android/documentsui/dirlist/AppsRowItemData.java b/src/com/android/documentsui/dirlist/AppsRowItemData.java
index fca68d656..305bde20b 100644
--- a/src/com/android/documentsui/dirlist/AppsRowItemData.java
+++ b/src/com/android/documentsui/dirlist/AppsRowItemData.java
@@ -36,16 +36,19 @@ import com.android.documentsui.sidebar.RootItem;
  */
 public abstract class AppsRowItemData {
 
-    private final UserId mUserId;
+    protected final UserId mUserId;
     private final String mTitle;
     private final @Nullable String mSummary;
     protected final ActionHandler mActionHandler;
+    protected final boolean mMaybeShowBadge;
 
-    public AppsRowItemData(Item item, ActionHandler actionHandler, boolean shouldShowSummary) {
+    public AppsRowItemData(Item item, ActionHandler actionHandler, boolean shouldShowSummary,
+            boolean maybeShowBadge) {
         mUserId = item.userId;
         mTitle = item.title;
         mSummary = shouldShowSummary ? item.getSummary() : null;
         mActionHandler = actionHandler;
+        mMaybeShowBadge = maybeShowBadge;
     }
 
     public final String getTitle() {
@@ -70,14 +73,15 @@ public abstract class AppsRowItemData {
 
         private final ResolveInfo mResolveInfo;
 
-        public AppData(AppItem item, ActionHandler actionHandler, boolean shouldShowSummary) {
-            super(item, actionHandler, shouldShowSummary);
+        public AppData(AppItem item, ActionHandler actionHandler, boolean shouldShowSummary,
+                boolean maybeShowBadge) {
+            super(item, actionHandler, shouldShowSummary, maybeShowBadge);
             mResolveInfo = item.info;
         }
 
         @Override
         protected Drawable getIconDrawable(Context context) {
-            return mResolveInfo.loadIcon(context.getPackageManager());
+            return mResolveInfo.loadIcon(mUserId.getPackageManager(context));
         }
 
         @Override
@@ -90,14 +94,15 @@ public abstract class AppsRowItemData {
 
         private final RootInfo mRootInfo;
 
-        public RootData(RootItem item, ActionHandler actionHandler, boolean shouldShowSummary) {
-            super(item, actionHandler, shouldShowSummary);
+        public RootData(RootItem item, ActionHandler actionHandler, boolean shouldShowSummary,
+                boolean maybeShowBadge) {
+            super(item, actionHandler, shouldShowSummary, maybeShowBadge);
             mRootInfo = item.root;
         }
 
         @Override
         protected Drawable getIconDrawable(Context context) {
-            return mRootInfo.loadIcon(context);
+            return mRootInfo.loadIcon(context, mMaybeShowBadge);
         }
 
         @Override
diff --git a/src/com/android/documentsui/dirlist/AppsRowManager.java b/src/com/android/documentsui/dirlist/AppsRowManager.java
index f2ef3ae33..7b7dbc8c4 100644
--- a/src/com/android/documentsui/dirlist/AppsRowManager.java
+++ b/src/com/android/documentsui/dirlist/AppsRowManager.java
@@ -46,10 +46,12 @@ public class AppsRowManager {
 
     private final ActionHandler mActionHandler;
     private final List<AppsRowItemData> mDataList;
+    private final boolean mMaybeShowBadge;
 
-    public AppsRowManager(ActionHandler handler) {
+    public AppsRowManager(ActionHandler handler, boolean maybeShowBadge) {
         mDataList = new ArrayList<>();
         mActionHandler = handler;
+        mMaybeShowBadge = maybeShowBadge;
     }
 
     public List<AppsRowItemData> updateList(List<Item> itemList) {
@@ -68,9 +70,11 @@ public class AppsRowManager {
         for (Item item : itemList) {
             boolean shouldShowSummary = packageNameCount.get(item.getPackageName()) > 1;
             if (item instanceof RootItem) {
-                mDataList.add(new RootData((RootItem) item, mActionHandler, shouldShowSummary));
+                mDataList.add(new RootData((RootItem) item, mActionHandler, shouldShowSummary,
+                        mMaybeShowBadge));
             } else {
-                mDataList.add(new AppData((AppItem) item, mActionHandler, shouldShowSummary));
+                mDataList.add(new AppData((AppItem) item, mActionHandler, shouldShowSummary,
+                        mMaybeShowBadge));
             }
         }
         return mDataList;
diff --git a/src/com/android/documentsui/dirlist/DirectoryFragment.java b/src/com/android/documentsui/dirlist/DirectoryFragment.java
index 8f994c6bd..9d574d4a8 100644
--- a/src/com/android/documentsui/dirlist/DirectoryFragment.java
+++ b/src/com/android/documentsui/dirlist/DirectoryFragment.java
@@ -304,7 +304,7 @@ public class DirectoryFragment extends Fragment implements SwipeRefreshLayout.On
             mLocalState.mSelectionId = Integer.toHexString(System.identityHashCode(mRecView));
         }
 
-        mIconHelper = new IconHelper(mActivity, MODE_GRID);
+        mIconHelper = new IconHelper(mActivity, MODE_GRID, mState.supportsCrossProfile());
 
         mAdapter = new DirectoryAddonsAdapter(
                 mAdapterEnv,
diff --git a/src/com/android/documentsui/dirlist/IconHelper.java b/src/com/android/documentsui/dirlist/IconHelper.java
index 5e58c185a..77e5e9c47 100644
--- a/src/com/android/documentsui/dirlist/IconHelper.java
+++ b/src/com/android/documentsui/dirlist/IconHelper.java
@@ -62,15 +62,17 @@ public class IconHelper {
     private int mMode;
     private Point mCurrentSize;
     private boolean mThumbnailsEnabled = true;
+    private final boolean mMaybeShowBadge;
 
     /**
      * @param context
      * @param mode MODE_GRID or MODE_LIST
      */
-    public IconHelper(Context context, int mode) {
+    public IconHelper(Context context, int mode, boolean maybeShowBadge) {
         mContext = context;
         setViewMode(mode);
         mThumbnailCache = DocumentsApplication.getThumbnailCache(context);
+        mMaybeShowBadge = maybeShowBadge;
     }
 
     /**
@@ -170,7 +172,7 @@ public class IconHelper {
                 loadThumbnail(uri, userId, docAuthority, docLastModified, iconThumb, iconMime);
         }
 
-        final Drawable mimeIcon = getDocumentIcon(mContext, docAuthority,
+        final Drawable mimeIcon = getDocumentIcon(mContext, userId, docAuthority,
                 DocumentsContract.getDocumentId(uri), mimeType, docIcon);
         if (subIconMime != null) {
             setMimeIcon(subIconMime, mimeIcon);
@@ -230,10 +232,10 @@ public class IconHelper {
         view.setAlpha(0f);
     }
 
-    private Drawable getDocumentIcon(
-        Context context, String authority, String id, String mimeType, int icon) {
+    private Drawable getDocumentIcon(Context context, UserId userId, String authority, String id,
+            String mimeType, int icon) {
         if (icon != 0) {
-            return IconUtils.loadPackageIcon(context, authority, icon);
+            return IconUtils.loadPackageIcon(context, userId, authority, icon, mMaybeShowBadge);
         } else {
             return IconUtils.loadMimeIcon(context, mimeType, authority, id, mMode);
         }
@@ -244,6 +246,6 @@ public class IconHelper {
      */
     public Drawable getDocumentIcon(Context context, DocumentInfo doc) {
         return getDocumentIcon(
-                context, doc.authority, doc.documentId, doc.mimeType, doc.icon);
+                context, doc.userId, doc.authority, doc.documentId, doc.mimeType, doc.icon);
     }
 }
diff --git a/src/com/android/documentsui/files/FilesActivity.java b/src/com/android/documentsui/files/FilesActivity.java
index 1d8909fbd..b281a7b8e 100644
--- a/src/com/android/documentsui/files/FilesActivity.java
+++ b/src/com/android/documentsui/files/FilesActivity.java
@@ -158,7 +158,7 @@ public class FilesActivity extends BaseActivity implements AbstractActionHandler
                 mInjector.selectionMgr,
                 mProfileTabsAddonsStub);
 
-        mAppsRowManager = new AppsRowManager(mInjector.actions);
+        mAppsRowManager = new AppsRowManager(mInjector.actions, mState.supportsCrossProfile());
         mInjector.appsRowManager = mAppsRowManager;
 
         mActivityInputHandler =
diff --git a/src/com/android/documentsui/picker/PickActivity.java b/src/com/android/documentsui/picker/PickActivity.java
index 2eb284565..a9965305e 100644
--- a/src/com/android/documentsui/picker/PickActivity.java
+++ b/src/com/android/documentsui/picker/PickActivity.java
@@ -143,7 +143,7 @@ public class PickActivity extends BaseActivity implements ActionHandler.Addons {
 
         Intent intent = getIntent();
 
-        mAppsRowManager = new AppsRowManager(mInjector.actions);
+        mAppsRowManager = new AppsRowManager(mInjector.actions, mState.supportsCrossProfile());
         mInjector.appsRowManager = mAppsRowManager;
 
         mSharedInputHandler =
diff --git a/src/com/android/documentsui/sidebar/AppItem.java b/src/com/android/documentsui/sidebar/AppItem.java
index 06fffaab1..8be1ae5bc 100644
--- a/src/com/android/documentsui/sidebar/AppItem.java
+++ b/src/com/android/documentsui/sidebar/AppItem.java
@@ -55,7 +55,8 @@ public class AppItem extends Item {
     }
 
     protected void bindIcon(ImageView icon) {
-        final PackageManager pm = icon.getContext().getPackageManager();
+        final PackageManager pm = userId.getPackageManager(icon.getContext());
+        // This always gives badged icon if package manager is from a managed profile.
         icon.setImageDrawable(info.loadIcon(pm));
     }
 
diff --git a/src/com/android/documentsui/sidebar/RootAndAppItem.java b/src/com/android/documentsui/sidebar/RootAndAppItem.java
index 6630d8bbd..e6b51f911 100644
--- a/src/com/android/documentsui/sidebar/RootAndAppItem.java
+++ b/src/com/android/documentsui/sidebar/RootAndAppItem.java
@@ -34,8 +34,9 @@ class RootAndAppItem extends RootItem {
 
     public final ResolveInfo resolveInfo;
 
-    public RootAndAppItem(RootInfo root, ResolveInfo info, ActionHandler actionHandler) {
-        super(root, actionHandler, info.activityInfo.packageName);
+    public RootAndAppItem(RootInfo root, ResolveInfo info, ActionHandler actionHandler,
+            boolean maybeShowBadge) {
+        super(root, actionHandler, info.activityInfo.packageName, maybeShowBadge);
         this.resolveInfo = info;
     }
 
diff --git a/src/com/android/documentsui/sidebar/RootItem.java b/src/com/android/documentsui/sidebar/RootItem.java
index 777334b49..9af396e4c 100644
--- a/src/com/android/documentsui/sidebar/RootItem.java
+++ b/src/com/android/documentsui/sidebar/RootItem.java
@@ -51,17 +51,20 @@ public class RootItem extends Item {
     public @Nullable DocumentInfo docInfo;
 
     protected final ActionHandler mActionHandler;
+    protected final boolean mMaybeShowBadge;
     private final String mPackageName;
 
-    public RootItem(RootInfo root, ActionHandler actionHandler) {
-        this(root, actionHandler, "" /* packageName */);
+    public RootItem(RootInfo root, ActionHandler actionHandler, boolean maybeShowBadge) {
+        this(root, actionHandler, "" /* packageName */, maybeShowBadge);
     }
 
-    public RootItem(RootInfo root, ActionHandler actionHandler, String packageName) {
+    public RootItem(RootInfo root, ActionHandler actionHandler, String packageName,
+            boolean maybeShowBadge) {
         super(R.layout.item_root, root.title, getStringId(root), root.userId);
         this.root = root;
         mActionHandler = actionHandler;
         mPackageName = packageName;
+        mMaybeShowBadge = maybeShowBadge;
     }
 
     private static String getStringId(RootInfo root) {
@@ -114,7 +117,7 @@ public class RootItem extends Item {
     }
 
     protected final void bindIconAndTitle(View view) {
-        bindIcon(view, root.loadDrawerIcon(view.getContext()));
+        bindIcon(view, root.loadDrawerIcon(view.getContext(), mMaybeShowBadge));
         bindTitle(view);
     }
 
@@ -213,7 +216,7 @@ public class RootItem extends Item {
     public static RootItem createDummyItem(RootItem item, UserId targetUser) {
         RootInfo dummyRootInfo = RootInfo.copyRootInfo(item.root);
         dummyRootInfo.userId = targetUser;
-        RootItem dummy = new RootItem(dummyRootInfo, item.mActionHandler);
+        RootItem dummy = new RootItem(dummyRootInfo, item.mActionHandler, item.mMaybeShowBadge);
         return dummy;
     }
 }
diff --git a/src/com/android/documentsui/sidebar/RootsFragment.java b/src/com/android/documentsui/sidebar/RootsFragment.java
index 04b2d057b..d34bceb83 100644
--- a/src/com/android/documentsui/sidebar/RootsFragment.java
+++ b/src/com/android/documentsui/sidebar/RootsFragment.java
@@ -247,10 +247,13 @@ public class RootsFragment extends Fragment {
                 final boolean excludeSelf =
                         intent.getBooleanExtra(DocumentsContract.EXTRA_EXCLUDE_SELF, false);
                 final String excludePackage = excludeSelf ? activity.getCallingPackage() : null;
+                final boolean maybeShowBadge =
+                        getBaseActivity().getDisplayState().supportsCrossProfile();
                 List<Item> sortedItems = sortLoadResult(roots, excludePackage, handlerAppIntent,
                         DocumentsApplication.getProvidersCache(getContext()),
                         getBaseActivity().getSelectedUser(),
-                        DocumentsApplication.getUserIdManager(getContext()).getUserIds());
+                        DocumentsApplication.getUserIdManager(getContext()).getUserIds(),
+                        maybeShowBadge);
 
                 // Get the first visible position and offset
                 final int firstPosition = mList.getFirstVisiblePosition();
@@ -294,7 +297,8 @@ public class RootsFragment extends Fragment {
             @Nullable Intent handlerAppIntent,
             ProvidersAccess providersAccess,
             UserId selectedUser,
-            List<UserId> userIds) {
+            List<UserId> userIds,
+            boolean maybeShowBadge) {
         final List<Item> result = new ArrayList<>();
 
         final RootItemListBuilder librariesBuilder = new RootItemListBuilder(selectedUser, userIds);
@@ -306,14 +310,15 @@ public class RootsFragment extends Fragment {
             final RootItem item;
 
             if (root.isLibrary() || root.isDownloads()) {
-                item = new RootItem(root, mActionHandler);
+                item = new RootItem(root, mActionHandler, maybeShowBadge);
                 librariesBuilder.add(item);
             } else if (root.isStorage()) {
-                item = new RootItem(root, mActionHandler);
+                item = new RootItem(root, mActionHandler, maybeShowBadge);
                 storageProvidersBuilder.add(item);
             } else {
                 item = new RootItem(root, mActionHandler,
-                        providersAccess.getPackageName(root.userId, root.authority));
+                        providersAccess.getPackageName(root.userId, root.authority),
+                        maybeShowBadge);
                 otherProviders.add(item);
             }
         }
@@ -337,7 +342,8 @@ public class RootsFragment extends Fragment {
 
         // Include apps that can handle this intent too.
         if (handlerAppIntent != null) {
-            includeHandlerApps(handlerAppIntent, excludePackage, result, otherProviders, userIds);
+            includeHandlerApps(handlerAppIntent, excludePackage, result, otherProviders, userIds,
+                    maybeShowBadge);
         } else {
             // Only add providers
             Collections.sort(otherProviders, comp);
@@ -361,7 +367,7 @@ public class RootsFragment extends Fragment {
      */
     private void includeHandlerApps(
             Intent handlerAppIntent, @Nullable String excludePackage, List<Item> result,
-            List<RootItem> otherProviders, List<UserId> userIds) {
+            List<RootItem> otherProviders, List<UserId> userIds, boolean maybeShowBadge) {
         if (VERBOSE) Log.v(TAG, "Adding handler apps for intent: " + handlerAppIntent);
 
         Context context = getContext();
@@ -417,7 +423,8 @@ public class RootsFragment extends Fragment {
 
             final Item item;
             if (resolveInfo != null) {
-                item = new RootAndAppItem(rootItem.root, resolveInfo, mActionHandler);
+                item = new RootAndAppItem(rootItem.root, resolveInfo, mActionHandler,
+                        maybeShowBadge);
                 appItems.remove(userPackage);
             } else {
                 item = rootItem;
diff --git a/tests/common/com/android/documentsui/testing/TestIconHelper.java b/tests/common/com/android/documentsui/testing/TestIconHelper.java
index ff7956de6..d9bf9ab5a 100644
--- a/tests/common/com/android/documentsui/testing/TestIconHelper.java
+++ b/tests/common/com/android/documentsui/testing/TestIconHelper.java
@@ -29,7 +29,7 @@ public class TestIconHelper extends IconHelper {
     public Drawable nextDocumentIcon;
 
     private TestIconHelper() {
-        super(null, 0);
+        super(null, 0, false);
     }
 
     @Override
diff --git a/tests/common/com/android/documentsui/testing/TestProvidersAccess.java b/tests/common/com/android/documentsui/testing/TestProvidersAccess.java
index 289338d58..4fae7a715 100644
--- a/tests/common/com/android/documentsui/testing/TestProvidersAccess.java
+++ b/tests/common/com/android/documentsui/testing/TestProvidersAccess.java
@@ -15,6 +15,7 @@
  */
 package com.android.documentsui.testing;
 
+import android.os.Process;
 import android.os.UserHandle;
 import android.provider.DocumentsContract.Root;
 
@@ -35,7 +36,8 @@ import javax.annotation.Nullable;
 
 public class TestProvidersAccess implements ProvidersAccess {
 
-    public static final UserId USER_ID = UserId.DEFAULT_USER;
+    public static final UserHandle USER_HANDLE = Process.myUserHandle();
+    public static final UserId USER_ID = UserId.of(USER_HANDLE);
 
     public static final RootInfo DOWNLOADS;
     public static final RootInfo HOME;
diff --git a/tests/unit/com/android/documentsui/IconUtilsTest.java b/tests/unit/com/android/documentsui/IconUtilsTest.java
index f5547bd64..335764235 100644
--- a/tests/unit/com/android/documentsui/IconUtilsTest.java
+++ b/tests/unit/com/android/documentsui/IconUtilsTest.java
@@ -2,11 +2,25 @@ package com.android.documentsui;
 
 import static com.google.common.truth.Truth.assertThat;
 
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyInt;
+import static org.mockito.ArgumentMatchers.anyString;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+import android.content.Context;
+import android.content.pm.ApplicationInfo;
+import android.content.pm.PackageManager;
+import android.content.pm.ProviderInfo;
+import android.graphics.drawable.Drawable;
+
 import androidx.test.InstrumentationRegistry;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
-import android.content.Context;
+import com.android.documentsui.testing.TestProvidersAccess;
 
 import org.junit.Before;
 import org.junit.Test;
@@ -55,4 +69,50 @@ public class IconUtilsTest {
     public void testLoadMimeIcon_isMimeTypeNull_shouldReturnNull() {
         assertThat(IconUtils.loadMimeIcon(mTargetContext, null)).isNull();
     }
+
+    @Test
+    public void testLoadPackageIcon() throws PackageManager.NameNotFoundException {
+        final String authority = "a";
+        final int icon = 1234;
+        final ProviderInfo providerInfo = new ProviderInfo();
+
+        Context context = mock(Context.class);
+        PackageManager packageManager = mock(PackageManager.class);
+        Drawable drawable = mock(Drawable.class);
+
+        when(context.getPackageManager()).thenReturn(packageManager);
+        when(packageManager.resolveContentProvider(eq(authority), anyInt())).thenReturn(
+                providerInfo);
+        when(packageManager.getDrawable(anyString(), eq(icon), any(ApplicationInfo.class)))
+                .thenReturn(drawable);
+
+        IconUtils.loadPackageIcon(context, TestProvidersAccess.USER_ID, authority, icon,
+                /* maybeShowBadge= */false);
+
+        verify(packageManager).getDrawable(any(), eq(icon), any());
+        verify(packageManager, never()).getUserBadgedIcon(drawable,
+                TestProvidersAccess.USER_HANDLE);
+    }
+
+    @Test
+    public void testLoadPackageIcon_maybeShowBadge() throws PackageManager.NameNotFoundException {
+        final String authority = "a";
+        final int icon = 1234;
+        final ProviderInfo providerInfo = new ProviderInfo();
+
+        Context context = mock(Context.class);
+        PackageManager packageManager = mock(PackageManager.class);
+        Drawable drawable = mock(Drawable.class);
+
+        when(context.getPackageManager()).thenReturn(packageManager);
+        when(packageManager.resolveContentProvider(eq(authority), anyInt())).thenReturn(
+                providerInfo);
+        when(packageManager.getDrawable(any(), eq(icon), any())).thenReturn(drawable);
+
+        IconUtils.loadPackageIcon(context, TestProvidersAccess.USER_ID, authority, icon,
+                /* maybeShowBadge= */true);
+
+        verify(packageManager).getDrawable(any(), eq(icon), any());
+        verify(packageManager).getUserBadgedIcon(drawable, TestProvidersAccess.USER_HANDLE);
+    }
 }
diff --git a/tests/unit/com/android/documentsui/dirlist/AppsRowManagerTest.java b/tests/unit/com/android/documentsui/dirlist/AppsRowManagerTest.java
index 7a2dde486..d5a7f5cdd 100644
--- a/tests/unit/com/android/documentsui/dirlist/AppsRowManagerTest.java
+++ b/tests/unit/com/android/documentsui/dirlist/AppsRowManagerTest.java
@@ -53,6 +53,7 @@ public class AppsRowManagerTest {
     private AppsRowManager mAppsRowManager;
 
     private ActionHandler mActionHandler;
+    private boolean mMaybeShowBadge;
     private BaseActivity mActivity;
     private State mState;
 
@@ -63,7 +64,7 @@ public class AppsRowManagerTest {
     public void setUp() {
         mActionHandler = new TestActionHandler();
 
-        mAppsRowManager = new AppsRowManager(mActionHandler);
+        mAppsRowManager = new AppsRowManager(mActionHandler, mMaybeShowBadge);
 
         Context context = InstrumentationRegistry.getInstrumentation().getTargetContext();
         LayoutInflater layoutInflater = LayoutInflater.from(context);
@@ -82,10 +83,12 @@ public class AppsRowManagerTest {
     @Test
     public void testUpdateList_byRootItem() {
         final List<Item> rootList = new ArrayList<>();
-        rootList.add(new RootItem(TestProvidersAccess.INSPECTOR, mActionHandler));
-        rootList.add(new RootItem(TestProvidersAccess.PICKLES, mActionHandler));
-        rootList.add(new RootItem(TestProvidersAccess.PICKLES, mActionHandler, "packageName"));
-        rootList.add(new RootItem(TestProvidersAccess.PICKLES, mActionHandler, "packageName"));
+        rootList.add(new RootItem(TestProvidersAccess.INSPECTOR, mActionHandler, mMaybeShowBadge));
+        rootList.add(new RootItem(TestProvidersAccess.PICKLES, mActionHandler, mMaybeShowBadge));
+        rootList.add(new RootItem(TestProvidersAccess.PICKLES, mActionHandler, "packageName",
+                mMaybeShowBadge));
+        rootList.add(new RootItem(TestProvidersAccess.PICKLES, mActionHandler, "packageName",
+                mMaybeShowBadge));
 
         final List<AppsRowItemData> chipDataList = mAppsRowManager.updateList(rootList);
 
@@ -105,7 +108,8 @@ public class AppsRowManagerTest {
         info.activityInfo.packageName = testPackageName;
 
         List<Item> hybridList = new ArrayList<>();
-        hybridList.add(new RootItem(TestProvidersAccess.INSPECTOR, mActionHandler));
+        hybridList.add(
+                new RootItem(TestProvidersAccess.INSPECTOR, mActionHandler, mMaybeShowBadge));
         hybridList.add(new AppItem(info, TestProvidersAccess.PICKLES.title, UserId.DEFAULT_USER,
                 mActionHandler));
 
@@ -124,7 +128,7 @@ public class AppsRowManagerTest {
         mState.action = State.ACTION_BROWSE;
         mState.stack.changeRoot(TestProvidersAccess.RECENTS);
         final List<Item> rootList = new ArrayList<>();
-        rootList.add(new RootItem(TestProvidersAccess.INSPECTOR, mActionHandler));
+        rootList.add(new RootItem(TestProvidersAccess.INSPECTOR, mActionHandler, mMaybeShowBadge));
         mAppsRowManager.updateList(rootList);
 
         mAppsRowManager.updateView(mActivity);
@@ -138,12 +142,14 @@ public class AppsRowManagerTest {
         mState.action = State.ACTION_GET_CONTENT;
         mState.stack.changeRoot(TestProvidersAccess.RECENTS);
         final List<Item> rootList = new ArrayList<>();
-        rootList.add(new RootItem(TestProvidersAccess.INSPECTOR, mActionHandler));
-        rootList.add(new RootItem(TestProvidersAccess.AUDIO, mActionHandler));
-        rootList.add(new RootItem(TestProvidersAccess.HAMMY, mActionHandler));
-        rootList.add(new RootItem(TestProvidersAccess.IMAGE, mActionHandler));
-        rootList.add(new RootItem(TestProvidersAccess.OtherUser.DOWNLOADS, mActionHandler));
-        rootList.add(new RootItem(TestProvidersAccess.OtherUser.PICKLES, mActionHandler));
+        rootList.add(new RootItem(TestProvidersAccess.INSPECTOR, mActionHandler, mMaybeShowBadge));
+        rootList.add(new RootItem(TestProvidersAccess.AUDIO, mActionHandler, mMaybeShowBadge));
+        rootList.add(new RootItem(TestProvidersAccess.HAMMY, mActionHandler, mMaybeShowBadge));
+        rootList.add(new RootItem(TestProvidersAccess.IMAGE, mActionHandler, mMaybeShowBadge));
+        rootList.add(new RootItem(TestProvidersAccess.OtherUser.DOWNLOADS, mActionHandler,
+                mMaybeShowBadge));
+        rootList.add(new RootItem(TestProvidersAccess.OtherUser.PICKLES, mActionHandler,
+                mMaybeShowBadge));
         mAppsRowManager.updateList(rootList);
 
         mAppsRowManager.updateView(mActivity);
@@ -157,12 +163,14 @@ public class AppsRowManagerTest {
         when(mActivity.getSelectedUser()).thenReturn(TestProvidersAccess.OtherUser.USER_ID);
         mState.stack.changeRoot(TestProvidersAccess.RECENTS);
         final List<Item> rootList = new ArrayList<>();
-        rootList.add(new RootItem(TestProvidersAccess.INSPECTOR, mActionHandler));
-        rootList.add(new RootItem(TestProvidersAccess.AUDIO, mActionHandler));
-        rootList.add(new RootItem(TestProvidersAccess.HAMMY, mActionHandler));
-        rootList.add(new RootItem(TestProvidersAccess.IMAGE, mActionHandler));
-        rootList.add(new RootItem(TestProvidersAccess.OtherUser.DOWNLOADS, mActionHandler));
-        rootList.add(new RootItem(TestProvidersAccess.OtherUser.PICKLES, mActionHandler));
+        rootList.add(new RootItem(TestProvidersAccess.INSPECTOR, mActionHandler, mMaybeShowBadge));
+        rootList.add(new RootItem(TestProvidersAccess.AUDIO, mActionHandler, mMaybeShowBadge));
+        rootList.add(new RootItem(TestProvidersAccess.HAMMY, mActionHandler, mMaybeShowBadge));
+        rootList.add(new RootItem(TestProvidersAccess.IMAGE, mActionHandler, mMaybeShowBadge));
+        rootList.add(new RootItem(TestProvidersAccess.OtherUser.DOWNLOADS, mActionHandler,
+                mMaybeShowBadge));
+        rootList.add(new RootItem(TestProvidersAccess.OtherUser.PICKLES, mActionHandler,
+                mMaybeShowBadge));
         mAppsRowManager.updateList(rootList);
 
         mAppsRowManager.updateView(mActivity);
@@ -174,7 +182,7 @@ public class AppsRowManagerTest {
     public void testUpdateView_notInRecent_hideRow() {
         mState.action = State.ACTION_BROWSE;
         final List<Item> rootList = new ArrayList<>();
-        rootList.add(new RootItem(TestProvidersAccess.INSPECTOR, mActionHandler));
+        rootList.add(new RootItem(TestProvidersAccess.INSPECTOR, mActionHandler, mMaybeShowBadge));
         mAppsRowManager.updateList(rootList);
 
         mState.stack.changeRoot(TestProvidersAccess.DOWNLOADS);
@@ -190,7 +198,7 @@ public class AppsRowManagerTest {
 
         mState.stack.changeRoot(TestProvidersAccess.RECENTS);
         final List<Item> rootList = new ArrayList<>();
-        rootList.add(new RootItem(TestProvidersAccess.INSPECTOR, mActionHandler));
+        rootList.add(new RootItem(TestProvidersAccess.INSPECTOR, mActionHandler, mMaybeShowBadge));
         mAppsRowManager.updateList(rootList);
 
         mAppsRowManager.updateView(mActivity);
@@ -218,10 +226,10 @@ public class AppsRowManagerTest {
         when(mActivity.getSelectedUser()).thenReturn(TestProvidersAccess.OtherUser.USER_ID);
 
         final List<Item> rootList = new ArrayList<>();
-        rootList.add(new RootItem(TestProvidersAccess.INSPECTOR, mActionHandler));
-        rootList.add(new RootItem(TestProvidersAccess.AUDIO, mActionHandler));
-        rootList.add(new RootItem(TestProvidersAccess.HAMMY, mActionHandler));
-        rootList.add(new RootItem(TestProvidersAccess.IMAGE, mActionHandler));
+        rootList.add(new RootItem(TestProvidersAccess.INSPECTOR, mActionHandler, mMaybeShowBadge));
+        rootList.add(new RootItem(TestProvidersAccess.AUDIO, mActionHandler, mMaybeShowBadge));
+        rootList.add(new RootItem(TestProvidersAccess.HAMMY, mActionHandler, mMaybeShowBadge));
+        rootList.add(new RootItem(TestProvidersAccess.IMAGE, mActionHandler, mMaybeShowBadge));
         mAppsRowManager.updateList(rootList);
 
         mAppsRowManager.updateView(mActivity);
diff --git a/tests/unit/com/android/documentsui/dirlist/DirectoryAddonsAdapterTest.java b/tests/unit/com/android/documentsui/dirlist/DirectoryAddonsAdapterTest.java
index 030c8b147..93a641dfb 100644
--- a/tests/unit/com/android/documentsui/dirlist/DirectoryAddonsAdapterTest.java
+++ b/tests/unit/com/android/documentsui/dirlist/DirectoryAddonsAdapterTest.java
@@ -55,9 +55,11 @@ public class DirectoryAddonsAdapterTest extends AndroidTestCase {
         DocumentsAdapter.Environment env = new TestEnvironment(testContext);
 
         mAdapter = new DirectoryAddonsAdapter(
-            env,
-            new ModelBackedDocumentsAdapter(
-                    env, new IconHelper(testContext, State.MODE_GRID), new TestFileTypeLookup()));
+                env,
+                new ModelBackedDocumentsAdapter(
+                        env,
+                        new IconHelper(testContext, State.MODE_GRID, /* maybeShowBadge= */ false),
+                        new TestFileTypeLookup()));
 
         mEnv.model.addUpdateListener(mAdapter.getModelUpdateListener());
     }
diff --git a/tests/unit/com/android/documentsui/dirlist/ModelBackedDocumentsAdapterTest.java b/tests/unit/com/android/documentsui/dirlist/ModelBackedDocumentsAdapterTest.java
index 7e339f40b..5b5d7f440 100644
--- a/tests/unit/com/android/documentsui/dirlist/ModelBackedDocumentsAdapterTest.java
+++ b/tests/unit/com/android/documentsui/dirlist/ModelBackedDocumentsAdapterTest.java
@@ -48,7 +48,9 @@ public class ModelBackedDocumentsAdapterTest extends AndroidTestCase {
         DocumentsAdapter.Environment env = new TestEnvironment(testContext);
 
         mAdapter = new ModelBackedDocumentsAdapter(
-                env, new IconHelper(testContext, State.MODE_GRID), new TestFileTypeLookup());
+                env,
+                new IconHelper(testContext, State.MODE_GRID, /* maybeShowBadge= */ false),
+                new TestFileTypeLookup());
         mAdapter.getModelUpdateListener().accept(Model.Update.UPDATE);
     }
 
diff --git a/tests/unit/com/android/documentsui/sidebar/RootItemListBuilderTest.java b/tests/unit/com/android/documentsui/sidebar/RootItemListBuilderTest.java
index 7f227c881..2574ecc5a 100644
--- a/tests/unit/com/android/documentsui/sidebar/RootItemListBuilderTest.java
+++ b/tests/unit/com/android/documentsui/sidebar/RootItemListBuilderTest.java
@@ -40,26 +40,26 @@ import java.util.List;
 public class RootItemListBuilderTest {
 
     private static final RootItem DOWNLOADS_DEFAULT_USER =
-            new RootItem(TestProvidersAccess.DOWNLOADS, null);
+            new RootItem(TestProvidersAccess.DOWNLOADS, null, false);
     private static final RootItem HAMMY_DEFAULT_USER =
-            new RootItem(TestProvidersAccess.HAMMY, null);
+            new RootItem(TestProvidersAccess.HAMMY, null, false);
     private static final RootItem HOME_DEFAULT_USER =
-            new RootItem(TestProvidersAccess.HOME, null);
+            new RootItem(TestProvidersAccess.HOME, null, false);
     private static final RootItem IMAGE_DEFAULT_USER =
-            new RootItem(TestProvidersAccess.IMAGE, null);
+            new RootItem(TestProvidersAccess.IMAGE, null, false);
     private static final RootItem PICKLES_DEFAULT_USER =
-            new RootItem(TestProvidersAccess.PICKLES, null);
+            new RootItem(TestProvidersAccess.PICKLES, null, false);
     private static final RootItem SDCARD_DEFAULT_USER =
-            new RootItem(TestProvidersAccess.SD_CARD, null);
+            new RootItem(TestProvidersAccess.SD_CARD, null, false);
 
     private static final RootItem DOWNLOADS_OTHER_USER =
-            new RootItem(TestProvidersAccess.OtherUser.DOWNLOADS, null);
+            new RootItem(TestProvidersAccess.OtherUser.DOWNLOADS, null, false);
     private static final RootItem HOME_OTHER_USER =
-            new RootItem(TestProvidersAccess.OtherUser.HOME, null);
+            new RootItem(TestProvidersAccess.OtherUser.HOME, null, false);
     private static final RootItem IMAGE_OTHER_USER =
-            new RootItem(TestProvidersAccess.OtherUser.IMAGE, null);
+            new RootItem(TestProvidersAccess.OtherUser.IMAGE, null, false);
     private static final RootItem PICKLES_OTHER_USER =
-            new RootItem(TestProvidersAccess.OtherUser.PICKLES, null);
+            new RootItem(TestProvidersAccess.OtherUser.PICKLES, null, false);
 
     private RootItemListBuilder mBuilder;
 
diff --git a/tests/unit/com/android/documentsui/sidebar/RootsFragmentTest.java b/tests/unit/com/android/documentsui/sidebar/RootsFragmentTest.java
index af9a66cc6..8a342ba51 100644
--- a/tests/unit/com/android/documentsui/sidebar/RootsFragmentTest.java
+++ b/tests/unit/com/android/documentsui/sidebar/RootsFragmentTest.java
@@ -70,7 +70,9 @@ public class RootsFragmentTest {
     public void testSortLoadResult_WithCorrectOrder() {
         List<Item> items = mRootsFragment.sortLoadResult(createFakeRootInfoList(),
                 null /* excludePackage */, null /* handlerAppIntent */, new TestProvidersAccess(),
-                UserId.DEFAULT_USER, Collections.singletonList(UserId.DEFAULT_USER));
+                UserId.DEFAULT_USER,
+                Collections.singletonList(UserId.DEFAULT_USER),
+                /* maybeShowBadge */ false);
         assertTrue(assertSortedResult(items));
     }
 
@@ -81,11 +83,11 @@ public class RootsFragmentTest {
         final RootsFragment.ItemComparator comp = new RootsFragment.ItemComparator(testPackageName);
         final List<Item> rootList = new ArrayList<>();
         rootList.add(new RootItem(TestProvidersAccess.HAMMY, null /* actionHandler */,
-                errorTestPackageName));
+                errorTestPackageName, /* maybeShowBadge= */ false));
         rootList.add(new RootItem(TestProvidersAccess.INSPECTOR, null /* actionHandler */,
-                errorTestPackageName));
+                errorTestPackageName, /* maybeShowBadge= */ false));
         rootList.add(new RootItem(TestProvidersAccess.PICKLES, null /* actionHandler */,
-                testPackageName));
+                testPackageName, /* maybeShowBadge= */ false));
         Collections.sort(rootList, comp);
 
         assertEquals(rootList.get(0).title, TestProvidersAccess.PICKLES.title);
@@ -99,7 +101,7 @@ public class RootsFragmentTest {
         final RootsFragment.ItemComparator comp = new RootsFragment.ItemComparator(testPackageName);
         final List<Item> rootList = new ArrayList<>();
         rootList.add(new RootItem(TestProvidersAccess.HAMMY, null /* actionHandler */,
-                testPackageName));
+                testPackageName, /* maybeShowBadge= */ false));
 
         final ResolveInfo info = TestResolveInfo.create();
         info.activityInfo.packageName = testPackageName;
@@ -107,7 +109,7 @@ public class RootsFragmentTest {
         rootList.add(new AppItem(info, TestProvidersAccess.PICKLES.title, UserId.DEFAULT_USER,
                 null /* actionHandler */));
         rootList.add(new RootAndAppItem(TestProvidersAccess.INSPECTOR, info,
-                null /* actionHandler */));
+                null /* actionHandler */, /* maybeShowBadge= */ false));
 
         Collections.sort(rootList, comp);