Persists bubbles to disk (part 4)

1. Decouple Bubble from NotificationEntry
2. Convert ShortcutInfo into Bubbles
3. Load overflow bubbles into memory

Test: follow the following steps
1. add a few bubbles from test app
2. restart the device
3. add a bubble from test app
4. click the plus sign to expand overview bubbles
5. verify overflow bubbles is displayed properly

Bug: 149713060
Change-Id: I6905d93fe409f4401f7d5d2d08597c36c028b5de
diff --git a/packages/SystemUI/src/com/android/systemui/bubbles/Bubble.java b/packages/SystemUI/src/com/android/systemui/bubbles/Bubble.java
index f1b401e..f044389 100644
--- a/packages/SystemUI/src/com/android/systemui/bubbles/Bubble.java
+++ b/packages/SystemUI/src/com/android/systemui/bubbles/Bubble.java
@@ -16,6 +16,7 @@
 package com.android.systemui.bubbles;
 
 
+import static android.app.Notification.FLAG_BUBBLE;
 import static android.os.AsyncTask.Status.FINISHED;
 import static android.view.Display.INVALID_DISPLAY;
 
@@ -27,6 +28,7 @@
 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.content.pm.ShortcutInfo;
@@ -55,6 +57,11 @@
 class Bubble implements BubbleViewProvider {
     private static final String TAG = "Bubble";
 
+    /**
+     * NotificationEntry associated with the bubble. A null value implies this bubble is loaded
+     * from disk.
+     */
+    @Nullable
     private NotificationEntry mEntry;
     private final String mKey;
     private final String mGroupId;
@@ -95,18 +102,56 @@
     private Bitmap mBadgedImage;
     private int mDotColor;
     private Path mDotPath;
+    private int mFlags;
 
-
+    /**
+     * Extract GroupId from {@link NotificationEntry}. See also {@link #groupId(ShortcutInfo)}.
+     */
     public static String groupId(NotificationEntry entry) {
         UserHandle user = entry.getSbn().getUser();
         return user.getIdentifier() + "|" + entry.getSbn().getPackageName();
     }
 
-    // TODO: Decouple Bubble from NotificationEntry and transform ShortcutInfo into Bubble
+    /**
+     * Extract GroupId from {@link ShortcutInfo}. This should match the one generated from
+     * {@link NotificationEntry}. See also {@link #groupId(NotificationEntry)}.
+     */
+    @NonNull
+    public static String groupId(@NonNull final ShortcutInfo shortcutInfo) {
+        return shortcutInfo.getUserId() + "|" + shortcutInfo.getPackage();
+    }
+
+    /**
+     * Generate a unique identifier for this bubble based on given {@link NotificationEntry}. If
+     * {@link ShortcutInfo} was found in the notification entry, the identifier would be generated
+     * from {@link ShortcutInfo} instead. See also {@link #key(ShortcutInfo)}.
+     */
+    @NonNull
+    public static String key(@NonNull final NotificationEntry entry) {
+        final ShortcutInfo shortcutInfo = entry.getRanking().getShortcutInfo();
+        if (shortcutInfo != null) return key(shortcutInfo);
+        return entry.getKey();
+    }
+
+    /**
+     * Generate a unique identifier for this bubble based on given {@link ShortcutInfo}.
+     * See also {@link #key(NotificationEntry)}.
+     */
+    @NonNull
+    public static String key(@NonNull final ShortcutInfo shortcutInfo) {
+        return shortcutInfo.getUserId() + "|" + shortcutInfo.getPackage() + "|"
+                + shortcutInfo.getId();
+    }
+
+    /**
+     * 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(ShortcutInfo shortcutInfo) {
         mShortcutInfo = shortcutInfo;
-        mKey = shortcutInfo.getId();
-        mGroupId = shortcutInfo.getId();
+        mKey = key(shortcutInfo);
+        mGroupId = groupId(shortcutInfo);
+        mFlags = 0;
     }
 
     /** Used in tests when no UI is required. */
@@ -114,10 +159,11 @@
     Bubble(NotificationEntry e,
             BubbleController.NotificationSuppressionChangedListener listener) {
         mEntry = e;
-        mKey = e.getKey();
+        mKey = key(e);
         mLastUpdated = e.getSbn().getPostTime();
         mGroupId = groupId(e);
         mSuppressionListener = listener;
+        mFlags = e.getSbn().getNotification().flags;
     }
 
     @Override
@@ -125,16 +171,26 @@
         return mKey;
     }
 
+    @Nullable
     public NotificationEntry getEntry() {
         return mEntry;
     }
 
+    @Nullable
+    public UserHandle getUser() {
+        if (mEntry != null) return mEntry.getSbn().getUser();
+        if (mShortcutInfo != null) return mShortcutInfo.getUserHandle();
+        return null;
+    }
+
     public String getGroupId() {
         return mGroupId;
     }
 
     public String getPackageName() {
-        return mEntry.getSbn().getPackageName();
+        return mEntry == null
+                ? mShortcutInfo == null ? null : mShortcutInfo.getPackage()
+                : mEntry.getSbn().getPackageName();
     }
 
     @Override
