From edddfb37ef05871f3d6411440f5140d23348db67 Mon Sep 17 00:00:00 2001
From: Mady Mellor <madym@google.com>
Date: Mon, 29 Mar 2021 16:48:30 -0700
Subject: [PATCH] DO NOT MERGE Fix bubbles for workprofile

- use getPackageManagerForUser when looking up app info
- use the userId on the bubble, not the currentUserId from
  BubbleController when writing to XML as the currentUserId
  in BubbleController does not account for workprofile.

Test: atest SystemUITests
Test: atest BubbleControllerTest
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
Bug: 184041127
Change-Id: Iefc4e4e80a569a4f3b4ac3b8a1268de1b475409b
---
 .../com/android/systemui/bubbles/Bubble.java  |  4 +-
 .../systemui/bubbles/BubbleController.java    |  6 +-
 .../systemui/bubbles/BubbleDataRepository.kt  | 15 ++--
 .../systemui/bubbles/BubbleViewInfoTask.java  |  4 +-
 .../bubbles/BubbleControllerTest.java         | 87 +++++++++++++++++++
 5 files changed, 103 insertions(+), 13 deletions(-)

diff --git a/packages/SystemUI/src/com/android/systemui/bubbles/Bubble.java b/packages/SystemUI/src/com/android/systemui/bubbles/Bubble.java
index b71e3adae8acf..ed9b904e5238d 100644
--- a/packages/SystemUI/src/com/android/systemui/bubbles/Bubble.java
+++ b/packages/SystemUI/src/com/android/systemui/bubbles/Bubble.java
@@ -43,6 +43,7 @@ import com.android.internal.annotations.VisibleForTesting;
 import com.android.internal.logging.InstanceId;
 import com.android.systemui.shared.system.SysUiStatsLog;
 import com.android.systemui.statusbar.notification.collection.NotificationEntry;
+import com.android.systemui.statusbar.phone.StatusBar;
 
 import java.io.FileDescriptor;
 import java.io.PrintWriter;
@@ -623,7 +624,8 @@ class Bubble implements BubbleViewProvider {
 
     private int getUid(final Context context) {
         if (mAppUid != -1) return mAppUid;
-        final PackageManager pm = context.getPackageManager();
+        final PackageManager pm = StatusBar.getPackageManagerForUser(context,
+                mUser.getIdentifier());
         if (pm == null) return -1;
         try {
             final ApplicationInfo info = pm.getApplicationInfo(mShortcutInfo.getPackage(), 0);
diff --git a/packages/SystemUI/src/com/android/systemui/bubbles/BubbleController.java b/packages/SystemUI/src/com/android/systemui/bubbles/BubbleController.java
index 5deae925ba30e..fb819f04e0e1f 100644
--- a/packages/SystemUI/src/com/android/systemui/bubbles/BubbleController.java
+++ b/packages/SystemUI/src/com/android/systemui/bubbles/BubbleController.java
@@ -1386,10 +1386,10 @@ public class BubbleController implements ConfigurationController.ConfigurationLi
                     }
                 }
             }
-            mDataRepository.removeBubbles(mCurrentUserId, bubblesToBeRemovedFromRepository);
+            mDataRepository.removeBubbles(bubblesToBeRemovedFromRepository);
 
             if (update.addedBubble != null && mStackView != null) {
-                mDataRepository.addBubble(mCurrentUserId, update.addedBubble);
+                mDataRepository.addBubble(update.addedBubble);
                 mStackView.addBubble(update.addedBubble);
             }
 
@@ -1400,7 +1400,7 @@ public class BubbleController implements ConfigurationController.ConfigurationLi
             // At this point, the correct bubbles are inflated in the stack.
             // Make sure the order in bubble data is reflected in bubble row.
             if (update.orderChanged && mStackView != null) {
-                mDataRepository.addBubbles(mCurrentUserId, update.bubbles);
+                mDataRepository.addBubbles(update.bubbles);
                 mStackView.updateBubbleOrder(update.bubbles);
             }
 
diff --git a/packages/SystemUI/src/com/android/systemui/bubbles/BubbleDataRepository.kt b/packages/SystemUI/src/com/android/systemui/bubbles/BubbleDataRepository.kt
index db64a13f3df34..a363208e8f101 100644
--- a/packages/SystemUI/src/com/android/systemui/bubbles/BubbleDataRepository.kt
+++ b/packages/SystemUI/src/com/android/systemui/bubbles/BubbleDataRepository.kt
@@ -16,7 +16,6 @@
 package com.android.systemui.bubbles
 
 import android.annotation.SuppressLint
-import android.annotation.UserIdInt
 import android.content.pm.LauncherApps
 import android.content.pm.LauncherApps.ShortcutQuery.FLAG_MATCH_DYNAMIC
 import android.content.pm.LauncherApps.ShortcutQuery.FLAG_MATCH_PINNED_BY_ANY_LAUNCHER
