Merge "Add multi-users support to app bubbles" into udc-dev am: c7211bdde2

Original change: https://googleplex-android-review.googlesource.com/c/platform/frameworks/base/+/22648901

Change-Id: Ie8459cddb8d49d46a6c3d069a195dd393b3ee8b8
Signed-off-by: Automerger Merge Worker <android-build-automerger-merge-worker@system.gserviceaccount.com>
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 79e0a48..d3f3958 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
@@ -64,7 +64,11 @@
 public class Bubble implements BubbleViewProvider {
     private static final String TAG = "Bubble";
 
-    public static final String KEY_APP_BUBBLE = "key_app_bubble";
+    /** A string suffix used in app bubbles' {@link #mKey}. */
+    private static final String KEY_APP_BUBBLE = "key_app_bubble";
+
+    /** Whether the bubble is an app bubble. */
+    private final boolean mIsAppBubble;
 
     private final String mKey;
     @Nullable
@@ -181,7 +185,7 @@
     private PendingIntent mDeleteIntent;
 
     /**
-     * Used only for a special bubble in the stack that has the key {@link #KEY_APP_BUBBLE}.
+     * Used only for a special bubble in the stack that has {@link #mIsAppBubble} set to true.
      * There can only be one of these bubbles in the stack and this intent will be populated for
      * that bubble.
      */
@@ -216,24 +220,54 @@
         mMainExecutor = mainExecutor;
         mTaskId = taskId;
         mBubbleMetadataFlagListener = listener;
+        mIsAppBubble = false;
     }
 
-    public Bubble(Intent intent,
+    private Bubble(
+            Intent intent,
             UserHandle user,
             @Nullable Icon icon,
+            boolean isAppBubble,
+            String key,
             Executor mainExecutor) {
-        mKey = KEY_APP_BUBBLE;
         mGroupKey = null;
         mLocusId = null;
         mFlags = 0;
         mUser = user;
         mIcon = icon;
+        mIsAppBubble = isAppBubble;
+        mKey = key;
         mShowBubbleUpdateDot = false;
         mMainExecutor = mainExecutor;
         mTaskId = INVALID_TASK_ID;
         mAppIntent = intent;
         mDesiredHeight = Integer.MAX_VALUE;
         mPackageName = intent.getPackage();
+
+    }
+
+    /** Creates an app bubble. */
+    public static Bubble createAppBubble(
+            Intent intent,
+            UserHandle user,
+            @Nullable Icon icon,
+            Executor mainExecutor) {
+        return new Bubble(intent,
+                user,
+                icon,
+                /* isAppBubble= */ true,
+                /* key= */ getAppBubbleKeyForApp(intent.getPackage(), user),
+                mainExecutor);
+    }
+
+    /**
+     * Returns the key for an app bubble from an app with package name, {@code packageName} on an
+     * Android user, {@code user}.
+     */
+    public static String getAppBubbleKeyForApp(String packageName, UserHandle user) {
+        Objects.requireNonNull(packageName);
+        Objects.requireNonNull(user);
+        return KEY_APP_BUBBLE + ":" + user.getIdentifier()  + ":" + packageName;
     }
 
     @VisibleForTesting(visibility = PRIVATE)
@@ -241,6 +275,7 @@
             final Bubbles.BubbleMetadataFlagListener listener,
             final Bubbles.PendingIntentCanceledListener intentCancelListener,
             Executor mainExecutor) {
+        mIsAppBubble = false;
         mKey = entry.getKey();
         mGroupKey = entry.getGroupKey();
         mLocusId = entry.getLocusId();
@@ -815,7 +850,7 @@
     }
 
     boolean isAppBubble() {
-        return KEY_APP_BUBBLE.equals(mKey);
+        return mIsAppBubble;
     }
 
     Intent getSettingsIntent(final Context context) {
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleController.java b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleController.java
index c407b06..21f02b1 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleController.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleController.java
@@ -24,7 +24,6 @@
 import static android.view.View.VISIBLE;
 import static android.view.WindowManager.LayoutParams.LAYOUT_IN_DISPLAY_CUTOUT_MODE_ALWAYS;
 
-import static com.android.wm.shell.bubbles.Bubble.KEY_APP_BUBBLE;
 import static com.android.wm.shell.bubbles.BubbleDebugConfig.DEBUG_BUBBLE_CONTROLLER;
 import static com.android.wm.shell.bubbles.BubbleDebugConfig.DEBUG_BUBBLE_GESTURE;
 import static com.android.wm.shell.bubbles.BubbleDebugConfig.TAG_BUBBLES;
@@ -1193,14 +1192,15 @@
             return;
         }
 