@@ -218,7 +274,8 @@
     void inflate(BubbleViewInfoTask.Callback callback,
             Context context,
             BubbleStackView stackView,
-            BubbleIconFactory iconFactory) {
+            BubbleIconFactory iconFactory,
+            boolean skipInflation) {
         if (isBubbleLoading()) {
             mInflationTask.cancel(true /* mayInterruptIfRunning */);
         }
@@ -226,6 +283,7 @@
                 context,
                 stackView,
                 iconFactory,
+                skipInflation,
                 callback);
         if (mInflateSynchronously) {
             mInflationTask.onPostExecute(mInflationTask.doInBackground());
@@ -345,6 +403,7 @@
      * Whether this notification should be shown in the shade.
      */
     boolean showInShade() {
+        if (mEntry == null) return false;
         return !shouldSuppressNotification() || !mEntry.isClearable();
     }
 
@@ -352,8 +411,8 @@
      * Sets whether this notification should be suppressed in the shade.
      */
     void setSuppressNotification(boolean suppressNotification) {
+        if (mEntry == null) return;
         boolean prevShowInShade = showInShade();
-
         Notification.BubbleMetadata data = mEntry.getBubbleMetadata();
         int flags = data.getFlags();
         if (suppressNotification) {
@@ -384,6 +443,7 @@
      */
     @Override
     public boolean showDot() {
+        if (mEntry == null) return false;
         return mShowBubbleUpdateDot
                 && !mEntry.shouldSuppressNotificationDot()
                 && !shouldSuppressNotification();
@@ -393,6 +453,7 @@
      * Whether the flyout for the bubble should be shown.
      */
     boolean showFlyout() {
+        if (mEntry == null) return false;
         return !mSuppressFlyout && !mEntry.shouldSuppressPeek()
                 && !shouldSuppressNotification()
                 && !mEntry.shouldSuppressNotificationList();
@@ -416,11 +477,13 @@
      * is an ongoing bubble.
      */
     boolean isOngoing() {
+        if (mEntry == null) return false;
         int flags = mEntry.getSbn().getNotification().flags;
         return (flags & Notification.FLAG_FOREGROUND_SERVICE) != 0;
     }
 
     float getDesiredHeight(Context context) {
+        if (mEntry == null) return 0;
         Notification.BubbleMetadata data = mEntry.getBubbleMetadata();
         boolean useRes = data.getDesiredHeightResId() != 0;
         if (useRes) {
@@ -434,6 +497,7 @@
     }
 
     String getDesiredHeightString() {
+        if (mEntry == null) return String.valueOf(0);
         Notification.BubbleMetadata data = mEntry.getBubbleMetadata();
         boolean useRes = data.getDesiredHeightResId() != 0;
         if (useRes) {
@@ -450,11 +514,13 @@
      * To populate the icon use {@link LauncherApps#getShortcutIconDrawable(ShortcutInfo, int)}.
      */
     boolean usingShortcutInfo() {
-        return mEntry.getBubbleMetadata().getShortcutId() != null;
+        return mEntry != null && mEntry.getBubbleMetadata().getShortcutId() != null
+                || mShortcutInfo != null;
     }
 
     @Nullable
     PendingIntent getBubbleIntent() {
+        if (mEntry == null) return null;
         Notification.BubbleMetadata data = mEntry.getBubbleMetadata();
         if (data != null) {
             return data.getIntent();
@@ -462,16 +528,32 @@
         return null;
     }
 
-    Intent getSettingsIntent() {
+    Intent getSettingsIntent(final Context context) {
         final Intent intent = new Intent(Settings.ACTION_APP_NOTIFICATION_BUBBLE_SETTINGS);
         intent.putExtra(Settings.EXTRA_APP_PACKAGE, getPackageName());
-        intent.putExtra(Settings.EXTRA_APP_UID, mEntry.getSbn().getUid());
+        final int uid = getUid(context);
+        if (uid != -1) {
+            intent.putExtra(Settings.EXTRA_APP_UID, uid);
+        }
         intent.addFlags(Intent.FLAG_ACTIVITY_MULTIPLE_TASK);
         intent.addFlags(Intent.FLAG_ACTIVITY_NEW_TASK);
         intent.addFlags(Intent.FLAG_ACTIVITY_SINGLE_TOP);
         return intent;
     }
 
+    private int getUid(final Context context) {
+        if (mEntry != null) return mEntry.getSbn().getUid();
+        final PackageManager pm = context.getPackageManager();
+        if (pm == null) return -1;
+        try {
+            final ApplicationInfo info = pm.getApplicationInfo(mShortcutInfo.getPackage(), 0);
+            return info.uid;
+        } catch (PackageManager.NameNotFoundException e) {
+            Log.e(TAG, "cannot find uid", e);
+        }
+        return -1;
+    }
+
     private int getDimenForPackageUser(Context context, int resId, String pkg, int userId) {
         PackageManager pm = context.getPackageManager();
         Resources r;
@@ -493,15 +575,30 @@
     }
 
     private boolean shouldSuppressNotification() {
+        if (mEntry == null) return false;
         return mEntry.getBubbleMetadata() != null
                 && mEntry.getBubbleMetadata().isNotificationSuppressed();
     }
 
     boolean shouldAutoExpand() {
+        if (mEntry == null) return false;
         Notification.BubbleMetadata metadata = mEntry.getBubbleMetadata();
         return metadata != null && metadata.getAutoExpandBubble();
     }
 
+    public boolean isBubble() {
+        if (mEntry == null) return (mFlags & FLAG_BUBBLE) != 0;
+        return (mEntry.getSbn().getNotification().flags & FLAG_BUBBLE) != 0;
+    }
+
+    public void enable(int option) {
+        mFlags |= option;
+    }
+
+    public void disable(int option) {
+        mFlags &= ~option;
+    }
+
     @Override
     public String toString() {
         return "Bubble{" + mKey + '}';
diff --git a/packages/SystemUI/src/com/android/systemui/bubbles/BubbleController.java b/packages/SystemUI/src/com/android/systemui/bubbles/BubbleController.java
index 6a5e94c..ec1216d 100644
--- a/packages/SystemUI/src/com/android/systemui/bubbles/BubbleController.java
+++ b/packages/SystemUI/src/com/android/systemui/bubbles/BubbleController.java
@@ -42,6 +42,7 @@
 import static java.lang.annotation.ElementType.PARAMETER;
 import static java.lang.annotation.RetentionPolicy.SOURCE;
 
+import android.annotation.NonNull;
 import android.annotation.UserIdInt;
 import android.app.ActivityManager.RunningTaskInfo;
 import android.app.INotificationManager;
@@ -237,7 +238,7 @@
          * This can happen when an app cancels a bubbled notification or when the user dismisses a
          * bubble.
          */
-        void removeNotification(NotificationEntry entry, int reason);
+        void removeNotification(@NonNull NotificationEntry entry, int reason);
 
         /**
          * Called when a bubbled notification has changed whether it should be
@@ -253,7 +254,7 @@
          * removes all remnants of the group's summary from the notification pipeline.
          * TODO: (b/145659174) Only old pipeline needs this - delete post-migration.
          */
-        void maybeCancelSummary(NotificationEntry entry);
+        void maybeCancelSummary(@NonNull NotificationEntry entry);
     }
 
     /**
@@ -476,7 +477,7 @@
 
         addNotifCallback(new NotifCallback() {
             @Override
-            public void removeNotification(NotificationEntry entry, int reason) {
+            public void removeNotification(@NonNull final NotificationEntry entry, int reason) {
                 mNotificationEntryManager.performRemoveNotification(entry.getSbn(),
                         reason);
             }
@@ -487,7 +488,7 @@
             }
 
             @Override
-            public void maybeCancelSummary(NotificationEntry entry) {
+            public void maybeCancelSummary(@NonNull final NotificationEntry entry) {
                 // Check if removed bubble has an associated suppressed group summary that needs
                 // to be removed now.
                 final String groupKey = entry.getSbn().getGroupKey();
@@ -696,10 +697,12 @@
         mBubbleIconFactory = new BubbleIconFactory(mContext);
         // Reload each bubble
         for (Bubble b: mBubbleData.getBubbles()) {
-            b.inflate(null /* callback */, mContext, mStackView, mBubbleIconFactory);
+            b.inflate(null /* callback */, mContext, mStackView, mBubbleIconFactory,
+                    false /* skipInflation */);
         }
         for (Bubble b: mBubbleData.getOverflowBubbles()) {
-            b.inflate(null /* callback */, mContext, mStackView, mBubbleIconFactory);
+            b.inflate(null /* callback */, mContext, mStackView, mBubbleIconFactory,
+                    false /* skipInflation */);
         }
     }
 
@@ -791,7 +794,7 @@
             if (bubble != null) {
                 mBubbleData.promoteBubbleFromOverflow(bubble, mStackView, mBubbleIconFactory);
             }
-        } else if (bubble.getEntry().isBubble()){
+        } else if (bubble.isBubble()) {
             mBubbleData.setSelectedBubble(bubble);
         }
         mBubbleData.setExpanded(true);
@@ -820,10 +823,33 @@
         updateBubble(notif, suppressFlyout, true /* showInShade */);
     }
 
+    /**
+     * Fills the overflow bubbles by loading them from disk.
+     */
+    void loadOverflowBubblesFromDisk() {
+        if (!mBubbleData.getOverflowBubbles().isEmpty()) {
+            // we don't need to load overflow bubbles from disk if it is already in memory
+            return;
+        }
+        mDataRepository.loadBubbles((bubbles) -> {
+            bubbles.forEach(bubble -> {
+                if (mBubbleData.getBubbles().contains(bubble)) {
+                    // if the bubble is already active, there's no need to push it to overflow
+                    return;
+                }
+                bubble.inflate((b) -> mBubbleData.overflowBubble(DISMISS_AGED, bubble),
+                        mContext, mStackView, mBubbleIconFactory, true /* skipInflation */);
+            });
+            return null;
+        });
+    }
+
     void updateBubble(NotificationEntry notif, boolean suppressFlyout, boolean showInShade) {
         if (mStackView == null) {
             // Lazy init stack view when a bubble is created
             ensureStackViewCreated();
+            // Lazy load overflow bubbles from disk
+            loadOverflowBubblesFromDisk();
         }
         // If this is an interruptive notif, mark that it's interrupted
         if (notif.getImportance() >= NotificationManager.IMPORTANCE_HIGH) {
@@ -843,11 +869,11 @@
                             return;
                         }
                         mHandler.post(
-                                () -> removeBubble(bubble.getEntry(),
+                                () -> removeBubble(bubble.getKey(),
                                         BubbleController.DISMISS_INVALID_INTENT));
                     });
                 },
-                mContext, mStackView, mBubbleIconFactory);
+                mContext, mStackView, mBubbleIconFactory, false /* skipInflation */);
     }
 
     /**
@@ -859,7 +885,10 @@
      * @param entry the notification to change bubble state for.
      * @param shouldBubble whether the notification should show as a bubble or not.
      */
-    public void onUserChangedBubble(NotificationEntry entry, boolean shouldBubble) {
+    public void onUserChangedBubble(@Nullable final NotificationEntry entry, boolean shouldBubble) {
+        if (entry == null) {
+            return;
+        }
         NotificationChannel channel = entry.getChannel();
         final String appPkg = entry.getSbn().getPackageName();
         final int appUid = entry.getSbn().getUid();
@@ -897,14 +926,14 @@
     }
 
     /**
-     * Removes the bubble with the given NotificationEntry.
+     * Removes the bubble with the given key.
      * <p>
      * Must be called from the main thread.
      */
     @MainThread
-    void removeBubble(NotificationEntry entry, int reason) {
-        if (mBubbleData.hasAnyBubbleWithKey(entry.getKey())) {
-            mBubbleData.notificationEntryRemoved(entry, reason);
+    void removeBubble(String key, int reason) {
+        if (mBubbleData.hasAnyBubbleWithKey(key)) {
+            mBubbleData.notificationEntryRemoved(key, reason);
         }
     }
 
@@ -920,7 +949,7 @@
                 && canLaunchInActivityView(mContext, entry);
         if (!shouldBubble && mBubbleData.hasAnyBubbleWithKey(entry.getKey())) {
             // It was previously a bubble but no longer a bubble -- lets remove it
-            removeBubble(entry, DISMISS_NO_LONGER_BUBBLE);
+            removeBubble(entry.getKey(), DISMISS_NO_LONGER_BUBBLE);
         } else if (shouldBubble) {
             updateBubble(entry);
         }
@@ -934,10 +963,10 @@
             // Remove any associated bubble children with the summary
             final List<Bubble> bubbleChildren = mBubbleData.getBubblesInGroup(groupKey);
             for (int i = 0; i < bubbleChildren.size(); i++) {
-                removeBubble(bubbleChildren.get(i).getEntry(), DISMISS_GROUP_CANCELLED);
+                removeBubble(bubbleChildren.get(i).getKey(), DISMISS_GROUP_CANCELLED);
             }
         } else {
-            removeBubble(entry, DISMISS_NOTIF_CANCEL);
+            removeBubble(entry.getKey(), DISMISS_NOTIF_CANCEL);
         }
     }
 
@@ -959,7 +988,8 @@
             rankingMap.getRanking(key, mTmpRanking);
             boolean isActiveBubble = mBubbleData.hasAnyBubbleWithKey(key);
             if (isActiveBubble && !mTmpRanking.canBubble()) {
-                mBubbleData.notificationEntryRemoved(entry, BubbleController.DISMISS_BLOCKED);
+                mBubbleData.notificationEntryRemoved(entry.getKey(),
+                        BubbleController.DISMISS_BLOCKED);
             } else if (entry != null && mTmpRanking.isBubble() && !isActiveBubble) {
                 entry.setFlagBubble(true);
                 onEntryUpdated(entry);
@@ -969,9 +999,15 @@
 
     private void setIsBubble(Bubble b, boolean isBubble) {
         if (isBubble) {
-            b.getEntry().getSbn().getNotification().flags |= FLAG_BUBBLE;
+            if (b.getEntry() != null) {
+                b.getEntry().getSbn().getNotification().flags |= FLAG_BUBBLE;
+            }
+            b.enable(FLAG_BUBBLE);
         } else {
-            b.getEntry().getSbn().getNotification().flags &= ~FLAG_BUBBLE;
+            if (b.getEntry() != null) {
+                b.getEntry().getSbn().getNotification().flags &= ~FLAG_BUBBLE;
+            }
+            b.disable(FLAG_BUBBLE);
         }
         try {
             mBarService.onNotificationBubbleChanged(b.getKey(), isBubble, 0);
@@ -1019,23 +1055,27 @@
                         // The bubble is now gone & the notification is hidden from the shade, so
                         // time to actually remove it
                         for (NotifCallback cb : mCallbacks) {
-                            cb.removeNotification(bubble.getEntry(), REASON_CANCEL);
+                            if (bubble.getEntry() != null) {
+                                cb.removeNotification(bubble.getEntry(), REASON_CANCEL);
+                            }
                         }
                     } else {
-                        if (bubble.getEntry().isBubble() && bubble.showInShade()) {
+                        if (bubble.isBubble() && bubble.showInShade()) {
                             setIsBubble(bubble, /* isBubble */ false);
                         }
-                        if (bubble.getEntry().getRow() != null) {
+                        if (bubble.getEntry() != null && bubble.getEntry().getRow() != null) {
                             bubble.getEntry().getRow().updateBubbleButton();
                         }
                     }
 
                 }
-                final String groupKey = bubble.getEntry().getSbn().getGroupKey();
-                if (mBubbleData.getBubblesInGroup(groupKey).isEmpty()) {
-                    // Time to potentially remove the summary
-                    for (NotifCallback cb : mCallbacks) {
-                        cb.maybeCancelSummary(bubble.getEntry());
+                if (bubble.getEntry() != null) {
+                    final String groupKey = bubble.getEntry().getSbn().getGroupKey();
+                    if (mBubbleData.getBubblesInGroup(groupKey).isEmpty()) {
+                        // Time to potentially remove the summary
+                        for (NotifCallback cb : mCallbacks) {
+                            cb.maybeCancelSummary(bubble.getEntry());
+                        }
                     }
                 }
             }
@@ -1060,7 +1100,7 @@
 
             if (update.selectionChanged) {
                 mStackView.setSelectedBubble(update.selectedBubble);
-                if (update.selectedBubble != null) {
+                if (update.selectedBubble != null && update.selectedBubble.getEntry() != null) {
                     mNotificationGroupManager.updateSuppression(
                             update.selectedBubble.getEntry());
                 }
diff --git a/packages/SystemUI/src/com/android/systemui/bubbles/BubbleData.java b/packages/SystemUI/src/com/android/systemui/bubbles/BubbleData.java
index 65d5beb..7b323ce 100644
--- a/packages/SystemUI/src/com/android/systemui/bubbles/BubbleData.java
+++ b/packages/SystemUI/src/com/android/systemui/bubbles/BubbleData.java
@@ -23,6 +23,7 @@
 
 import static java.util.stream.Collectors.toList;
 
+import android.annotation.NonNull;
 import android.app.Notification;
 import android.app.PendingIntent;
 import android.content.Context;
@@ -223,7 +224,7 @@
                             false, /* showInShade */ true);
                     setSelectedBubble(bubble);
                 },
-                mContext, stack, factory);
+                mContext, stack, factory, false /* skipInflation */);
         dispatchPendingChanges();
     }
 
@@ -278,7 +279,8 @@
         }
         mPendingBubbles.remove(bubble); // No longer pending once we're here
         Bubble prevBubble = getBubbleInStackWithKey(bubble.getKey());
-        suppressFlyout |= !bubble.getEntry().getRanking().visuallyInterruptive();
+        suppressFlyout |= bubble.getEntry() == null
+                || !bubble.getEntry().getRanking().visuallyInterruptive();
 
         if (prevBubble == null) {
             // Create a new bubble
@@ -307,11 +309,14 @@
         dispatchPendingChanges();
     }
 
-    public void notificationEntryRemoved(NotificationEntry entry, @DismissReason int reason) {
+    /**
+     * Called when a notification associated with a bubble is removed.
+     */
+    public void notificationEntryRemoved(String key, @DismissReason int reason) {
         if (DEBUG_BUBBLE_DATA) {
-            Log.d(TAG, "notificationEntryRemoved: entry=" + entry + " reason=" + reason);
+            Log.d(TAG, "notificationEntryRemoved: key=" + key + " reason=" + reason);
         }
-        doRemove(entry.getKey(), reason);
+        doRemove(key, reason);
         dispatchPendingChanges();
     }
 
@@ -359,7 +364,7 @@
             return bubbleChildren;
         }
         for (Bubble b : mBubbles) {
-            if (groupKey.equals(b.getEntry().getSbn().getGroupKey())) {
+            if (b.getEntry() != null && groupKey.equals(b.getEntry().getSbn().getGroupKey())) {
                 bubbleChildren.add(b);
             }
         }
@@ -470,7 +475,9 @@
             Bubble newSelected = mBubbles.get(newIndex);
             setSelectedBubbleInternal(newSelected);
         }
-        maybeSendDeleteIntent(reason, bubbleToRemove.getEntry());
+        if (bubbleToRemove.getEntry() != null) {
+            maybeSendDeleteIntent(reason, bubbleToRemove.getEntry());
+        }
     }
 
     void overflowBubble(@DismissReason int reason, Bubble bubble) {
@@ -744,7 +751,8 @@
         return true;
     }
 
-    private void maybeSendDeleteIntent(@DismissReason int reason, NotificationEntry entry) {
+    private void maybeSendDeleteIntent(@DismissReason int reason,
+            @NonNull final NotificationEntry entry) {
         if (reason == BubbleController.DISMISS_USER_GESTURE) {
             Notification.BubbleMetadata bubbleMetadata = entry.getBubbleMetadata();
             PendingIntent deleteIntent = bubbleMetadata != null
diff --git a/packages/SystemUI/src/com/android/systemui/bubbles/BubbleDataRepository.kt b/packages/SystemUI/src/com/android/systemui/bubbles/BubbleDataRepository.kt
index ba93f41..6df3132 100644
--- a/packages/SystemUI/src/com/android/systemui/bubbles/BubbleDataRepository.kt
+++ b/packages/SystemUI/src/com/android/systemui/bubbles/BubbleDataRepository.kt
@@ -74,7 +74,10 @@
 
     private fun transform(userId: Int, bubbles: List<Bubble>): List<BubbleEntity> {
         return bubbles.mapNotNull { b ->
-            val shortcutId = b.shortcutInfo?.id ?: return@mapNotNull null
+            var shortcutId = b.shortcutInfo?.id
+            if (shortcutId == null) shortcutId = b.entry?.bubbleMetadata?.shortcutId
+            if (shortcutId == null) shortcutId = b.entry?.ranking?.shortcutInfo?.id
+            if (shortcutId == null) return@mapNotNull null
             BubbleEntity(userId, b.packageName, shortcutId)
         }
     }
@@ -108,7 +111,6 @@
     /**
      * Load bubbles from disk.
      */
-    // TODO: call this method from BubbleController and update UI
     @SuppressLint("WrongConstant")
     fun loadBubbles(cb: (List<Bubble>) -> Unit) = ioScope.launch {
         /**
diff --git a/packages/SystemUI/src/com/android/systemui/bubbles/BubbleExpandedView.java b/packages/SystemUI/src/com/android/systemui/bubbles/BubbleExpandedView.java
index baf92dc..0a1a0f7 100644
--- a/packages/SystemUI/src/com/android/systemui/bubbles/BubbleExpandedView.java
+++ b/packages/SystemUI/src/com/android/systemui/bubbles/BubbleExpandedView.java
@@ -65,7 +65,6 @@
 import com.android.systemui.R;
 import com.android.systemui.recents.TriangleShape;
 import com.android.systemui.statusbar.AlphaOptimizedButton;
-import com.android.systemui.statusbar.notification.collection.NotificationEntry;
 
 /**
  * Container for the expanded bubble view, handles rendering the caret and settings icon.
@@ -161,7 +160,7 @@
                             // the bubble again so we'll just remove it.
                             Log.w(TAG, "Exception while displaying bubble: " + getBubbleKey()
                                     + ", " + e.getMessage() + "; removing bubble");
-                            mBubbleController.removeBubble(getBubbleEntry(),
+                            mBubbleController.removeBubble(getBubbleKey(),
                                     BubbleController.DISMISS_INVALID_INTENT);
                         }
                     });
@@ -205,7 +204,7 @@
             }
             if (mBubble != null) {
                 // Must post because this is called from a binder thread.
-                post(() -> mBubbleController.removeBubble(mBubble.getEntry(),
+                post(() -> mBubbleController.removeBubble(mBubble.getKey(),
                         BubbleController.DISMISS_TASK_FINISHED));
             }
         }
@@ -292,10 +291,6 @@
         return mBubble != null ? mBubble.getKey() : "null";
     }
 
-    private NotificationEntry getBubbleEntry() {
-        return mBubble != null ? mBubble.getEntry() : null;
-    }
-
     void setManageClickListener(OnClickListener manageClickListener) {
         findViewById(R.id.settings_button).setOnClickListener(manageClickListener);
     }
diff --git a/packages/SystemUI/src/com/android/systemui/bubbles/BubbleStackView.java b/packages/SystemUI/src/com/android/systemui/bubbles/BubbleStackView.java
index 2cb097f..c06063b 100644
--- a/packages/SystemUI/src/com/android/systemui/bubbles/BubbleStackView.java
+++ b/packages/SystemUI/src/com/android/systemui/bubbles/BubbleStackView.java
@@ -49,7 +49,6 @@
 import android.graphics.Region;
 import android.os.Bundle;
 import android.os.Vibrator;
-import android.service.notification.StatusBarNotification;
 import android.util.Log;
 import android.view.Choreographer;
 import android.view.DisplayCutout;
@@ -916,10 +915,10 @@
                     showManageMenu(false /* show */);
                     final Bubble bubble = mBubbleData.getSelectedBubble();
                     if (bubble != null && mBubbleData.hasBubbleInStackWithKey(bubble.getKey())) {
-                        final Intent intent = bubble.getSettingsIntent();
+                        final Intent intent = bubble.getSettingsIntent(mContext);
                         collapseStack(() -> {
-                            mContext.startActivityAsUser(
-                                    intent, bubble.getEntry().getSbn().getUser());
+
+                            mContext.startActivityAsUser(intent, bubble.getUser());
                             logBubbleClickEvent(
                                     bubble,
                                     SysUiStatsLog.BUBBLE_UICHANGED__ACTION__HEADER_GO_TO_SETTINGS);
@@ -1153,13 +1152,19 @@
         for (int i = 0; i < mBubbleData.getBubbles().size(); i++) {
             final Bubble bubble = mBubbleData.getBubbles().get(i);
             final String appName = bubble.getAppName();
-            final Notification notification = bubble.getEntry().getSbn().getNotification();
-            final CharSequence titleCharSeq =
-                    notification.extras.getCharSequence(Notification.EXTRA_TITLE);
 
-            String titleStr = getResources().getString(R.string.notification_bubble_title);
+            final CharSequence titleCharSeq;
+            if (bubble.getEntry() == null) {
+                titleCharSeq = null;
+            } else {
+                titleCharSeq = bubble.getEntry().getSbn().getNotification().extras.getCharSequence(
+                        Notification.EXTRA_TITLE);
+            }
+            final String titleStr;
             if (titleCharSeq != null) {
                 titleStr = titleCharSeq.toString();
+            } else {
+                titleStr = getResources().getString(R.string.notification_bubble_title);
             }
 
             if (bubble.getIconView() != null) {
@@ -1777,7 +1782,7 @@
     private void dismissBubbleIfExists(@Nullable Bubble bubble) {
         if (bubble != null && mBubbleData.hasBubbleInStackWithKey(bubble.getKey())) {
             mBubbleData.notificationEntryRemoved(
-                    bubble.getEntry(), BubbleController.DISMISS_USER_GESTURE);
+                    bubble.getKey(), BubbleController.DISMISS_USER_GESTURE);
         }
     }
 
@@ -2275,18 +2280,12 @@
      * @param action the user interaction enum.
      */
     private void logBubbleClickEvent(Bubble bubble, int action) {
-        StatusBarNotification notification = bubble.getEntry().getSbn();
-        SysUiStatsLog.write(SysUiStatsLog.BUBBLE_UI_CHANGED,
-                notification.getPackageName(),
-                notification.getNotification().getChannelId(),
-                notification.getId(),
-                getBubbleIndex(getExpandedBubble()),
+        bubble.logUIEvent(
                 getBubbleCount(),
                 action,
                 getNormalizedXPosition(),
                 getNormalizedYPosition(),
-                bubble.showInShade(),
-                bubble.isOngoing(),
-                false /* isAppForeground (unused) */);
+                getBubbleIndex(getExpandedBubble())
+        );
     }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/bubbles/BubbleViewInfoTask.java b/packages/SystemUI/src/com/android/systemui/bubbles/BubbleViewInfoTask.java
index 8a57a73..525d5b5 100644
--- a/packages/SystemUI/src/com/android/systemui/bubbles/BubbleViewInfoTask.java
+++ b/packages/SystemUI/src/com/android/systemui/bubbles/BubbleViewInfoTask.java
@@ -37,6 +37,7 @@
 import android.graphics.drawable.Icon;
 import android.os.AsyncTask;
 import android.os.Parcelable;
+import android.os.UserHandle;
 import android.service.notification.StatusBarNotification;
 import android.text.TextUtils;
 import android.util.Log;
@@ -74,6 +75,7 @@
     private WeakReference<Context> mContext;
     private WeakReference<BubbleStackView> mStackView;
     private BubbleIconFactory mIconFactory;
+    private boolean mSkipInflation;
     private Callback mCallback;
 
     /**
@@ -84,17 +86,20 @@
             Context context,
             BubbleStackView stackView,
             BubbleIconFactory factory,
+            boolean skipInflation,
             Callback c) {
         mBubble = b;
         mContext = new WeakReference<>(context);
         mStackView = new WeakReference<>(stackView);
         mIconFactory = factory;
+        mSkipInflation = skipInflation;
         mCallback = c;
     }
 
     @Override
     protected BubbleViewInfo doInBackground(Void... voids) {
-        return BubbleViewInfo.populate(mContext.get(), mStackView.get(), mIconFactory, mBubble);
+        return BubbleViewInfo.populate(mContext.get(), mStackView.get(), mIconFactory, mBubble,
+                mSkipInflation);
     }
 
     @Override
@@ -123,11 +128,36 @@
 
         @Nullable
         static BubbleViewInfo populate(Context c, BubbleStackView stackView,
-                BubbleIconFactory iconFactory, Bubble b) {
+                BubbleIconFactory iconFactory, Bubble b, boolean skipInflation) {
+            final NotificationEntry entry = b.getEntry();
+            if (entry == null) {
+                // populate from ShortcutInfo when NotificationEntry is not available
+                final ShortcutInfo s = b.getShortcutInfo();
+                return populate(c, stackView, iconFactory, skipInflation || b.isInflated(),
+                        s.getPackage(), s.getUserHandle(), s, null);
+            }
+            final StatusBarNotification sbn = entry.getSbn();
+            final String bubbleShortcutId =  entry.getBubbleMetadata().getShortcutId();
+            final ShortcutInfo si = bubbleShortcutId == null
+                    ? null : entry.getRanking().getShortcutInfo();
+            return populate(
+                    c, stackView, iconFactory, skipInflation || b.isInflated(),
+                    sbn.getPackageName(), sbn.getUser(), si, entry);
+        }
+
+        private static BubbleViewInfo populate(
+                @NonNull final Context c,
+                @NonNull final BubbleStackView stackView,
+                @NonNull final BubbleIconFactory iconFactory,
+                final boolean isInflated,
+                @NonNull final String packageName,
+                @NonNull final UserHandle user,
+                @Nullable final ShortcutInfo shortcutInfo,
+                @Nullable final NotificationEntry entry) {
             BubbleViewInfo info = new BubbleViewInfo();
 
             // View inflation: only should do this once per bubble
-            if (!b.isInflated()) {
+            if (!isInflated) {
                 LayoutInflater inflater = LayoutInflater.from(c);
                 info.imageView = (BadgedImageView) inflater.inflate(
                         R.layout.bubble_view, stackView, false /* attachToRoot */);
@@ -137,12 +167,8 @@
                 info.expandedView.setStackView(stackView);
             }
 
-            StatusBarNotification sbn = b.getEntry().getSbn();
-            String packageName = sbn.getPackageName();
-
-            String bubbleShortcutId =  b.getEntry().getBubbleMetadata().getShortcutId();
-            if (bubbleShortcutId != null) {
-                info.shortcutInfo = b.getEntry().getRanking().getShortcutInfo();
+            if (shortcutInfo != null) {
+                info.shortcutInfo = shortcutInfo;
             }
 
             // App name & app icon
@@ -161,7 +187,7 @@
                     info.appName = String.valueOf(pm.getApplicationLabel(appInfo));
                 }
                 appIcon = pm.getApplicationIcon(packageName);
-                badgedIcon = pm.getUserBadgedIcon(appIcon, sbn.getUser());
+                badgedIcon = pm.getUserBadgedIcon(appIcon, user);
             } catch (PackageManager.NameNotFoundException exception) {
                 // If we can't find package... don't think we should show the bubble.
                 Log.w(TAG, "Unable to find package: " + packageName);
@@ -170,7 +196,7 @@
 
             // Badged bubble image
             Drawable bubbleDrawable = iconFactory.getBubbleDrawable(c, info.shortcutInfo,
-                    b.getEntry().getBubbleMetadata());
+                    entry == null ? null : entry.getBubbleMetadata());
             if (bubbleDrawable == null) {
                 // Default to app icon
                 bubbleDrawable = appIcon;
@@ -196,7 +222,9 @@
                     Color.WHITE, WHITE_SCRIM_ALPHA);
 
             // Flyout
-            info.flyoutMessage = extractFlyoutMessage(c, b.getEntry());
+            if (entry != null) {
+                info.flyoutMessage = extractFlyoutMessage(c, entry);
+            }
             return info;
         }
     }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/coordinator/BubbleCoordinator.java b/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/coordinator/BubbleCoordinator.java
index 92426e5..57be1a5 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/coordinator/BubbleCoordinator.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/coordinator/BubbleCoordinator.java
@@ -19,6 +19,8 @@
 import static android.service.notification.NotificationStats.DISMISSAL_OTHER;
 import static android.service.notification.NotificationStats.DISMISS_SENTIMENT_NEUTRAL;
 
+import android.annotation.NonNull;
+
 import com.android.internal.statusbar.NotificationVisibility;
 import com.android.systemui.bubbles.BubbleController;
 import com.android.systemui.statusbar.notification.collection.NotifCollection;
@@ -121,7 +123,7 @@
     private final BubbleController.NotifCallback mNotifCallback =
             new BubbleController.NotifCallback() {
         @Override
-        public void removeNotification(NotificationEntry entry, int reason) {
+        public void removeNotification(@NonNull final NotificationEntry entry, int reason) {
             if (isInterceptingDismissal(entry)) {
                 mInterceptedDismissalEntries.remove(entry.getKey());
                 mOnEndDismissInterception.onEndDismissInterception(mDismissInterceptor, entry,
@@ -141,7 +143,7 @@
         }
 
         @Override
-        public void maybeCancelSummary(NotificationEntry entry) {
+        public void maybeCancelSummary(@NonNull final NotificationEntry entry) {
             // no-op
         }
     };
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/phone/NotificationGroupManager.java b/packages/SystemUI/src/com/android/systemui/statusbar/phone/NotificationGroupManager.java
index 80785db..d0a872e 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/phone/NotificationGroupManager.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/phone/NotificationGroupManager.java
@@ -16,6 +16,7 @@
 
 package com.android.systemui.statusbar.phone;
 
+import android.annotation.NonNull;
 import android.annotation.Nullable;
 import android.service.notification.StatusBarNotification;
 import android.util.ArraySet;
@@ -454,7 +455,7 @@
      * If there is a {@link NotificationGroup} associated with the provided entry, this method
      * will update the suppression of that group.
      */
-    public void updateSuppression(NotificationEntry entry) {
+    public void updateSuppression(@NonNull final NotificationEntry entry) {
         NotificationGroup group = mGroupMap.get(getGroupKey(entry.getSbn()));
         if (group != null) {
             updateSuppression(group);
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 96e868d..52923ab 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/bubbles/BubbleControllerTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/bubbles/BubbleControllerTest.java
@@ -317,7 +317,7 @@
         verify(mNotificationEntryManager).updateNotifications(any());
 
         mBubbleController.removeBubble(
-                mRow.getEntry(), BubbleController.DISMISS_USER_GESTURE);
+                mRow.getEntry().getKey(), BubbleController.DISMISS_USER_GESTURE);
         assertNull(mBubbleData.getBubbleInStackWithKey(mRow.getEntry().getKey()));
         verify(mNotificationEntryManager, times(2)).updateNotifications(anyString());
 
@@ -329,7 +329,7 @@
         mBubbleController.updateBubble(mRow2.getEntry());
         mBubbleController.updateBubble(mRow.getEntry());
         mBubbleController.removeBubble(
-                mRow.getEntry(), BubbleController.DISMISS_USER_GESTURE);
+                mRow.getEntry().getKey(), BubbleController.DISMISS_USER_GESTURE);
 
         Bubble b = mBubbleData.getOverflowBubbleWithKey(mRow.getEntry().getKey());
         assertThat(mBubbleData.getOverflowBubbles()).isEqualTo(ImmutableList.of(b));
@@ -350,9 +350,10 @@
         mBubbleController.updateBubble(mRow.getEntry(), /* suppressFlyout */
                 false, /* showInShade */ true);
         mBubbleController.removeBubble(
-                mRow.getEntry(), BubbleController.DISMISS_USER_GESTURE);
+                mRow.getEntry().getKey(), BubbleController.DISMISS_USER_GESTURE);
 
-        mBubbleController.removeBubble(mRow.getEntry(), BubbleController.DISMISS_NOTIF_CANCEL);
+        mBubbleController.removeBubble(
+                mRow.getEntry().getKey(), BubbleController.DISMISS_NOTIF_CANCEL);
         verify(mNotificationEntryManager, times(1)).performRemoveNotification(
                 eq(mRow.getEntry().getSbn()), anyInt());
         assertThat(mBubbleData.getOverflowBubbles()).isEmpty();
@@ -365,7 +366,7 @@
         assertTrue(mBubbleController.hasBubbles());
 
         mBubbleController.removeBubble(
-                mRow.getEntry(), BubbleController.DISMISS_USER_CHANGED);
+                mRow.getEntry().getKey(), BubbleController.DISMISS_USER_CHANGED);
         verify(mNotificationEntryManager, never()).performRemoveNotification(
                 eq(mRow.getEntry().getSbn()), anyInt());
         assertFalse(mBubbleController.hasBubbles());
@@ -563,7 +564,8 @@
 
         // Dismiss currently expanded
         mBubbleController.removeBubble(
-                mBubbleData.getBubbleInStackWithKey(stackView.getExpandedBubble().getKey()).getEntry(),
+                mBubbleData.getBubbleInStackWithKey(stackView.getExpandedBubble().getKey())
+                        .getEntry().getKey(),
                 BubbleController.DISMISS_USER_GESTURE);
         verify(mBubbleExpandListener).onBubbleExpandChanged(false, mRow2.getEntry().getKey());
 
@@ -574,7 +576,8 @@
 
         // Dismiss that one
         mBubbleController.removeBubble(
-                mBubbleData.getBubbleInStackWithKey(stackView.getExpandedBubble().getKey()).getEntry(),
+                mBubbleData.getBubbleInStackWithKey(stackView.getExpandedBubble().getKey())
+                        .getEntry().getKey(),
                 BubbleController.DISMISS_USER_GESTURE);
 
         // Make sure state changes and collapse happens
@@ -700,7 +703,7 @@
     @Test
     public void testDeleteIntent_removeBubble_aged() throws PendingIntent.CanceledException {
         mBubbleController.updateBubble(mRow.getEntry());
-        mBubbleController.removeBubble(mRow.getEntry(), BubbleController.DISMISS_AGED);
+        mBubbleController.removeBubble(mRow.getEntry().getKey(), BubbleController.DISMISS_AGED);
         verify(mDeleteIntent, never()).send();
     }
 
@@ -708,7 +711,7 @@
     public void testDeleteIntent_removeBubble_user() throws PendingIntent.CanceledException {
         mBubbleController.updateBubble(mRow.getEntry());
         mBubbleController.removeBubble(
-                mRow.getEntry(), BubbleController.DISMISS_USER_GESTURE);
+                mRow.getEntry().getKey(), BubbleController.DISMISS_USER_GESTURE);
         verify(mDeleteIntent, times(1)).send();
     }
 
@@ -808,7 +811,7 @@
 
         // Dismiss the bubble into overflow.
         mBubbleController.removeBubble(
-                mRow.getEntry(), BubbleController.DISMISS_USER_GESTURE);
+                mRow.getEntry().getKey(), BubbleController.DISMISS_USER_GESTURE);
         assertFalse(mBubbleController.hasBubbles());
 
         boolean intercepted = mRemoveInterceptor.onNotificationRemoveRequested(
@@ -829,7 +832,7 @@
                 mRow.getEntry()));
 
         mBubbleController.removeBubble(
-                mRow.getEntry(), BubbleController.DISMISS_NO_LONGER_BUBBLE);
+                mRow.getEntry().getKey(), BubbleController.DISMISS_NO_LONGER_BUBBLE);
         assertFalse(mBubbleController.hasBubbles());
 
         boolean intercepted = mRemoveInterceptor.onNotificationRemoveRequested(
@@ -851,12 +854,12 @@
 
         mBubbleData.setMaxOverflowBubbles(1);
         mBubbleController.removeBubble(
-                mRow.getEntry(), BubbleController.DISMISS_USER_GESTURE);
+                mRow.getEntry().getKey(), BubbleController.DISMISS_USER_GESTURE);
         assertEquals(mBubbleData.getBubbles().size(), 2);
         assertEquals(mBubbleData.getOverflowBubbles().size(), 1);
 
         mBubbleController.removeBubble(
-                mRow2.getEntry(), BubbleController.DISMISS_USER_GESTURE);
+                mRow2.getEntry().getKey(), BubbleController.DISMISS_USER_GESTURE);
         // Overflow max of 1 is reached; mRow is oldest, so it gets removed
         verify(mNotificationEntryManager, times(1)).performRemoveNotification(
                 mRow.getEntry().getSbn(), REASON_CANCEL);
diff --git a/packages/SystemUI/tests/src/com/android/systemui/bubbles/BubbleDataTest.java b/packages/SystemUI/tests/src/com/android/systemui/bubbles/BubbleDataTest.java
index 66f119a..882504b 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/bubbles/BubbleDataTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/bubbles/BubbleDataTest.java
@@ -20,11 +20,8 @@
 
 import static com.google.common.truth.Truth.assertThat;
 
-import static org.mockito.ArgumentMatchers.anyInt;
-import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.reset;
-import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.verifyZeroInteractions;
 import static org.mockito.Mockito.when;
@@ -180,7 +177,8 @@
         mBubbleData.setListener(mListener);
 
         // Test
-        mBubbleData.notificationEntryRemoved(mEntryA1, BubbleController.DISMISS_USER_GESTURE);
+        mBubbleData.notificationEntryRemoved(
+                mEntryA1.getKey(), BubbleController.DISMISS_USER_GESTURE);
 
         // Verify
         verifyUpdateReceived();
@@ -297,12 +295,14 @@
         mBubbleData.setListener(mListener);
 
         mBubbleData.setMaxOverflowBubbles(1);
-        mBubbleData.notificationEntryRemoved(mEntryA1, BubbleController.DISMISS_USER_GESTURE);
+        mBubbleData.notificationEntryRemoved(
+                mEntryA1.getKey(), BubbleController.DISMISS_USER_GESTURE);
         verifyUpdateReceived();
         assertOverflowChangedTo(ImmutableList.of(mBubbleA1));
 
         // Overflow max of 1 is reached; A1 is oldest, so it gets removed
-        mBubbleData.notificationEntryRemoved(mEntryA2, BubbleController.DISMISS_USER_GESTURE);
+        mBubbleData.notificationEntryRemoved(
+                mEntryA2.getKey(), BubbleController.DISMISS_USER_GESTURE);
         verifyUpdateReceived();
         assertOverflowChangedTo(ImmutableList.of(mBubbleA2));
     }
@@ -449,7 +449,8 @@
         mBubbleData.setListener(mListener);
 
         // Test
-        mBubbleData.notificationEntryRemoved(mEntryA2, BubbleController.DISMISS_USER_GESTURE);
+        mBubbleData.notificationEntryRemoved(
+                mEntryA2.getKey(), BubbleController.DISMISS_USER_GESTURE);
         verifyUpdateReceived();
         assertOrderChangedTo(mBubbleB2, mBubbleB1, mBubbleA1);
     }
