From 26d226301c56fb53f6eddb60dc89212a747c06a2 Mon Sep 17 00:00:00 2001
From: Mady Mellor <madym@google.com>
Date: Wed, 31 Mar 2021 12:18:32 -0700
Subject: [PATCH] Fix bubbles for workprofile

- use getPackageManagerForUser when looking up app info
- add test for it

Bug: 184041127
Test: atest BubblesTest
Test: - have a managed work profile
      - install the bubbles app *only* for the work profile
      - make some bubbles
      => Notice bubbles appear
      - dismiss all the bubbles
      - restart the device, add a bubble
      - open the bubble, navigate to the bubble overflow
      => notice the previously dismissed workprofile bubbles
        are in the overflow
Change-Id: I479bb717b3c365346682331b0def7170ed1a791b
---
 .../com/android/wm/shell/bubbles/Bubble.java  |  8 +-
 .../wm/shell/bubbles/BubbleIconFactory.java   |  5 +-
 .../wm/shell/bubbles/BubbleViewInfoTask.java  | 10 +-
 .../android/systemui/wmshell/BubblesTest.java | 92 ++++++++++++++++++-
 4 files changed, 107 insertions(+), 8 deletions(-)

diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/Bubble.java b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/Bubble.java
index 0a15d8468983b..f6e92ef0e8eae 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/Bubble.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/Bubble.java
@@ -165,7 +165,8 @@ public class Bubble implements BubbleViewProvider {
      * Create a bubble with limited information based on given {@link ShortcutInfo}.
      * Note: Currently this is only being used when the bubble is persisted to disk.
      */
-    Bubble(@NonNull final String key, @NonNull final ShortcutInfo shortcutInfo,
+    @VisibleForTesting(visibility = PRIVATE)
+    public Bubble(@NonNull final String key, @NonNull final ShortcutInfo shortcutInfo,
             final int desiredHeight, final int desiredHeightResId, @Nullable final String title,
             int taskId, @Nullable final String locus, Executor mainExecutor) {
         Objects.requireNonNull(key);
@@ -188,7 +189,7 @@ public class Bubble implements BubbleViewProvider {
     }
 
     @VisibleForTesting(visibility = PRIVATE)
-    Bubble(@NonNull final BubbleEntry entry,
+    public Bubble(@NonNull final BubbleEntry entry,
             @Nullable final Bubbles.SuppressionChangedListener listener,
             final Bubbles.PendingIntentCanceledListener intentCancelListener,
             Executor mainExecutor) {
@@ -718,7 +719,8 @@ public class Bubble implements BubbleViewProvider {
 
     private int getUid(final Context context) {
         if (mAppUid != -1) return mAppUid;
-        final PackageManager pm = context.getPackageManager();
+        final PackageManager pm = BubbleController.getPackageManagerForUser(context,
+                mUser.getIdentifier());
         if (pm == null) return -1;
         try {
             final ApplicationInfo info = pm.getApplicationInfo(mShortcutInfo.getPackage(), 0);
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleIconFactory.java b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleIconFactory.java
index fe3f9ef6aa5f2..e64ed6a0836c1 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleIconFactory.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleIconFactory.java
@@ -29,6 +29,8 @@ import android.graphics.drawable.AdaptiveIconDrawable;
 import android.graphics.drawable.Drawable;
 import android.graphics.drawable.Icon;
 
+import androidx.annotation.VisibleForTesting;
+
 import com.android.launcher3.icons.BaseIconFactory;
 import com.android.launcher3.icons.BitmapInfo;
 import com.android.launcher3.icons.ShadowGenerator;
@@ -39,11 +41,12 @@ import com.android.wm.shell.R;
  * We are not using Launcher's IconFactory because bubbles only runs on the UI thread,
  * so there is no need to manage a pool across multiple threads.
  */
+@VisibleForTesting
 public class BubbleIconFactory extends BaseIconFactory {
 
     private int mBadgeSize;
 
-    protected BubbleIconFactory(Context context) {
+    public BubbleIconFactory(Context context) {
         super(context, context.getResources().getConfiguration().densityDpi,
                 context.getResources().getDimensionPixelSize(R.dimen.individual_bubble_size));
         mBadgeSize = mContext.getResources().getDimensionPixelSize(
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleViewInfoTask.java b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleViewInfoTask.java
index 5cdce886a9e8b..fc53ef26dbd95 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleViewInfoTask.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleViewInfoTask.java
@@ -39,6 +39,7 @@ import android.util.PathParser;
 import android.view.LayoutInflater;
 
 import androidx.annotation.Nullable;
+import androidx.annotation.VisibleForTesting;
 
 import com.android.internal.graphics.ColorUtils;
 import com.android.launcher3.icons.BitmapInfo;
@@ -118,7 +119,8 @@ public class BubbleViewInfoTask extends AsyncTask<Void, Void, BubbleViewInfoTask
     /**
      * Info necessary to render a bubble.
      */
-    static class BubbleViewInfo {
+    @VisibleForTesting
+    public static class BubbleViewInfo {
         BadgedImageView imageView;
         BubbleExpandedView expandedView;
         ShortcutInfo shortcutInfo;
@@ -129,8 +131,9 @@ public class BubbleViewInfoTask extends AsyncTask<Void, Void, BubbleViewInfoTask
         Path dotPath;
         Bubble.FlyoutMessage flyoutMessage;
 
+        @VisibleForTesting
         @Nullable
-        static BubbleViewInfo populate(Context c, BubbleController controller,
+        public static BubbleViewInfo populate(Context c, BubbleController controller,
                 BubbleStackView stackView, BubbleIconFactory iconFactory, Bubble b,
                 boolean skipInflation) {
             BubbleViewInfo info = new BubbleViewInfo();
@@ -152,7 +155,8 @@ public class BubbleViewInfoTask extends AsyncTask<Void, Void, BubbleViewInfoTask
             }
 
             // App name & app icon
-            PackageManager pm = c.getPackageManager();
+            PackageManager pm = BubbleController.getPackageManagerForUser(c,
+                    b.getUser().getIdentifier());
             ApplicationInfo appInfo;
             Drawable badgedIcon;
             Drawable appIcon;
diff --git a/packages/SystemUI/tests/src/com/android/systemui/wmshell/BubblesTest.java b/packages/SystemUI/tests/src/com/android/systemui/wmshell/BubblesTest.java
index 3e5fb7347b136..9fc1df73d8b35 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/wmshell/BubblesTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/wmshell/BubblesTest.java
@@ -17,6 +17,7 @@
 package com.android.systemui.wmshell;
 
 import static android.app.Notification.FLAG_BUBBLE;
+import static android.app.PendingIntent.FLAG_MUTABLE;
 import static android.service.notification.NotificationListenerService.REASON_APP_CANCEL;
 import static android.service.notification.NotificationListenerService.REASON_CANCEL;
 import static android.service.notification.NotificationListenerService.REASON_CANCEL_ALL;
@@ -44,7 +45,14 @@ import android.app.IActivityManager;
 import android.app.INotificationManager;
 import android.app.Notification;
 import android.app.PendingIntent;
+import android.content.Context;
+import android.content.Intent;
+import android.content.pm.ApplicationInfo;
 import android.content.pm.LauncherApps;
+import android.content.pm.PackageManager;
+import android.graphics.Rect;
+import android.graphics.drawable.Drawable;
+import android.graphics.drawable.Icon;
 import android.hardware.display.AmbientDisplayConfiguration;
 import android.hardware.face.FaceManager;
 import android.os.Handler;
@@ -78,6 +86,7 @@ import com.android.systemui.statusbar.notification.NotificationEntryManager;
 import com.android.systemui.statusbar.notification.NotificationFilter;
 import com.android.systemui.statusbar.notification.collection.NotifPipeline;
 import com.android.systemui.statusbar.notification.collection.NotificationEntry;
+import com.android.systemui.statusbar.notification.collection.NotificationEntryBuilder;
 import com.android.systemui.statusbar.notification.collection.legacy.NotificationGroupManagerLegacy;
 import com.android.systemui.statusbar.notification.row.ExpandableNotificationRow;
 import com.android.systemui.statusbar.notification.row.NotificationTestHelper;
@@ -90,15 +99,18 @@ import com.android.systemui.statusbar.policy.BatteryController;
 import com.android.systemui.statusbar.policy.ConfigurationController;
 import com.android.systemui.statusbar.policy.HeadsUpManager;
 import com.android.systemui.statusbar.policy.ZenModeController;
+import com.android.wm.shell.R;
 import com.android.wm.shell.ShellTaskOrganizer;
 import com.android.wm.shell.WindowManagerShellWrapper;
 import com.android.wm.shell.bubbles.Bubble;
 import com.android.wm.shell.bubbles.BubbleData;
 import com.android.wm.shell.bubbles.BubbleDataRepository;
 import com.android.wm.shell.bubbles.BubbleEntry;
+import com.android.wm.shell.bubbles.BubbleIconFactory;
 import com.android.wm.shell.bubbles.BubbleLogger;
 import com.android.wm.shell.bubbles.BubbleOverflow;
 import com.android.wm.shell.bubbles.BubbleStackView;
+import com.android.wm.shell.bubbles.BubbleViewInfoTask;
 import com.android.wm.shell.bubbles.Bubbles;
 import com.android.wm.shell.common.FloatingContentCoordinator;
 import com.android.wm.shell.common.ShellExecutor;
@@ -1132,7 +1144,7 @@ public class BubblesTest extends SysuiTestCase {
 
         // Switch users
         mBubbleController.onUserChanged(secondUserId);
-        assertThat(mBubbleData.getOverflowBubbles().isEmpty());
+        assertThat(mBubbleData.getOverflowBubbles()).isEmpty();
 
         // Give this user some bubbles
         mBubbleController.updateBubble(mBubbleEntryUser11);
@@ -1145,6 +1157,75 @@ public class BubblesTest extends SysuiTestCase {
         assertThat(mBubbleData.getOverflowBubbleWithKey(mBubbleEntry2User11.getKey())).isNotNull();
     }
 
+
+    /**
+     * Verifies that the package manager for the user is used when loading info for the bubble.
+     */
+    @Test
+    public void test_bubbleViewInfoGetPackageForUser() throws Exception {
+        final int workProfileUserId = 10;
+        final UserHandle workUser = new UserHandle(workProfileUserId);
+        final String workPkg = "work.pkg";
+
+        final Bubble bubble = createBubble(workProfileUserId, workPkg);
+        assertEquals(workProfileUserId, bubble.getUser().getIdentifier());
+
+        final Context context = setUpContextWithPackageManager(workPkg, null /* AppInfo */);
+        when(context.getResources()).thenReturn(mContext.getResources());
+        final Context userContext = setUpContextWithPackageManager(workPkg,
+                mock(ApplicationInfo.class));
+
+        // If things are working correctly, StatusBar.getPackageManagerForUser will call this
+        when(context.createPackageContextAsUser(eq(workPkg), anyInt(), eq(workUser)))
+                .thenReturn(userContext);
+
+        BubbleViewInfoTask.BubbleViewInfo info = BubbleViewInfoTask.BubbleViewInfo.populate(context,
+                mBubbleController,
+                mBubbleController.getStackView(),
+                new BubbleIconFactory(mContext),
+                bubble,
+                true /* skipInflation */);
+        verify(userContext, times(1)).getPackageManager();
+        verify(context, times(1)).createPackageContextAsUser(eq(workPkg),
+                eq(Context.CONTEXT_RESTRICTED),
+                eq(workUser));
+        assertNotNull(info);
+    }
+
+    /** Creates a bubble using the userId and package. */
+    private Bubble createBubble(int userId, String pkg) {
+        final UserHandle userHandle = new UserHandle(userId);
+        NotificationEntry workEntry = new NotificationEntryBuilder()
+                .setPkg(pkg)
+                .setUser(userHandle)
+                .build();
+        workEntry.setBubbleMetadata(getMetadata());
+        workEntry.setFlagBubble(true);
+
+        return new Bubble(BubblesManager.notifToBubbleEntry(workEntry),
+                null,
+                mock(Bubbles.PendingIntentCanceledListener.class), new SyncExecutor());
+    }
+
+    /** Creates a context that will return a PackageManager with specific AppInfo. */
+    private Context setUpContextWithPackageManager(String pkg, ApplicationInfo info)
+            throws Exception {
+        final PackageManager pm = mock(PackageManager.class);
+        when(pm.getApplicationInfo(eq(pkg), anyInt())).thenReturn(info);
+
+        if (info != null) {
+            Drawable d = mock(Drawable.class);
+            when(d.getBounds()).thenReturn(new Rect());
+            when(pm.getApplicationIcon(anyString())).thenReturn(d);
+            when(pm.getUserBadgedIcon(any(), any())).thenReturn(d);
+        }
+
+        final Context context = mock(Context.class);
+        when(context.getPackageName()).thenReturn(pkg);
+        when(context.getPackageManager()).thenReturn(pm);
+        return context;
+    }
+
     /**
      * Sets the bubble metadata flags for this entry. These ]flags are normally set by
      * NotificationManagerService when the notification is sent, however, these tests do not
@@ -1161,4 +1242,13 @@ public class BubblesTest extends SysuiTestCase {
         }
         bubbleMetadata.setFlags(flags);
     }
+
+    private Notification.BubbleMetadata getMetadata() {
+        Intent target = new Intent(mContext, BubblesTestActivity.class);
+        PendingIntent bubbleIntent = PendingIntent.getActivity(mContext, 0, target, FLAG_MUTABLE);
+
+        return new Notification.BubbleMetadata.Builder(bubbleIntent,
+                Icon.createWithResource(mContext, R.drawable.bubble_ic_create_bubble))
+                .build();
+    }
 }