@@ -51,31 +50,31 @@ internal class BubbleDataRepository @Inject constructor(
      * Adds the bubble in memory, then persists the snapshot after adding the bubble to disk
      * asynchronously.
      */
-    fun addBubble(@UserIdInt userId: Int, bubble: Bubble) = addBubbles(userId, listOf(bubble))
+    fun addBubble(bubble: Bubble) = addBubbles(listOf(bubble))
 
     /**
      * Adds the bubble in memory, then persists the snapshot after adding the bubble to disk
      * asynchronously.
      */
-    fun addBubbles(@UserIdInt userId: Int, bubbles: List<Bubble>) {
+    fun addBubbles(bubbles: List<Bubble>) {
         if (DEBUG) Log.d(TAG, "adding ${bubbles.size} bubbles")
-        val entities = transform(userId, bubbles).also(volatileRepository::addBubbles)
+        val entities = transform(bubbles).also(volatileRepository::addBubbles)
         if (entities.isNotEmpty()) persistToDisk()
     }
 
     /**
      * Removes the bubbles from memory, then persists the snapshot to disk asynchronously.
      */
-    fun removeBubbles(@UserIdInt userId: Int, bubbles: List<Bubble>) {
+    fun removeBubbles(bubbles: List<Bubble>) {
         if (DEBUG) Log.d(TAG, "removing ${bubbles.size} bubbles")
-        val entities = transform(userId, bubbles).also(volatileRepository::removeBubbles)
+        val entities = transform(bubbles).also(volatileRepository::removeBubbles)
         if (entities.isNotEmpty()) persistToDisk()
     }
 
-    private fun transform(userId: Int, bubbles: List<Bubble>): List<BubbleEntity> {
+    private fun transform(bubbles: List<Bubble>): List<BubbleEntity> {
         return bubbles.mapNotNull { b ->
             BubbleEntity(
-                    userId,
+                    b.user.identifier,
                     b.packageName,
                     b.metadataShortcutId ?: return@mapNotNull null,
                     b.key,
diff --git a/packages/SystemUI/src/com/android/systemui/bubbles/BubbleViewInfoTask.java b/packages/SystemUI/src/com/android/systemui/bubbles/BubbleViewInfoTask.java
index 1929fc4e9dbfe..a6cf69a0f0a9b 100644
--- a/packages/SystemUI/src/com/android/systemui/bubbles/BubbleViewInfoTask.java
+++ b/packages/SystemUI/src/com/android/systemui/bubbles/BubbleViewInfoTask.java
@@ -48,6 +48,7 @@ import com.android.internal.graphics.ColorUtils;
 import com.android.launcher3.icons.BitmapInfo;
 import com.android.systemui.R;
 import com.android.systemui.statusbar.notification.collection.NotificationEntry;
+import com.android.systemui.statusbar.phone.StatusBar;
 
 import java.lang.ref.WeakReference;
 import java.util.List;
@@ -146,7 +147,8 @@ public class BubbleViewInfoTask extends AsyncTask<Void, Void, BubbleViewInfoTask
             }
 
             // App name & app icon
-            PackageManager pm = c.getPackageManager();
+            PackageManager pm = StatusBar.getPackageManagerForUser(
+                    c, b.getUser().getIdentifier());
             ApplicationInfo appInfo;
             Drawable badgedIcon;
             Drawable appIcon;
diff --git a/packages/SystemUI/tests/src/com/android/systemui/bubbles/BubbleControllerTest.java b/packages/SystemUI/tests/src/com/android/systemui/bubbles/BubbleControllerTest.java
index 1ad88560bf9c7..8fc2d1ad0fe5e 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/bubbles/BubbleControllerTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/bubbles/BubbleControllerTest.java
@@ -44,11 +44,19 @@ 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;
 import android.os.PowerManager;
+import android.os.UserHandle;
 import android.service.dreams.IDreamManager;
 import android.service.notification.NotificationListenerService;
 import android.service.notification.ZenModeConfig;
@@ -77,6 +85,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.row.ExpandableNotificationRow;
 import com.android.systemui.statusbar.notification.row.NotificationTestHelper;
 import com.android.systemui.statusbar.phone.DozeParameters;
@@ -89,6 +98,7 @@ 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.systemui.tests.R;
 import com.android.systemui.util.FloatingContentCoordinator;
 
 import com.google.common.collect.ImmutableList;
@@ -1022,6 +1032,74 @@ public class BubbleControllerTest extends SysuiTestCase {
         assertThat(mBubbleData.hasBubbleInStackWithKey(mRow.getEntry().getKey())).isFalse();
     }
 
+    /**
+     * 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.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(workEntry,
+                null,
+                mock(BubbleController.PendingIntentCanceledListener.class));
+    }
+
+    /** 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
@@ -1038,4 +1116,13 @@ public class BubbleControllerTest extends SysuiTestCase {
         }
         bubbleMetadata.setFlags(flags);
     }
+
+    private Notification.BubbleMetadata getMetadata() {
+        Intent target = new Intent(mContext, BubblesTestActivity.class);
+        PendingIntent bubbleIntent = PendingIntent.getActivity(mContext, 0, target, 0);
+
+        return new Notification.BubbleMetadata.Builder(bubbleIntent,
+                Icon.createWithResource(mContext, R.drawable.android))
+                .build();
+    }
 }