@@ -469,7 +470,8 @@
         mBubbleData.setListener(mListener);
 
         // Test
-        mBubbleData.notificationEntryRemoved(mEntryB1, BubbleController.DISMISS_USER_GESTURE);
+        mBubbleData.notificationEntryRemoved(
+                mEntryB1.getKey(), BubbleController.DISMISS_USER_GESTURE);
         verifyUpdateReceived();
         assertOrderNotChanged();
     }
@@ -489,7 +491,8 @@
         mBubbleData.setListener(mListener);
 
         // Test
-        mBubbleData.notificationEntryRemoved(mEntryA1, BubbleController.DISMISS_NOTIF_CANCEL);
+        mBubbleData.notificationEntryRemoved(
+                mEntryA1.getKey(), BubbleController.DISMISS_NOTIF_CANCEL);
         verifyUpdateReceived();
         assertOrderChangedTo(mBubbleB2, mBubbleB1, mBubbleA2);
     }
@@ -510,12 +513,14 @@
         mBubbleData.setListener(mListener);
 
         // Test
-        mBubbleData.notificationEntryRemoved(mEntryA1, BubbleController.DISMISS_NOTIF_CANCEL);
+        mBubbleData.notificationEntryRemoved(
+                mEntryA1.getKey(), BubbleController.DISMISS_NOTIF_CANCEL);
         verifyUpdateReceived();
         assertOverflowChangedTo(ImmutableList.of(mBubbleA2));
 
         // Test