+        String appBubbleKey = Bubble.getAppBubbleKeyForApp(intent.getPackage(), user);
         PackageManager packageManager = getPackageManagerForUser(mContext, user.getIdentifier());
-        if (!isResizableActivity(intent, packageManager, KEY_APP_BUBBLE)) return;
+        if (!isResizableActivity(intent, packageManager, appBubbleKey)) return;
 
-        Bubble existingAppBubble = mBubbleData.getBubbleInStackWithKey(KEY_APP_BUBBLE);
+        Bubble existingAppBubble = mBubbleData.getBubbleInStackWithKey(appBubbleKey);
         if (existingAppBubble != null) {
             BubbleViewProvider selectedBubble = mBubbleData.getSelectedBubble();
             if (isStackExpanded()) {
-                if (selectedBubble != null && KEY_APP_BUBBLE.equals(selectedBubble.getKey())) {
+                if (selectedBubble != null && appBubbleKey.equals(selectedBubble.getKey())) {
                     // App bubble is expanded, lets collapse
                     collapseStack();
                 } else {
@@ -1214,7 +1214,7 @@
             }
         } else {
             // App bubble does not exist, lets add and expand it
-            Bubble b = new Bubble(intent, user, icon, mMainExecutor);
+            Bubble b = Bubble.createAppBubble(intent, user, icon, mMainExecutor);
             b.setShouldAutoExpand(true);
             inflateAndAdd(b, /* suppressFlyout= */ true, /* showInShade= */ false);
         }
@@ -1247,8 +1247,8 @@
     }
 
     /** Sets the app bubble's taskId which is cached for SysUI. */
-    public void setAppBubbleTaskId(int taskId) {
-        mImpl.mCachedState.setAppBubbleTaskId(taskId);
+    public void setAppBubbleTaskId(String key, int taskId) {
+        mImpl.mCachedState.setAppBubbleTaskId(key, taskId);
     }
 
     /**
@@ -2045,7 +2045,8 @@
             private HashSet<String> mSuppressedBubbleKeys = new HashSet<>();
             private HashMap<String, String> mSuppressedGroupToNotifKeys = new HashMap<>();
             private HashMap<String, Bubble> mShortcutIdToBubble = new HashMap<>();
-            private int mAppBubbleTaskId = INVALID_TASK_ID;
+
+            private HashMap<String, Integer> mAppBubbleTaskIds = new HashMap();
 
             private ArrayList<Bubble> mTmpBubbles = new ArrayList<>();
 
@@ -2077,20 +2078,20 @@
 
                 mSuppressedBubbleKeys.clear();
                 mShortcutIdToBubble.clear();
-                mAppBubbleTaskId = INVALID_TASK_ID;
+                mAppBubbleTaskIds.clear();
                 for (Bubble b : mTmpBubbles) {
                     mShortcutIdToBubble.put(b.getShortcutId(), b);
                     updateBubbleSuppressedState(b);
 
-                    if (KEY_APP_BUBBLE.equals(b.getKey())) {
-                        mAppBubbleTaskId = b.getTaskId();
+                    if (b.isAppBubble()) {
+                        mAppBubbleTaskIds.put(b.getKey(), b.getTaskId());
                     }
                 }
             }
 
             /** Sets the app bubble's taskId which is cached for SysUI. */
-            synchronized void setAppBubbleTaskId(int taskId) {
-                mAppBubbleTaskId = taskId;
+            synchronized void setAppBubbleTaskId(String key, int taskId) {
+                mAppBubbleTaskIds.put(key, taskId);
             }
 
             /**
@@ -2143,7 +2144,7 @@
                     pw.println("   suppressing: " + key);
                 }
 
-                pw.print("mAppBubbleTaskId: " + mAppBubbleTaskId);
+                pw.print("mAppBubbleTaskIds: " + mAppBubbleTaskIds.values());
             }
         }
 
@@ -2205,7 +2206,7 @@
 
         @Override
         public boolean isAppBubbleTaskId(int taskId) {
-            return mCachedState.mAppBubbleTaskId == taskId;
+            return mCachedState.mAppBubbleTaskIds.values().contains(taskId);
         }
 
         @Override
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleData.java b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleData.java
index 92b969b..cc8f50e 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleData.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleData.java
@@ -17,7 +17,6 @@
 
 import static com.android.internal.annotations.VisibleForTesting.Visibility.PACKAGE;
 import static com.android.internal.annotations.VisibleForTesting.Visibility.PRIVATE;
-import static com.android.wm.shell.bubbles.Bubble.KEY_APP_BUBBLE;
 import static com.android.wm.shell.bubbles.BubbleDebugConfig.DEBUG_BUBBLE_DATA;
 import static com.android.wm.shell.bubbles.BubbleDebugConfig.TAG_BUBBLES;
 import static com.android.wm.shell.bubbles.BubbleDebugConfig.TAG_WITH_CLASS_NAME;
@@ -780,7 +779,7 @@
                 || !(reason == Bubbles.DISMISS_AGED
                 || reason == Bubbles.DISMISS_USER_GESTURE
                 || reason == Bubbles.DISMISS_RELOAD_FROM_DISK)
-                || KEY_APP_BUBBLE.equals(bubble.getKey())) {
+                || bubble.isAppBubble()) {
             return;
         }
         if (DEBUG_BUBBLE_DATA) {
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleExpandedView.java b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleExpandedView.java
index 684a23a..6c482c8 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleExpandedView.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleExpandedView.java
@@ -287,9 +287,9 @@
             // The taskId is saved to use for removeTask, preventing appearance in recent tasks.
             mTaskId = taskId;
 
-            if (Bubble.KEY_APP_BUBBLE.equals(getBubbleKey())) {
+            if (mBubble != null && mBubble.isAppBubble()) {
                 // Let the controller know sooner what the taskId is.
-                mController.setAppBubbleTaskId(mTaskId);
+                mController.setAppBubbleTaskId(mBubble.getKey(), mTaskId);
             }
 
             // With the task org, the taskAppeared callback will only happen once the task has
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/bubbles/BubbleDataTest.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/bubbles/BubbleDataTest.java
index 919bf06..4a55429 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/bubbles/BubbleDataTest.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/bubbles/BubbleDataTest.java
@@ -16,8 +16,6 @@
 
 package com.android.wm.shell.bubbles;
 
-import static com.android.wm.shell.bubbles.Bubble.KEY_APP_BUBBLE;
-
 import static com.google.common.truth.Truth.assertThat;
 import static com.google.common.truth.Truth.assertWithMessage;
 
@@ -185,7 +183,10 @@
 
         Intent appBubbleIntent = new Intent(mContext, BubblesTestActivity.class);
         appBubbleIntent.setPackage(mContext.getPackageName());
-        mAppBubble = new Bubble(appBubbleIntent, new UserHandle(1), mock(Icon.class),
+        mAppBubble = Bubble.createAppBubble(
+                appBubbleIntent,
+                new UserHandle(1),
+                mock(Icon.class),
                 mMainExecutor);
 
         mPositioner = new TestableBubblePositioner(mContext,
@@ -1101,14 +1102,15 @@
 
     @Test
     public void test_removeAppBubble_skipsOverflow() {
+        String appBubbleKey = mAppBubble.getKey();
         mBubbleData.notificationEntryUpdated(mAppBubble, true /* suppressFlyout*/,
                 false /* showInShade */);
-        assertThat(mBubbleData.getBubbleInStackWithKey(KEY_APP_BUBBLE)).isEqualTo(mAppBubble);
+        assertThat(mBubbleData.getBubbleInStackWithKey(appBubbleKey)).isEqualTo(mAppBubble);
 
-        mBubbleData.dismissBubbleWithKey(KEY_APP_BUBBLE, Bubbles.DISMISS_USER_GESTURE);
+        mBubbleData.dismissBubbleWithKey(appBubbleKey, Bubbles.DISMISS_USER_GESTURE);
 
-        assertThat(mBubbleData.getOverflowBubbleWithKey(KEY_APP_BUBBLE)).isNull();
-        assertThat(mBubbleData.getBubbleInStackWithKey(KEY_APP_BUBBLE)).isNull();
+        assertThat(mBubbleData.getOverflowBubbleWithKey(appBubbleKey)).isNull();
+        assertThat(mBubbleData.getBubbleInStackWithKey(appBubbleKey)).isNull();
     }
 
     private void verifyUpdateReceived() {
diff --git a/packages/SystemUI/src/com/android/systemui/notetask/NoteTaskController.kt b/packages/SystemUI/src/com/android/systemui/notetask/NoteTaskController.kt
index 8f70376..aab898e 100644
--- a/packages/SystemUI/src/com/android/systemui/notetask/NoteTaskController.kt
+++ b/packages/SystemUI/src/com/android/systemui/notetask/NoteTaskController.kt
@@ -85,12 +85,12 @@
     fun onBubbleExpandChanged(isExpanding: Boolean, key: String?) {
         if (!isEnabled) return
 
-        if (key != Bubble.KEY_APP_BUBBLE) return
+        val info = infoReference.getAndSet(null) ?: return
 
-        val info = infoReference.getAndSet(null)
+        if (key != Bubble.getAppBubbleKeyForApp(info.packageName, info.user)) return
 
         // Safe guard mechanism, this callback should only be called for app bubbles.
-        if (info?.launchMode != NoteTaskLaunchMode.AppBubble) return
+        if (info.launchMode != NoteTaskLaunchMode.AppBubble) return
 
         if (isExpanding) {
             logDebug { "onBubbleExpandChanged - expanding: $info" }
@@ -173,7 +173,7 @@
             return
         }
 
-        val info = resolver.resolveInfo(entryPoint, isKeyguardLocked)
+        val info = resolver.resolveInfo(entryPoint, isKeyguardLocked, user)
 
         if (info == null) {
             logDebug { "Default notes app isn't set" }
diff --git a/packages/SystemUI/src/com/android/systemui/notetask/NoteTaskInfo.kt b/packages/SystemUI/src/com/android/systemui/notetask/NoteTaskInfo.kt
index 2b9f0af..a758347 100644
--- a/packages/SystemUI/src/com/android/systemui/notetask/NoteTaskInfo.kt
+++ b/packages/SystemUI/src/com/android/systemui/notetask/NoteTaskInfo.kt
@@ -15,10 +15,13 @@
  */
 package com.android.systemui.notetask
 
+import android.os.UserHandle
+
 /** Contextual information required to launch a Note Task by [NoteTaskController]. */
 data class NoteTaskInfo(
     val packageName: String,
     val uid: Int,
+    val user: UserHandle,
     val entryPoint: NoteTaskEntryPoint? = null,
     val isKeyguardLocked: Boolean = false,
 ) {
diff --git a/packages/SystemUI/src/com/android/systemui/notetask/NoteTaskInfoResolver.kt b/packages/SystemUI/src/com/android/systemui/notetask/NoteTaskInfoResolver.kt
index 616f9b5..89a8526 100644
--- a/packages/SystemUI/src/com/android/systemui/notetask/NoteTaskInfoResolver.kt
+++ b/packages/SystemUI/src/com/android/systemui/notetask/NoteTaskInfoResolver.kt
@@ -25,7 +25,6 @@
 import android.os.UserHandle
 import android.util.Log
 import com.android.systemui.notetask.NoteTaskRoleManagerExt.getDefaultRoleHolderAsUser
-import com.android.systemui.settings.UserTracker
 import javax.inject.Inject
 
 class NoteTaskInfoResolver
@@ -33,15 +32,13 @@
 constructor(
     private val roleManager: RoleManager,
     private val packageManager: PackageManager,
-    private val userTracker: UserTracker,
 ) {
 
     fun resolveInfo(
         entryPoint: NoteTaskEntryPoint? = null,
         isKeyguardLocked: Boolean = false,
+        user: UserHandle,
     ): NoteTaskInfo? {
-        val user = userTracker.userHandle
-
         val packageName = roleManager.getDefaultRoleHolderAsUser(ROLE_NOTES, user)
 
         if (packageName.isNullOrEmpty()) return null
@@ -49,6 +46,7 @@
         return NoteTaskInfo(
             packageName = packageName,
             uid = packageManager.getUidOf(packageName, user),
+            user = user,
             entryPoint = entryPoint,
             isKeyguardLocked = isKeyguardLocked,
         )
diff --git a/packages/SystemUI/tests/src/com/android/systemui/notetask/NoteTaskControllerTest.kt b/packages/SystemUI/tests/src/com/android/systemui/notetask/NoteTaskControllerTest.kt
index a03bc1e..7dc622b 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/notetask/NoteTaskControllerTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/notetask/NoteTaskControllerTest.kt
@@ -98,7 +98,7 @@
         whenever(context.getString(R.string.note_task_button_label))
             .thenReturn(NOTE_TASK_SHORT_LABEL)
         whenever(context.packageManager).thenReturn(packageManager)
-        whenever(resolver.resolveInfo(any(), any())).thenReturn(NOTE_TASK_INFO)
+        whenever(resolver.resolveInfo(any(), any(), any())).thenReturn(NOTE_TASK_INFO)
         whenever(userManager.isUserUnlocked).thenReturn(true)
         whenever(
                 devicePolicyManager.getKeyguardDisabledFeatures(
@@ -142,7 +142,7 @@
             .apply { infoReference.set(expectedInfo) }
             .onBubbleExpandChanged(
                 isExpanding = true,
-                key = Bubble.KEY_APP_BUBBLE,
+                key = Bubble.getAppBubbleKeyForApp(expectedInfo.packageName, expectedInfo.user),
             )
 
         verify(eventLogger).logNoteTaskOpened(expectedInfo)
@@ -157,7 +157,7 @@
             .apply { infoReference.set(expectedInfo) }
             .onBubbleExpandChanged(
                 isExpanding = false,
-                key = Bubble.KEY_APP_BUBBLE,
+                key = Bubble.getAppBubbleKeyForApp(expectedInfo.packageName, expectedInfo.user),
             )
 
         verify(eventLogger).logNoteTaskClosed(expectedInfo)
@@ -172,7 +172,7 @@
             .apply { infoReference.set(expectedInfo) }
             .onBubbleExpandChanged(
                 isExpanding = true,
-                key = Bubble.KEY_APP_BUBBLE,
+                key = Bubble.getAppBubbleKeyForApp(expectedInfo.packageName, expectedInfo.user),
             )
 
         verifyZeroInteractions(context, bubbles, keyguardManager, userManager, eventLogger)
@@ -186,7 +186,7 @@
             .apply { infoReference.set(expectedInfo) }
             .onBubbleExpandChanged(
                 isExpanding = false,
-                key = Bubble.KEY_APP_BUBBLE,
+                key = Bubble.getAppBubbleKeyForApp(expectedInfo.packageName, expectedInfo.user),
             )
 
         verifyZeroInteractions(context, bubbles, keyguardManager, userManager, eventLogger)
@@ -208,7 +208,7 @@
         createNoteTaskController(isEnabled = false)
             .onBubbleExpandChanged(
                 isExpanding = true,
-                key = Bubble.KEY_APP_BUBBLE,
+                key = Bubble.getAppBubbleKeyForApp(NOTE_TASK_INFO.packageName, NOTE_TASK_INFO.user),
             )
 
         verifyZeroInteractions(context, bubbles, keyguardManager, userManager, eventLogger)
@@ -224,7 +224,7 @@
                 isKeyguardLocked = true,
             )
         whenever(keyguardManager.isKeyguardLocked).thenReturn(expectedInfo.isKeyguardLocked)
-        whenever(resolver.resolveInfo(any(), any())).thenReturn(expectedInfo)
+        whenever(resolver.resolveInfo(any(), any(), any())).thenReturn(expectedInfo)
 
         createNoteTaskController()
             .showNoteTask(
@@ -256,9 +256,10 @@
             NOTE_TASK_INFO.copy(
                 entryPoint = NoteTaskEntryPoint.TAIL_BUTTON,
                 isKeyguardLocked = true,
+                user = user10,
             )
         whenever(keyguardManager.isKeyguardLocked).thenReturn(expectedInfo.isKeyguardLocked)
-        whenever(resolver.resolveInfo(any(), any())).thenReturn(expectedInfo)
+        whenever(resolver.resolveInfo(any(), any(), any())).thenReturn(expectedInfo)
 
         createNoteTaskController()
             .showNoteTaskAsUser(
@@ -292,7 +293,7 @@
                 isKeyguardLocked = true,
             )
         whenever(keyguardManager.isKeyguardLocked).thenReturn(expectedInfo.isKeyguardLocked)
-        whenever(resolver.resolveInfo(any(), any())).thenReturn(expectedInfo)
+        whenever(resolver.resolveInfo(any(), any(), any())).thenReturn(expectedInfo)
         whenever(activityManager.getRunningTasks(anyInt()))
             .thenReturn(listOf(NOTE_RUNNING_TASK_INFO))
 
@@ -318,7 +319,7 @@
                 entryPoint = NoteTaskEntryPoint.TAIL_BUTTON,
                 isKeyguardLocked = false,
             )
-        whenever(resolver.resolveInfo(any(), any())).thenReturn(expectedInfo)
+        whenever(resolver.resolveInfo(any(), any(), any())).thenReturn(expectedInfo)
         whenever(keyguardManager.isKeyguardLocked).thenReturn(expectedInfo.isKeyguardLocked)
 
         createNoteTaskController()
@@ -344,7 +345,7 @@
 
     @Test
     fun showNoteTask_intentResolverReturnsNull_shouldShowToast() {
-        whenever(resolver.resolveInfo(any(), any())).thenReturn(null)
+        whenever(resolver.resolveInfo(any(), any(), any())).thenReturn(null)
         val noteTaskController = spy(createNoteTaskController())
         doNothing().whenever(noteTaskController).showNoDefaultNotesAppToast()
 
@@ -384,7 +385,7 @@
                 isKeyguardLocked = true,
             )
         whenever(keyguardManager.isKeyguardLocked).thenReturn(expectedInfo.isKeyguardLocked)
-        whenever(resolver.resolveInfo(any(), any())).thenReturn(expectedInfo)
+        whenever(resolver.resolveInfo(any(), any(), any())).thenReturn(expectedInfo)
 
         createNoteTaskController()
             .showNoteTask(
@@ -717,6 +718,7 @@
             NoteTaskInfo(
                 packageName = NOTE_TASK_PACKAGE_NAME,
                 uid = NOTE_TASK_UID,
+                user = UserHandle.of(0),
             )
         private val NOTE_RUNNING_TASK_INFO =
             ActivityManager.RunningTaskInfo().apply {
diff --git a/packages/SystemUI/tests/src/com/android/systemui/notetask/NoteTaskEventLoggerTest.kt b/packages/SystemUI/tests/src/com/android/systemui/notetask/NoteTaskEventLoggerTest.kt
index a4df346..b4f5528 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/notetask/NoteTaskEventLoggerTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/notetask/NoteTaskEventLoggerTest.kt
@@ -15,6 +15,7 @@
  */
 package com.android.systemui.notetask
 
+import android.os.UserHandle
 import android.test.suitebuilder.annotation.SmallTest
 import androidx.test.runner.AndroidJUnit4
 import com.android.internal.logging.UiEventLogger
@@ -44,7 +45,7 @@
         NoteTaskEventLogger(uiEventLogger)
 
     private fun createNoteTaskInfo(): NoteTaskInfo =
-        NoteTaskInfo(packageName = NOTES_PACKAGE_NAME, uid = NOTES_UID)
+        NoteTaskInfo(packageName = NOTES_PACKAGE_NAME, uid = NOTES_UID, UserHandle.of(0))
 
     @Before
     fun setUp() {
diff --git a/packages/SystemUI/tests/src/com/android/systemui/notetask/NoteTaskInfoResolverTest.kt b/packages/SystemUI/tests/src/com/android/systemui/notetask/NoteTaskInfoResolverTest.kt
index 0c945df..e09c804 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/notetask/NoteTaskInfoResolverTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/notetask/NoteTaskInfoResolverTest.kt
@@ -22,8 +22,6 @@
 import android.test.suitebuilder.annotation.SmallTest
 import androidx.test.runner.AndroidJUnit4
 import com.android.systemui.SysuiTestCase
-import com.android.systemui.settings.FakeUserTracker
-import com.android.systemui.settings.UserTracker
 import com.android.systemui.util.mockito.eq
 import com.android.systemui.util.mockito.whenever
 import com.google.common.truth.Truth.assertThat
@@ -46,14 +44,13 @@
 
     @Mock lateinit var packageManager: PackageManager
     @Mock lateinit var roleManager: RoleManager
-    private val userTracker: UserTracker = FakeUserTracker()
 
     private lateinit var underTest: NoteTaskInfoResolver
 
     @Before
     fun setUp() {
         MockitoAnnotations.initMocks(this)
-        underTest = NoteTaskInfoResolver(roleManager, packageManager, userTracker)
+        underTest = NoteTaskInfoResolver(roleManager, packageManager)
     }
 
     @Test
@@ -72,11 +69,12 @@
             )
             .thenReturn(ApplicationInfo().apply { this.uid = uid })
 
-        val actual = underTest.resolveInfo()
+        val actual = underTest.resolveInfo(user = context.user)
 
         requireNotNull(actual) { "Note task info must not be null" }
         assertThat(actual.packageName).isEqualTo(packageName)
         assertThat(actual.uid).isEqualTo(uid)
+        assertThat(actual.user).isEqualTo(context.user)
     }
 
     @Test
@@ -94,11 +92,12 @@
             )
             .thenThrow(PackageManager.NameNotFoundException(packageName))
 
-        val actual = underTest.resolveInfo()
+        val actual = underTest.resolveInfo(user = context.user)
 
         requireNotNull(actual) { "Note task info must not be null" }
         assertThat(actual.packageName).isEqualTo(packageName)
         assertThat(actual.uid).isEqualTo(0)
+        assertThat(actual.user).isEqualTo(context.user)
     }
 
     @Test
@@ -107,7 +106,7 @@
             emptyList<String>()
         }
 
-        val actual = underTest.resolveInfo()
+        val actual = underTest.resolveInfo(user = context.user)
 
         assertThat(actual).isNull()
     }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/notetask/NoteTaskInfoTest.kt b/packages/SystemUI/tests/src/com/android/systemui/notetask/NoteTaskInfoTest.kt
index 91cd6ae..3435450 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/notetask/NoteTaskInfoTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/notetask/NoteTaskInfoTest.kt
@@ -15,6 +15,7 @@
  */
 package com.android.systemui.notetask
 
+import android.os.UserHandle
 import android.test.suitebuilder.annotation.SmallTest
 import androidx.test.runner.AndroidJUnit4
 import com.android.systemui.SysuiTestCase
@@ -28,7 +29,7 @@
 internal class NoteTaskInfoTest : SysuiTestCase() {
 
     private fun createNoteTaskInfo(): NoteTaskInfo =
-        NoteTaskInfo(packageName = NOTES_PACKAGE_NAME, uid = NOTES_UID)
+        NoteTaskInfo(packageName = NOTES_PACKAGE_NAME, uid = NOTES_UID, UserHandle.of(0))
 
     @Test
     fun launchMode_keyguardLocked_launchModeActivity() {
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 9a99538..bc3a5b7 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/wmshell/BubblesTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/wmshell/BubblesTest.java
@@ -24,7 +24,6 @@
 import static android.service.notification.NotificationListenerService.REASON_GROUP_SUMMARY_CANCELED;
 
 import static com.android.dx.mockito.inline.extended.ExtendedMockito.spyOn;
-import static com.android.wm.shell.bubbles.Bubble.KEY_APP_BUBBLE;
 
 import static com.google.common.truth.Truth.assertThat;
 
@@ -1734,13 +1733,13 @@
     @Test
     public void testShowOrHideAppBubble_addsAndExpand() {
         assertThat(mBubbleController.isStackExpanded()).isFalse();
-        assertThat(mBubbleData.getBubbleInStackWithKey(KEY_APP_BUBBLE)).isNull();
 
         mBubbleController.showOrHideAppBubble(mAppBubbleIntent, mUser0, mAppBubbleIcon);
 
         verify(mBubbleController).inflateAndAdd(any(Bubble.class), /* suppressFlyout= */ eq(true),
                 /* showInShade= */ eq(false));
-        assertThat(mBubbleData.getSelectedBubble().getKey()).isEqualTo(KEY_APP_BUBBLE);
+        assertThat(mBubbleData.getSelectedBubble().getKey()).isEqualTo(
+                Bubble.getAppBubbleKeyForApp(mContext.getPackageName(), mUser0));
         assertThat(mBubbleController.isStackExpanded()).isTrue();
     }
 
@@ -1754,7 +1753,8 @@
         // Calling this while collapsed will expand the app bubble
         mBubbleController.showOrHideAppBubble(mAppBubbleIntent, mUser0, mAppBubbleIcon);
 
-        assertThat(mBubbleData.getSelectedBubble().getKey()).isEqualTo(KEY_APP_BUBBLE);
+        assertThat(mBubbleData.getSelectedBubble().getKey()).isEqualTo(
+                Bubble.getAppBubbleKeyForApp(mContext.getPackageName(), mUser0));
         assertThat(mBubbleController.isStackExpanded()).isTrue();
         assertThat(mBubbleData.getBubbles().size()).isEqualTo(2);
     }
@@ -1762,13 +1762,15 @@
     @Test
     public void testShowOrHideAppBubble_collapseIfSelected() {
         mBubbleController.showOrHideAppBubble(mAppBubbleIntent, mUser0, mAppBubbleIcon);
-        assertThat(mBubbleData.getSelectedBubble().getKey()).isEqualTo(KEY_APP_BUBBLE);
+        assertThat(mBubbleData.getSelectedBubble().getKey()).isEqualTo(
+                Bubble.getAppBubbleKeyForApp(mContext.getPackageName(), mUser0));
         assertThat(mBubbleController.isStackExpanded()).isTrue();
 
         // Calling this while the app bubble is expanded should collapse the stack
         mBubbleController.showOrHideAppBubble(mAppBubbleIntent, mUser0, mAppBubbleIcon);
 
-        assertThat(mBubbleData.getSelectedBubble().getKey()).isEqualTo(KEY_APP_BUBBLE);
+        assertThat(mBubbleData.getSelectedBubble().getKey()).isEqualTo(
+                Bubble.getAppBubbleKeyForApp(mContext.getPackageName(), mUser0));
         assertThat(mBubbleController.isStackExpanded()).isFalse();
         assertThat(mBubbleData.getBubbles().size()).isEqualTo(1);
         assertThat(mBubbleData.getBubbles().get(0).getUser()).isEqualTo(mUser0);
@@ -1777,8 +1779,9 @@
     @Test
     public void testShowOrHideAppBubbleWithNonPrimaryUser_bubbleCollapsedWithExpectedUser() {
         UserHandle user10 = createUserHandle(/* userId = */ 10);
+        String appBubbleKey = Bubble.getAppBubbleKeyForApp(mContext.getPackageName(), user10);
         mBubbleController.showOrHideAppBubble(mAppBubbleIntent, user10, mAppBubbleIcon);
-        assertThat(mBubbleData.getSelectedBubble().getKey()).isEqualTo(KEY_APP_BUBBLE);
+        assertThat(mBubbleData.getSelectedBubble().getKey()).isEqualTo(appBubbleKey);
         assertThat(mBubbleController.isStackExpanded()).isTrue();
         assertThat(mBubbleData.getBubbles().size()).isEqualTo(1);
         assertThat(mBubbleData.getBubbles().get(0).getUser()).isEqualTo(user10);
@@ -1786,13 +1789,28 @@
         // Calling this while the app bubble is expanded should collapse the stack
         mBubbleController.showOrHideAppBubble(mAppBubbleIntent, user10, mAppBubbleIcon);
 
-        assertThat(mBubbleData.getSelectedBubble().getKey()).isEqualTo(KEY_APP_BUBBLE);
+        assertThat(mBubbleData.getSelectedBubble().getKey()).isEqualTo(appBubbleKey);
         assertThat(mBubbleController.isStackExpanded()).isFalse();
         assertThat(mBubbleData.getBubbles().size()).isEqualTo(1);
         assertThat(mBubbleData.getBubbles().get(0).getUser()).isEqualTo(user10);
     }
 
     @Test
+    public void testShowOrHideAppBubbleOnUser10AndThenUser0_user0BubbleExpanded() {
+        UserHandle user10 = createUserHandle(/* userId = */ 10);
+        mBubbleController.showOrHideAppBubble(mAppBubbleIntent, user10, mAppBubbleIcon);
+
+        String appBubbleUser0Key = Bubble.getAppBubbleKeyForApp(mContext.getPackageName(), mUser0);
+        mBubbleController.showOrHideAppBubble(mAppBubbleIntent, mUser0, mAppBubbleIcon);
+
+        assertThat(mBubbleData.getSelectedBubble().getKey()).isEqualTo(appBubbleUser0Key);
+        assertThat(mBubbleController.isStackExpanded()).isTrue();
+        assertThat(mBubbleData.getBubbles()).hasSize(2);
+        assertThat(mBubbleData.getBubbles().get(0).getUser()).isEqualTo(mUser0);
+        assertThat(mBubbleData.getBubbles().get(1).getUser()).isEqualTo(user10);
+    }
+
+    @Test
     public void testShowOrHideAppBubble_selectIfNotSelected() {
         mBubbleController.showOrHideAppBubble(mAppBubbleIntent, mUser0, mAppBubbleIcon);
         mBubbleController.updateBubble(mBubbleEntry);
@@ -1801,7 +1819,8 @@
         assertThat(mBubbleController.isStackExpanded()).isTrue();
 
         mBubbleController.showOrHideAppBubble(mAppBubbleIntent, mUser0, mAppBubbleIcon);
-        assertThat(mBubbleData.getSelectedBubble().getKey()).isEqualTo(KEY_APP_BUBBLE);
+        assertThat(mBubbleData.getSelectedBubble().getKey()).isEqualTo(
+                Bubble.getAppBubbleKeyForApp(mContext.getPackageName(), mUser0));
         assertThat(mBubbleController.isStackExpanded()).isTrue();
         assertThat(mBubbleData.getBubbles().size()).isEqualTo(2);
     }