Merge changes I610545fd,I3c0d2677 into main

* changes:
  Update NoMan's ShortcutHelper to use ShortcutChangeCallback
  Add way to remove shortcut changed callbacks to shortcut service

6 files changed, 240 insertions, 194 deletions

diff --git a/core/java/android/content/pm/ShortcutServiceInternal.java b/core/java/android/content/pm/ShortcutServiceInternal.java
index 55d0bc1deedc..c811a47e5b05 100644
--- a/core/java/android/content/pm/ShortcutServiceInternal.java
+++ b/core/java/android/content/pm/ShortcutServiceInternal.java
@@ -91,6 +91,9 @@ public abstract class ShortcutServiceInternal {
     public abstract void addShortcutChangeCallback(
             @NonNull LauncherApps.ShortcutChangeCallback callback);
 
+    public abstract void removeShortcutChangeCallback(
+            @NonNull LauncherApps.ShortcutChangeCallback callback);
+
     public abstract int getShortcutIconResId(int launcherUserId, @NonNull String callingPackage,
             @NonNull String packageName, @NonNull String shortcutId, int userId);
 
diff --git a/services/core/java/com/android/server/notification/NotificationManagerService.java b/services/core/java/com/android/server/notification/NotificationManagerService.java
index 1c40f44b7b78..9e53cc357ea4 100644
--- a/services/core/java/com/android/server/notification/NotificationManagerService.java
+++ b/services/core/java/com/android/server/notification/NotificationManagerService.java
@@ -8658,8 +8658,7 @@ public class NotificationManagerService extends SystemService {
                     mAttentionHelper.updateLightsLocked();
                     if (mShortcutHelper != null) {
                         mShortcutHelper.maybeListenForShortcutChangesForBubbles(r,
-                                true /* isRemoved */,
-                                mHandler);
+                                true /* isRemoved */);
                     }
                 } else {
                     if (notificationForceGrouping()) {
@@ -9116,8 +9115,7 @@ public class NotificationManagerService extends SystemService {
 
                     if (mShortcutHelper != null) {
                         mShortcutHelper.maybeListenForShortcutChangesForBubbles(r,
-                                false /* isRemoved */,
-                                mHandler);
+                                false /* isRemoved */);
                     }
 
                     maybeRecordInterruptionLocked(r);
diff --git a/services/core/java/com/android/server/notification/ShortcutHelper.java b/services/core/java/com/android/server/notification/ShortcutHelper.java
index 86dcecf9290a..857e3198d55b 100644
--- a/services/core/java/com/android/server/notification/ShortcutHelper.java
+++ b/services/core/java/com/android/server/notification/ShortcutHelper.java
@@ -27,7 +27,6 @@ import android.content.pm.LauncherApps;
 import android.content.pm.ShortcutInfo;
 import android.content.pm.ShortcutServiceInternal;
 import android.os.Binder;
-import android.os.Handler;
 import android.os.UserHandle;
 import android.os.UserManager;
 import android.text.TextUtils;
@@ -65,85 +64,34 @@ public class ShortcutHelper {
         void onShortcutRemoved(String key);
     }
 
+    private final ShortcutListener mShortcutListener;
     private LauncherApps mLauncherAppsService;
-    private ShortcutListener mShortcutListener;
     private ShortcutServiceInternal mShortcutServiceInternal;
     private UserManager mUserManager;
 
-    // Key: packageName Value: <shortcutId, notifId>
-    private HashMap<String, HashMap<String, String>> mActiveShortcutBubbles = new HashMap<>();
-    private boolean mLauncherAppsCallbackRegistered;
+    // Key: packageName|userId Value: <shortcutId, notifId>
+    private final HashMap<String, HashMap<String, String>> mActiveShortcutBubbles = new HashMap<>();
+    private boolean mShortcutChangedCallbackRegistered;
 
     // Bubbles can be created based on a shortcut, we need to listen for changes to
     // that shortcut so that we may update the bubble appropriately.
-    private final LauncherApps.Callback mLauncherAppsCallback = new LauncherApps.Callback() {
-        @Override
-        public void onPackageRemoved(String packageName, UserHandle user) {
-        }
-
-        @Override
-        public void onPackageAdded(String packageName, UserHandle user) {
-        }
-
-        @Override
-        public void onPackageChanged(String packageName, UserHandle user) {
-        }
-
-        @Override
-        public void onPackagesAvailable(String[] packageNames, UserHandle user,
-                boolean replacing) {
-        }
-
-        @Override
-        public void onPackagesUnavailable(String[] packageNames, UserHandle user,
-                boolean replacing) {
-        }
+    private final LauncherApps.ShortcutChangeCallback mShortcutChangeCallback =
+            new LauncherApps.ShortcutChangeCallback() {
 
-        @Override
-        public void onShortcutsChanged(@NonNull String packageName,
-                @NonNull List<ShortcutInfo> shortcuts, @NonNull UserHandle user) {
-            HashMap<String, String> shortcutBubbles = mActiveShortcutBubbles.get(packageName);
-            ArrayList<String> bubbleKeysToRemove = new ArrayList<>();
-            if (shortcutBubbles != null) {
-                // Copy to avoid a concurrent modification exception when we remove bubbles from
-                // shortcutBubbles.
-                final Set<String> shortcutIds = new HashSet<>(shortcutBubbles.keySet());
-
-                // If we can't find one of our bubbles in the shortcut list, that bubble needs
-                // to be removed.
-                for (String shortcutId : shortcutIds) {
-                    boolean foundShortcut = false;
-                    for (int i = 0; i < shortcuts.size(); i++) {
-                        if (shortcuts.get(i).getId().equals(shortcutId)) {
-                            foundShortcut = true;
-                            break;
-                        }
-                    }
-                    if (!foundShortcut) {
-                        bubbleKeysToRemove.add(shortcutBubbles.get(shortcutId));
-                        shortcutBubbles.remove(shortcutId);
-                        if (shortcutBubbles.isEmpty()) {
-                            mActiveShortcutBubbles.remove(packageName);
-                            if (mLauncherAppsCallbackRegistered
-                                    && mActiveShortcutBubbles.isEmpty()) {
-                                mLauncherAppsService.unregisterCallback(mLauncherAppsCallback);
-                                mLauncherAppsCallbackRegistered = false;
-                            }
-                        }
-                    }
+                @Override
+                public void onShortcutsAddedOrUpdated(@NonNull String packageName,
+                        @NonNull List<ShortcutInfo> shortcuts, @NonNull UserHandle user) {
                 }
-            }
 
-            // Let NoMan know about the updates
-            for (int i = 0; i < bubbleKeysToRemove.size(); i++) {
-                // update flag bubble
-                String bubbleKey = bubbleKeysToRemove.get(i);
-                if (mShortcutListener != null) {
-                    mShortcutListener.onShortcutRemoved(bubbleKey);
+                public void onShortcutsRemoved(@NonNull String packageName,
+                        @NonNull List<ShortcutInfo> removedShortcuts, @NonNull UserHandle user) {
+                    final String packageUserKey = getPackageUserKey(packageName, user);
+                    if (mActiveShortcutBubbles.get(packageUserKey) == null) return;
+                    for (ShortcutInfo info : removedShortcuts) {
+                        onShortcutRemoved(packageUserKey, info.getId());
+                    }
                 }
-            }
-        }
-    };
+            };
 
     ShortcutHelper(LauncherApps launcherApps, ShortcutListener listener,
             ShortcutServiceInternal shortcutServiceInternal, UserManager userManager) {
@@ -172,14 +120,14 @@ public class ShortcutHelper {
      * Returns whether the given shortcut info is a conversation shortcut.
      */
     public static boolean isConversationShortcut(
-            ShortcutInfo shortcutInfo, ShortcutServiceInternal mShortcutServiceInternal,
+            ShortcutInfo shortcutInfo, ShortcutServiceInternal shortcutServiceInternal,
             int callingUserId) {
         if (shortcutInfo == null || !shortcutInfo.isLongLived() || !shortcutInfo.isEnabled()) {
             return false;
         }
         // TODO (b/155016294) uncomment when sharing shortcuts are required
         /*
-        mShortcutServiceInternal.isSharingShortcut(callingUserId, "android",
+        shortcutServiceInternal.isSharingShortcut(callingUserId, "android",
                 shortcutInfo.getPackage(), shortcutInfo.getId(), shortcutInfo.getUserId(),
                 SHARING_FILTER);
          */
@@ -233,34 +181,30 @@ public class ShortcutHelper {
      *
      * @param r the notification record to check
      * @param removedNotification true if this notification is being removed
-     * @param handler handler to register the callback with
      */
     void maybeListenForShortcutChangesForBubbles(NotificationRecord r,
-            boolean removedNotification,
-            Handler handler) {
+            boolean removedNotification) {
         final String shortcutId = r.getNotification().getBubbleMetadata() != null
                 ? r.getNotification().getBubbleMetadata().getShortcutId()
                 : null;
+        final String packageUserKey = getPackageUserKey(r.getSbn().getPackageName(), r.getUser());
         if (!removedNotification
                 && !TextUtils.isEmpty(shortcutId)
                 && r.getShortcutInfo() != null
                 && r.getShortcutInfo().getId().equals(shortcutId)) {
             // Must track shortcut based bubbles in case the shortcut is removed
             HashMap<String, String> packageBubbles = mActiveShortcutBubbles.get(
-                    r.getSbn().getPackageName());
+                    packageUserKey);
             if (packageBubbles == null) {
                 packageBubbles = new HashMap<>();
             }
             packageBubbles.put(shortcutId, r.getKey());
-            mActiveShortcutBubbles.put(r.getSbn().getPackageName(), packageBubbles);
-            if (!mLauncherAppsCallbackRegistered) {
-                mLauncherAppsService.registerCallback(mLauncherAppsCallback, handler);
-                mLauncherAppsCallbackRegistered = true;
-            }
+            mActiveShortcutBubbles.put(packageUserKey, packageBubbles);
+            registerCallbackIfNeeded();
         } else {
             // No longer track shortcut
             HashMap<String, String> packageBubbles = mActiveShortcutBubbles.get(
-                    r.getSbn().getPackageName());
+                    packageUserKey);
             if (packageBubbles != null) {
                 if (!TextUtils.isEmpty(shortcutId)) {
                     packageBubbles.remove(shortcutId);
@@ -278,20 +222,62 @@ public class ShortcutHelper {
                     }
                 }
                 if (packageBubbles.isEmpty()) {
-                    mActiveShortcutBubbles.remove(r.getSbn().getPackageName());
+                    mActiveShortcutBubbles.remove(packageUserKey);
                 }
             }
-            if (mLauncherAppsCallbackRegistered && mActiveShortcutBubbles.isEmpty()) {
-                mLauncherAppsService.unregisterCallback(mLauncherAppsCallback);
-                mLauncherAppsCallbackRegistered = false;
+            unregisterCallbackIfNeeded();
+        }
+    }
+
+    private String getPackageUserKey(String packageName, UserHandle user) {
+        return packageName + "|" + user.getIdentifier();
+    }
+
+    private void onShortcutRemoved(String packageUserKey, String shortcutId) {
+        HashMap<String, String> shortcutBubbles = mActiveShortcutBubbles.get(packageUserKey);
+        ArrayList<String> bubbleKeysToRemove = new ArrayList<>();
+        if (shortcutBubbles != null) {
+            if (shortcutBubbles.containsKey(shortcutId)) {
+                bubbleKeysToRemove.add(shortcutBubbles.get(shortcutId));
+                shortcutBubbles.remove(shortcutId);
+                if (shortcutBubbles.isEmpty()) {
+                    mActiveShortcutBubbles.remove(packageUserKey);
+                    unregisterCallbackIfNeeded();
+                }
             }
+            notifyNoMan(bubbleKeysToRemove);
+        }
+    }
+
+    private void registerCallbackIfNeeded() {
+        if (!mShortcutChangedCallbackRegistered) {
+            mShortcutChangedCallbackRegistered = true;
+            mShortcutServiceInternal.addShortcutChangeCallback(mShortcutChangeCallback);
+        }
+    }
+
+    private void unregisterCallbackIfNeeded() {
+        if (mShortcutChangedCallbackRegistered && mActiveShortcutBubbles.isEmpty()) {
+            mShortcutServiceInternal.removeShortcutChangeCallback(mShortcutChangeCallback);
+            mShortcutChangedCallbackRegistered = false;
         }
     }
 
     void destroy() {
-        if (mLauncherAppsCallbackRegistered) {
-            mLauncherAppsService.unregisterCallback(mLauncherAppsCallback);
-            mLauncherAppsCallbackRegistered = false;
+        if (mShortcutChangedCallbackRegistered) {
+            mShortcutServiceInternal.removeShortcutChangeCallback(mShortcutChangeCallback);
+            mShortcutChangedCallbackRegistered = false;
+        }
+    }
+
+    private void notifyNoMan(List<String> bubbleKeysToRemove) {
+        // Let NoMan know about the updates
+        for (int i = 0; i < bubbleKeysToRemove.size(); i++) {
+            // update flag bubble
+            String bubbleKey = bubbleKeysToRemove.get(i);
+            if (mShortcutListener != null) {
+                mShortcutListener.onShortcutRemoved(bubbleKey);
+            }
         }
     }
 }
diff --git a/services/core/java/com/android/server/pm/ShortcutService.java b/services/core/java/com/android/server/pm/ShortcutService.java
index 1cd77ffcedaa..4a60e452919f 100644
--- a/services/core/java/com/android/server/pm/ShortcutService.java
+++ b/services/core/java/com/android/server/pm/ShortcutService.java
@@ -3443,6 +3443,14 @@ public class ShortcutService extends IShortcutService.Stub {
         }
 
         @Override
+        public void removeShortcutChangeCallback(
+                @NonNull LauncherApps.ShortcutChangeCallback callback) {
+            synchronized (mServiceLock) {
+                mShortcutChangeCallbacks.remove(Objects.requireNonNull(callback));
+            }
+        }
+
+        @Override
         public int getShortcutIconResId(int launcherUserId, @NonNull String callingPackage,
                 @NonNull String packageName, @NonNull String shortcutId, int userId) {
             Objects.requireNonNull(callingPackage, "callingPackage");
diff --git a/services/tests/uiservicestests/src/com/android/server/notification/NotificationManagerServiceTest.java b/services/tests/uiservicestests/src/com/android/server/notification/NotificationManagerServiceTest.java
index 5d306e152ad7..c1f5a01a8c47 100644
--- a/services/tests/uiservicestests/src/com/android/server/notification/NotificationManagerServiceTest.java
+++ b/services/tests/uiservicestests/src/com/android/server/notification/NotificationManagerServiceTest.java
@@ -89,7 +89,6 @@ import static android.os.UserManager.USER_TYPE_FULL_SECONDARY;
 import static android.os.UserManager.USER_TYPE_PROFILE_CLONE;
 import static android.os.UserManager.USER_TYPE_PROFILE_MANAGED;
 import static android.os.UserManager.USER_TYPE_PROFILE_PRIVATE;
-import static android.platform.test.flag.junit.SetFlagsRule.DefaultInitValueType.DEVICE_DEFAULT;
 import static android.provider.Settings.Global.ZEN_MODE_IMPORTANT_INTERRUPTIONS;
 import static android.service.notification.Adjustment.KEY_CONTEXTUAL_ACTIONS;
 import static android.service.notification.Adjustment.KEY_IMPORTANCE;
@@ -838,13 +837,7 @@ public class NotificationManagerServiceTest extends UiServiceTestCase {
 
         // Pretend the shortcut exists
         List<ShortcutInfo> shortcutInfos = new ArrayList<>();
-        ShortcutInfo info = mock(ShortcutInfo.class);
-        when(info.getPackage()).thenReturn(mPkg);
-        when(info.getId()).thenReturn(VALID_CONVO_SHORTCUT_ID);
-        when(info.getUserId()).thenReturn(USER_SYSTEM);
-        when(info.isLongLived()).thenReturn(true);
-        when(info.isEnabled()).thenReturn(true);
-        shortcutInfos.add(info);
+        shortcutInfos.add(createMockConvoShortcut());
         when(mLauncherApps.getShortcuts(any(), any())).thenReturn(shortcutInfos);
         when(mShortcutServiceInternal.isSharingShortcut(anyInt(), anyString(), anyString(),
                 anyString(), anyInt(), any())).thenReturn(true);
@@ -11109,8 +11102,8 @@ public class NotificationManagerServiceTest extends UiServiceTestCase {
                 BUBBLE_PREFERENCE_ALL /* app */,
                 true /* channel */);
 
-        ArgumentCaptor<LauncherApps.Callback> launcherAppsCallback =
-                ArgumentCaptor.forClass(LauncherApps.Callback.class);
+        ArgumentCaptor<LauncherApps.ShortcutChangeCallback> shortcutChangeCallback =
+                ArgumentCaptor.forClass(LauncherApps.ShortcutChangeCallback.class);
 
         // Messaging notification with shortcut info
         Notification.BubbleMetadata metadata =
@@ -11131,7 +11124,8 @@ public class NotificationManagerServiceTest extends UiServiceTestCase {
         // Verify:
 
         // Make sure we register the callback for shortcut changes
-        verify(mLauncherApps, times(1)).registerCallback(launcherAppsCallback.capture(), any());
+        verify(mShortcutServiceInternal, times(1)).addShortcutChangeCallback(
+                shortcutChangeCallback.capture());
 
         // yes allowed, yes messaging w/shortcut, yes bubble
         Notification notif = mService.getNotificationRecord(nr.getSbn().getKey()).getNotification();
@@ -11144,14 +11138,17 @@ public class NotificationManagerServiceTest extends UiServiceTestCase {
 
         // Test: Remove the shortcut
         when(mLauncherApps.getShortcuts(any(), any())).thenReturn(null);
-        launcherAppsCallback.getValue().onShortcutsChanged(mPkg, emptyList(),
+        ArrayList<ShortcutInfo> removedShortcuts = new ArrayList<>();
+        removedShortcuts.add(createMockConvoShortcut());
+        shortcutChangeCallback.getValue().onShortcutsRemoved(mPkg, removedShortcuts,
                 UserHandle.getUserHandleForUid(mUid));
         waitForIdle();
 
         // Verify:
 
         // Make sure callback is unregistered
-        verify(mLauncherApps, times(1)).unregisterCallback(launcherAppsCallback.getValue());
+        verify(mShortcutServiceInternal, times(1)).removeShortcutChangeCallback(
+                shortcutChangeCallback.getValue());
 
         // We're no longer a bubble
         NotificationRecord notif2 = mService.getNotificationRecord(
@@ -11169,8 +11166,8 @@ public class NotificationManagerServiceTest extends UiServiceTestCase {
                 BUBBLE_PREFERENCE_ALL /* app */,
                 true /* channel */);
 
-        ArgumentCaptor<LauncherApps.Callback> launcherAppsCallback =
-                ArgumentCaptor.forClass(LauncherApps.Callback.class);
+        ArgumentCaptor<LauncherApps.ShortcutChangeCallback> shortcutChangeCallback =
+                ArgumentCaptor.forClass(LauncherApps.ShortcutChangeCallback.class);
 
         // Messaging notification with shortcut info
         Notification.BubbleMetadata metadata = new Notification.BubbleMetadata.Builder(
@@ -11204,7 +11201,8 @@ public class NotificationManagerServiceTest extends UiServiceTestCase {
         // Verify:
 
         // Make sure we register the callback for shortcut changes
-        verify(mLauncherApps, times(1)).registerCallback(launcherAppsCallback.capture(), any());
+        verify(mShortcutServiceInternal, times(1)).addShortcutChangeCallback(
+                shortcutChangeCallback.capture());
 
         // yes allowed, yes messaging w/shortcut, yes bubble
         Notification notif = mService.getNotificationRecord(nr.getSbn().getKey()).getNotification();
@@ -11223,7 +11221,8 @@ public class NotificationManagerServiceTest extends UiServiceTestCase {
         // Verify:
 
         // Make sure callback is unregistered
-        verify(mLauncherApps, times(1)).unregisterCallback(launcherAppsCallback.getValue());
+        verify(mShortcutServiceInternal, times(1)).removeShortcutChangeCallback(
+                shortcutChangeCallback.getValue());
     }
 
     @Test
@@ -16263,4 +16262,14 @@ public class NotificationManagerServiceTest extends UiServiceTestCase {
 
         assertThat(r.getChannel().getId()).isEqualTo(NEWS_ID);
     }
+
+    private ShortcutInfo createMockConvoShortcut() {
+        ShortcutInfo info = mock(ShortcutInfo.class);
+        when(info.getPackage()).thenReturn(mPkg);
+        when(info.getId()).thenReturn(VALID_CONVO_SHORTCUT_ID);
+        when(info.getUserId()).thenReturn(USER_SYSTEM);
+        when(info.isLongLived()).thenReturn(true);
+        when(info.isEnabled()).thenReturn(true);
+        return info;
+    }
 }
diff --git a/services/tests/uiservicestests/src/com/android/server/notification/ShortcutHelperTest.java b/services/tests/uiservicestests/src/com/android/server/notification/ShortcutHelperTest.java
index a4fb16dc1adc..f008cb671e78 100644
--- a/services/tests/uiservicestests/src/com/android/server/notification/ShortcutHelperTest.java
+++ b/services/tests/uiservicestests/src/com/android/server/notification/ShortcutHelperTest.java
@@ -22,17 +22,22 @@ import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyInt;
 import static org.mockito.ArgumentMatchers.anyString;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.reset;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
 import android.app.Notification;
+import android.app.NotificationChannel;
+import android.app.PendingIntent;
 import android.app.Person;
 import android.content.pm.LauncherApps;
 import android.content.pm.LauncherApps.ShortcutQuery;
 import android.content.pm.ShortcutInfo;
 import android.content.pm.ShortcutQueryWrapper;
 import android.content.pm.ShortcutServiceInternal;
+import android.graphics.drawable.Icon;
 import android.os.UserHandle;
 import android.os.UserManager;
 import android.service.notification.StatusBarNotification;
@@ -50,11 +55,9 @@ import org.junit.runner.RunWith;
 import org.mockito.ArgumentCaptor;
 import org.mockito.Captor;
 import org.mockito.Mock;
-import org.mockito.Mockito;
 import org.mockito.MockitoAnnotations;
 
 import java.util.ArrayList;
-import java.util.Collections;
 import java.util.List;
 
 @SmallTest
@@ -64,7 +67,6 @@ public class ShortcutHelperTest extends UiServiceTestCase {
 
     private static final String SHORTCUT_ID = "shortcut";
     private static final String PKG = "pkg";
-    private static final String KEY = "key";
     private static final Person PERSON = mock(Person.class);
 
     @Mock
@@ -75,19 +77,11 @@ public class ShortcutHelperTest extends UiServiceTestCase {
     UserManager mUserManager;
     @Mock
     ShortcutServiceInternal mShortcutServiceInternal;
-    @Mock
-    NotificationRecord mNr;
-    @Mock
-    Notification mNotif;
-    @Mock
-    StatusBarNotification mSbn;
-    @Mock
-    Notification.BubbleMetadata mBubbleMetadata;
-    @Mock
-    ShortcutInfo mShortcutInfo;
 
     @Captor private ArgumentCaptor<ShortcutQuery> mShortcutQueryCaptor;
 
+    NotificationRecord mNr;
+
     ShortcutHelper mShortcutHelper;
 
     @Before
@@ -96,137 +90,186 @@ public class ShortcutHelperTest extends UiServiceTestCase {
 
         mShortcutHelper = new ShortcutHelper(
                 mLauncherApps, mShortcutListener, mShortcutServiceInternal, mUserManager);
-        when(mSbn.getPackageName()).thenReturn(PKG);
-        when(mShortcutInfo.getId()).thenReturn(SHORTCUT_ID);
-        when(mNotif.getBubbleMetadata()).thenReturn(mBubbleMetadata);
-        when(mBubbleMetadata.getShortcutId()).thenReturn(SHORTCUT_ID);
         when(mUserManager.isUserUnlocked(any(UserHandle.class))).thenReturn(true);
 
-        setUpMockNotificationRecord(mNr, KEY);
+        mNr = setUpNotificationRecord(SHORTCUT_ID, PKG, UserHandle.of(UserHandle.USER_SYSTEM));
     }
 
-    private void setUpMockNotificationRecord(NotificationRecord mockRecord, String key) {
-        when(mockRecord.getKey()).thenReturn(key);
-        when(mockRecord.getSbn()).thenReturn(mSbn);
-        when(mockRecord.getNotification()).thenReturn(mNotif);
-        when(mockRecord.getShortcutInfo()).thenReturn(mShortcutInfo);
+    private NotificationRecord setUpNotificationRecord(String shortcutId,
+            String pkg,
+            UserHandle user) {
+        ShortcutInfo shortcutInfo = mock(ShortcutInfo.class);
+        when(shortcutInfo.getId()).thenReturn(shortcutId);
+        when(shortcutInfo.getUserHandle()).thenReturn(user);
+        when(shortcutInfo.isLongLived()).thenReturn(true);
+
+        Notification notification = new Notification.Builder(getContext())
+                .setContentTitle("title")
+                .setShortcutId(shortcutId)
+                .setBubbleMetadata(new Notification.BubbleMetadata.Builder(shortcutId).build())
+                .build();
+
+        StatusBarNotification sbn = new StatusBarNotification(pkg, pkg, 0, null,
+                1000, 2000, notification, user, null, System.currentTimeMillis());
+        NotificationRecord record = new NotificationRecord(mContext, sbn,
+                mock(NotificationChannel.class));
+        record.setShortcutInfo(shortcutInfo);
+        return record;
     }
 
-    private LauncherApps.Callback addShortcutBubbleAndVerifyListener() {
-        mShortcutHelper.maybeListenForShortcutChangesForBubbles(mNr,
-                false /* removed */,
-                null /* handler */);
+    private LauncherApps.ShortcutChangeCallback addShortcutBubbleAndVerifyListener(
+            NotificationRecord record) {
+        mShortcutHelper.maybeListenForShortcutChangesForBubbles(record, false /* removed */);
 
-        ArgumentCaptor<LauncherApps.Callback> launcherAppsCallback =
-                ArgumentCaptor.forClass(LauncherApps.Callback.class);
+        ArgumentCaptor<LauncherApps.ShortcutChangeCallback> launcherAppsCallback =
+                ArgumentCaptor.forClass(LauncherApps.ShortcutChangeCallback.class);
 
-        verify(mLauncherApps, times(1)).registerCallback(
-                launcherAppsCallback.capture(), any());
+        verify(mShortcutServiceInternal, times(1)).addShortcutChangeCallback(
+                launcherAppsCallback.capture());
         return launcherAppsCallback.getValue();
     }
 
     @Test
     public void testBubbleAdded_listenedAdded() {
-        addShortcutBubbleAndVerifyListener();
+        addShortcutBubbleAndVerifyListener(mNr);
     }
 
     @Test
+    public void testListenerNotifiedOnShortcutRemoved() {
+        LauncherApps.ShortcutChangeCallback callback = addShortcutBubbleAndVerifyListener(mNr);
+
+        List<ShortcutInfo> removedShortcuts = new ArrayList<>();
+        removedShortcuts.add(mNr.getShortcutInfo());
+
+        callback.onShortcutsRemoved(PKG, removedShortcuts, mNr.getUser());
+        verify(mShortcutListener).onShortcutRemoved(mNr.getKey());
+    }
+
+    @Test
+    public void testListenerNotNotified_notMatchingPackage() {
+        LauncherApps.ShortcutChangeCallback callback = addShortcutBubbleAndVerifyListener(mNr);
+
+        List<ShortcutInfo> removedShortcuts = new ArrayList<>();
+        removedShortcuts.add(mNr.getShortcutInfo());
+
+        callback.onShortcutsRemoved("differentPackage", removedShortcuts, mNr.getUser());
+        verify(mShortcutListener, never()).onShortcutRemoved(anyString());
+    }
+
+    @Test
+    public void testListenerNotNotified_notMatchingUser() {
+        LauncherApps.ShortcutChangeCallback callback = addShortcutBubbleAndVerifyListener(mNr);
+
+        List<ShortcutInfo> removedShortcuts = new ArrayList<>();
+        removedShortcuts.add(mNr.getShortcutInfo());
+
+        callback.onShortcutsRemoved(PKG, removedShortcuts, UserHandle.of(10));
+        verify(mShortcutListener, never()).onShortcutRemoved(anyString());
+    }
+
+    @Test
+    public void testListenerNotifiedDifferentUser() {
+        LauncherApps.ShortcutChangeCallback callback = addShortcutBubbleAndVerifyListener(mNr);
+        NotificationRecord diffUserRecord = setUpNotificationRecord(SHORTCUT_ID, PKG,
+                UserHandle.of(10));
+        mShortcutHelper.maybeListenForShortcutChangesForBubbles(diffUserRecord,
+                false /* removed */);
+
+        List<ShortcutInfo> removedShortcuts = new ArrayList<>();
+        removedShortcuts.add(mNr.getShortcutInfo());
+
+        callback.onShortcutsRemoved(PKG, removedShortcuts, mNr.getUser());
+        verify(mShortcutListener).onShortcutRemoved(mNr.getKey());
+
+        reset(mShortcutListener);
+        removedShortcuts.clear();
+        removedShortcuts.add(diffUserRecord.getShortcutInfo());
+
+        callback.onShortcutsRemoved(PKG, removedShortcuts, diffUserRecord.getUser());
+        verify(mShortcutListener).onShortcutRemoved(diffUserRecord.getKey());
+    }
+
+
+    @Test
     public void testBubbleRemoved_listenerRemoved() {
         // First set it up to listen
-        addShortcutBubbleAndVerifyListener();
+        addShortcutBubbleAndVerifyListener(mNr);
 
         // Then remove the notif
         mShortcutHelper.maybeListenForShortcutChangesForBubbles(mNr,
-                true /* removed */,
-                null /* handler */);
+                true /* removed */);
 
-        verify(mLauncherApps, times(1)).unregisterCallback(any());
+        verify(mShortcutServiceInternal, times(1)).removeShortcutChangeCallback(any());
     }
 
     @Test
     public void testBubbleNoLongerHasBubbleMetadata_listenerRemoved() {
         // First set it up to listen
-        addShortcutBubbleAndVerifyListener();
+        addShortcutBubbleAndVerifyListener(mNr);
 
         // Then make it not a bubble
-        when(mNotif.getBubbleMetadata()).thenReturn(null);
+        mNr.getNotification().setBubbleMetadata(null);
         mShortcutHelper.maybeListenForShortcutChangesForBubbles(mNr,
-                false /* removed */,
-                null /* handler */);
+                false /* removed */);
 
-        verify(mLauncherApps, times(1)).unregisterCallback(any());
+        verify(mShortcutServiceInternal, times(1)).removeShortcutChangeCallback(any());
     }
 
     @Test
     public void testBubbleNoLongerHasShortcutId_listenerRemoved() {
         // First set it up to listen
-        addShortcutBubbleAndVerifyListener();
+        addShortcutBubbleAndVerifyListener(mNr);
 
         // Clear out shortcutId
-        when(mBubbleMetadata.getShortcutId()).thenReturn(null);
+        mNr.getNotification().setBubbleMetadata(new Notification.BubbleMetadata.Builder(
+                mock(PendingIntent.class), mock(Icon.class)).build());
         mShortcutHelper.maybeListenForShortcutChangesForBubbles(mNr,
-                false /* removed */,
-                null /* handler */);
+                false /* removed */);
 
-        verify(mLauncherApps, times(1)).unregisterCallback(any());
+        verify(mShortcutServiceInternal, times(1)).removeShortcutChangeCallback(any());
     }
 
     @Test
     public void testNotifNoLongerHasShortcut_listenerRemoved() {
         // First set it up to listen
-        addShortcutBubbleAndVerifyListener();
-
-        NotificationRecord validMock1 = Mockito.mock(NotificationRecord.class);
-        setUpMockNotificationRecord(validMock1, "KEY1");
-
-        NotificationRecord validMock2 = Mockito.mock(NotificationRecord.class);
-        setUpMockNotificationRecord(validMock2, "KEY2");
+        addShortcutBubbleAndVerifyListener(mNr);
 
-        NotificationRecord validMock3 = Mockito.mock(NotificationRecord.class);
-        setUpMockNotificationRecord(validMock3, "KEY3");
+        NotificationRecord record1 = setUpNotificationRecord(SHORTCUT_ID, PKG,
+                UserHandle.of(UserHandle.USER_SYSTEM));
+        NotificationRecord record2 = setUpNotificationRecord(SHORTCUT_ID, PKG,
+                UserHandle.of(UserHandle.USER_SYSTEM));
+        NotificationRecord record3 = setUpNotificationRecord(SHORTCUT_ID, PKG,
+                UserHandle.of(UserHandle.USER_SYSTEM));
 
-        mShortcutHelper.maybeListenForShortcutChangesForBubbles(validMock1,
-                false /* removed */,
-                null /* handler */);
+        mShortcutHelper.maybeListenForShortcutChangesForBubbles(record1,
+                false /* removed */);
 
-        mShortcutHelper.maybeListenForShortcutChangesForBubbles(validMock2,
-                false /* removed */,
-                null /* handler */);
+        mShortcutHelper.maybeListenForShortcutChangesForBubbles(record2,
+                false /* removed */);
 
-        mShortcutHelper.maybeListenForShortcutChangesForBubbles(validMock3,
-                false /* removed */,
-                null /* handler */);
+        mShortcutHelper.maybeListenForShortcutChangesForBubbles(record3,
+                false /* removed */);
 
-        // Clear out shortcutId of the bubble in the middle, to double check that we don't hit a
+        // Clear out shortcutId of the bubble in the middle, to double-check that we don't hit a
         // concurrent modification exception (removing the last bubble would sidestep that check).
-        when(validMock2.getShortcutInfo()).thenReturn(null);
-        mShortcutHelper.maybeListenForShortcutChangesForBubbles(validMock2,
-                false /* removed */,
-                null /* handler */);
+        record2.setShortcutInfo(null);
+        mShortcutHelper.maybeListenForShortcutChangesForBubbles(record2,
+                false /* removed */);
 
-        verify(mLauncherApps, times(1)).unregisterCallback(any());
+        verify(mShortcutServiceInternal, times(1)).removeShortcutChangeCallback(any());
     }
 
     @Test
     public void testOnShortcutsChanged_listenerRemoved() {
         // First set it up to listen
-        LauncherApps.Callback callback = addShortcutBubbleAndVerifyListener();
+        LauncherApps.ShortcutChangeCallback callback = addShortcutBubbleAndVerifyListener(mNr);
 
         // App shortcuts are removed:
-        callback.onShortcutsChanged(PKG, Collections.emptyList(),  mock(UserHandle.class));
+        List<ShortcutInfo> removedShortcuts = new ArrayList<>();
+        removedShortcuts.add(mNr.getShortcutInfo());
+        callback.onShortcutsRemoved(PKG, removedShortcuts,  mNr.getUser());
 
-        verify(mLauncherApps, times(1)).unregisterCallback(any());
-    }
-
-    @Test
-    public void testListenerNotifiedOnShortcutRemoved() {
-        LauncherApps.Callback callback = addShortcutBubbleAndVerifyListener();
-
-        List<ShortcutInfo> shortcutInfos = new ArrayList<>();
-        when(mLauncherApps.getShortcuts(any(), any())).thenReturn(shortcutInfos);
-
-        callback.onShortcutsChanged(PKG, shortcutInfos, mock(UserHandle.class));
-        verify(mShortcutListener).onShortcutRemoved(mNr.getKey());
+        verify(mShortcutServiceInternal, times(1)).removeShortcutChangeCallback(any());
     }
 
     @Test
@@ -321,7 +364,6 @@ public class ShortcutHelperTest extends UiServiceTestCase {
                 .isSameInstanceAs(si);
     }
 
-
     @Test
     public void testGetValidShortcutInfo_isValidButUserLocked() {
         ShortcutInfo si = mock(ShortcutInfo.class);