-        mBubbleData.notificationEntryRemoved(mEntryA2, BubbleController.DISMISS_GROUP_CANCELLED);
+        mBubbleData.notificationEntryRemoved(
+                mEntryA2.getKey(), BubbleController.DISMISS_GROUP_CANCELLED);
         verifyUpdateReceived();
         assertOverflowChangedTo(ImmutableList.of());
     }
@@ -534,7 +539,8 @@
         mBubbleData.setListener(mListener);
 
         // Test
-        mBubbleData.notificationEntryRemoved(mEntryA2, BubbleController.DISMISS_NOTIF_CANCEL);
+        mBubbleData.notificationEntryRemoved(
+                mEntryA2.getKey(), BubbleController.DISMISS_NOTIF_CANCEL);
         verifyUpdateReceived();
         assertSelectionChangedTo(mBubbleB2);
     }
@@ -625,7 +631,8 @@
         mBubbleData.setListener(mListener);
 
         // Test
-        mBubbleData.notificationEntryRemoved(mEntryA1, BubbleController.DISMISS_USER_GESTURE);
+        mBubbleData.notificationEntryRemoved(
+                mEntryA1.getKey(), BubbleController.DISMISS_USER_GESTURE);
 
         // Verify the selection was cleared.
         verifyUpdateReceived();
@@ -779,7 +786,8 @@
         mBubbleData.setListener(mListener);
 
         // Test
-        mBubbleData.notificationEntryRemoved(mEntryB2, BubbleController.DISMISS_USER_GESTURE);
+        mBubbleData.notificationEntryRemoved(
+                mEntryB2.getKey(), BubbleController.DISMISS_USER_GESTURE);
         verifyUpdateReceived();
         assertOrderChangedTo(mBubbleB1, mBubbleA2, mBubbleA1);
     }
@@ -804,11 +812,13 @@
         mBubbleData.setListener(mListener);
 
         // Test
-        mBubbleData.notificationEntryRemoved(mEntryA2, BubbleController.DISMISS_USER_GESTURE);
+        mBubbleData.notificationEntryRemoved(
+                mEntryA2.getKey(), BubbleController.DISMISS_USER_GESTURE);
         verifyUpdateReceived();
         assertSelectionChangedTo(mBubbleA1);
 
-        mBubbleData.notificationEntryRemoved(mEntryA1, BubbleController.DISMISS_USER_GESTURE);
+        mBubbleData.notificationEntryRemoved(
+                mEntryA1.getKey(), BubbleController.DISMISS_USER_GESTURE);
         verifyUpdateReceived();
         assertSelectionChangedTo(mBubbleB1);
     }
@@ -923,7 +933,8 @@
         mBubbleData.setListener(mListener);
 
         // Test
-        mBubbleData.notificationEntryRemoved(mEntryA1, BubbleController.DISMISS_USER_GESTURE);
+        mBubbleData.notificationEntryRemoved(
+                mEntryA1.getKey(), BubbleController.DISMISS_USER_GESTURE);
         verifyUpdateReceived();
         assertExpandedChangedTo(false);
     }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/bubbles/NewNotifPipelineBubbleControllerTest.java b/packages/SystemUI/tests/src/com/android/systemui/bubbles/NewNotifPipelineBubbleControllerTest.java
index 73b8760..c16801c 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/bubbles/NewNotifPipelineBubbleControllerTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/bubbles/NewNotifPipelineBubbleControllerTest.java
@@ -285,7 +285,8 @@
         assertTrue(mBubbleController.hasBubbles());
         verify(mNotifCallback, times(1)).invalidateNotifications(anyString());
 
-        mBubbleController.removeBubble(mRow.getEntry(), BubbleController.DISMISS_USER_GESTURE);
+        mBubbleController.removeBubble(
+                mRow.getEntry().getKey(), BubbleController.DISMISS_USER_GESTURE);
         assertNull(mBubbleData.getBubbleInStackWithKey(mRow.getEntry().getKey()));
         verify(mNotifCallback, times(2)).invalidateNotifications(anyString());
     }
@@ -302,7 +303,8 @@
         mBubbleData.getBubbleInStackWithKey(mRow.getEntry().getKey()).setSuppressNotification(true);
 
         // Now remove the bubble
-        mBubbleController.removeBubble(mRow.getEntry(), BubbleController.DISMISS_USER_GESTURE);
+        mBubbleController.removeBubble(
+                mRow.getEntry().getKey(), BubbleController.DISMISS_USER_GESTURE);
         assertTrue(mBubbleData.hasOverflowBubbleWithKey(mRow.getEntry().getKey()));
 
         // We don't remove the notification since the bubble is still in overflow.
@@ -322,7 +324,8 @@
         mBubbleData.getBubbleInStackWithKey(mRow.getEntry().getKey()).setSuppressNotification(true);
 
         // Now remove the bubble
-        mBubbleController.removeBubble(mRow.getEntry(), BubbleController.DISMISS_NOTIF_CANCEL);
+        mBubbleController.removeBubble(
+                mRow.getEntry().getKey(), BubbleController.DISMISS_NOTIF_CANCEL);
         assertFalse(mBubbleData.hasOverflowBubbleWithKey(mRow.getEntry().getKey()));
 
         // Since the notif is dismissed and not in overflow, once the bubble is removed,
@@ -502,7 +505,8 @@
 
         // Dismiss currently expanded
         mBubbleController.removeBubble(
-                mBubbleData.getBubbleInStackWithKey(stackView.getExpandedBubble().getKey()).getEntry(),
+                mBubbleData.getBubbleInStackWithKey(
+                        stackView.getExpandedBubble().getKey()).getEntry().getKey(),
                 BubbleController.DISMISS_USER_GESTURE);
         verify(mBubbleExpandListener).onBubbleExpandChanged(false, mRow2.getEntry().getKey());
 
@@ -513,7 +517,8 @@
 
         // Dismiss that one
         mBubbleController.removeBubble(
-                mBubbleData.getBubbleInStackWithKey(stackView.getExpandedBubble().getKey()).getEntry(),
+                mBubbleData.getBubbleInStackWithKey(
+                        stackView.getExpandedBubble().getKey()).getEntry().getKey(),
                 BubbleController.DISMISS_USER_GESTURE);
 
         // Make sure state changes and collapse happens
@@ -613,7 +618,7 @@
     @Test
     public void testDeleteIntent_removeBubble_aged() throws PendingIntent.CanceledException {
         mBubbleController.updateBubble(mRow.getEntry());
-        mBubbleController.removeBubble(mRow.getEntry(), BubbleController.DISMISS_AGED);
+        mBubbleController.removeBubble(mRow.getEntry().getKey(), BubbleController.DISMISS_AGED);
         verify(mDeleteIntent, never()).send();
     }
 
@@ -621,7 +626,7 @@
     public void testDeleteIntent_removeBubble_user() throws PendingIntent.CanceledException {
         mBubbleController.updateBubble(mRow.getEntry());
         mBubbleController.removeBubble(
-                mRow.getEntry(), BubbleController.DISMISS_USER_GESTURE);
+                mRow.getEntry().getKey(), BubbleController.DISMISS_USER_GESTURE);
         verify(mDeleteIntent, times(1)).send();
     }
 
@@ -686,7 +691,7 @@
 
         // Dismiss the bubble
         mBubbleController.removeBubble(
-                mRow.getEntry(), BubbleController.DISMISS_USER_GESTURE);
+                mRow.getEntry().getKey(), BubbleController.DISMISS_USER_GESTURE);
         assertFalse(mBubbleController.hasBubbles());
 
         // Dismiss the notification
@@ -707,7 +712,7 @@
 
         // Dismiss the bubble
         mBubbleController.removeBubble(
-                mRow.getEntry(), BubbleController.DISMISS_NOTIF_CANCEL);
+                mRow.getEntry().getKey(), BubbleController.DISMISS_NOTIF_CANCEL);
         assertFalse(mBubbleController.hasBubbles());
 
         // Dismiss the notification