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

This reverts commit b28e0568a062e779816a7b5bf0e5e1eac7ca005d.
Reason for revert: Rolling-forward with fixed unit test

Bug: 156697897
Bug: 149713060
Change-Id: I21e4b17ee53ee7b6698115e1b497253bb54b355f
diff --git a/packages/SystemUI/src/com/android/systemui/bubbles/Bubble.java b/packages/SystemUI/src/com/android/systemui/bubbles/Bubble.java
index ecd8b45..8cfa2d6 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;
 
@@ -96,11 +103,38 @@
     private Bitmap mBadgedImage;
     private int mDotColor;
     private Path mDotPath;
+    private int mFlags;
 
-    // TODO: Decouple Bubble from NotificationEntry and transform ShortcutInfo into Bubble
+    /**
+     * 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();
+        mKey = key(shortcutInfo);
+        mFlags = 0;
     }
 
     /** Used in tests when no UI is required. */
@@ -108,9 +142,10 @@
     Bubble(NotificationEntry e,
             BubbleController.NotificationSuppressionChangedListener listener) {
         mEntry = e;
-        mKey = e.getKey();
+        mKey = key(e);
         mLastUpdated = e.getSbn().getPostTime();
         mSuppressionListener = listener;
+        mFlags = e.getSbn().getNotification().flags;
     }
 
     @Override
@@ -118,12 +153,22 @@
         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 getPackageName() {
-        return mEntry.getSbn().getPackageName();
+        return mEntry == null
+                ? mShortcutInfo == null ? null : mShortcutInfo.getPackage()
+                : mEntry.getSbn().getPackageName();
     }
 
     @Override
@@ -167,6 +212,18 @@
         return mExpandedView;
     }
 
+    @Nullable
+    public String getTitle() {
+        final CharSequence titleCharSeq;
+        if (mEntry == null) {
+            titleCharSeq = null;
+        } else {
+            titleCharSeq = mEntry.getSbn().getNotification().extras.getCharSequence(
+                    Notification.EXTRA_TITLE);
+        }
+        return titleCharSeq != null ? titleCharSeq.toString() : null;
+    }
+
     /**
      * Call when the views should be removed, ensure this is called to clean up ActivityView
      * content.
@@ -207,7 +264,8 @@
     void inflate(BubbleViewInfoTask.Callback callback,
             Context context,
             BubbleStackView stackView,
-            BubbleIconFactory iconFactory) {
+            BubbleIconFactory iconFactory,
+            boolean skipInflation) {
         if (isBubbleLoading()) {
             mInflationTask.cancel(true /* mayInterruptIfRunning */);
         }
@@ -215,6 +273,7 @@
                 context,
                 stackView,
                 iconFactory,
+                skipInflation,
                 callback);
         if (mInflateSynchronously) {
             mInflationTask.onPostExecute(mInflationTask.doInBackground());
@@ -327,6 +386,7 @@
      * Whether this notification should be shown in the shade.
      */
     boolean showInShade() {
+        if (mEntry == null) return false;
         return !shouldSuppressNotification() || !mEntry.isClearable();
     }
 
@@ -334,8 +394,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) {
@@ -366,6 +426,7 @@
      */
     @Override
     public boolean showDot() {
+        if (mEntry == null) return false;
         return mShowBubbleUpdateDot
                 && !mEntry.shouldSuppressNotificationDot()
                 && !shouldSuppressNotification();
@@ -375,6 +436,7 @@
      * Whether the flyout for the bubble should be shown.
      */
     boolean showFlyout() {
+        if (mEntry == null) return false;
         return !mSuppressFlyout && !mEntry.shouldSuppressPeek()
                 && !shouldSuppressNotification()
                 && !mEntry.shouldSuppressNotificationList();
@@ -394,6 +456,7 @@
     }
 
     float getDesiredHeight(Context context) {
+        if (mEntry == null) return 0;
         Notification.BubbleMetadata data = mEntry.getBubbleMetadata();
         boolean useRes = data.getDesiredHeightResId() != 0;
         if (useRes) {
@@ -407,6 +470,7 @@
     }
 
     String getDesiredHeightString() {
+        if (mEntry == null) return String.valueOf(0);
         Notification.BubbleMetadata data = mEntry.getBubbleMetadata();
         boolean useRes = data.getDesiredHeightResId() != 0;
         if (useRes) {
@@ -423,11 +487,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();
@@ -435,16 +501,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;
@@ -466,11 +548,13 @@
     }
 
     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()) ||  mShouldAutoExpand;
     }
@@ -479,6 +563,19 @@
         mShouldAutoExpand = shouldAutoExpand;
     }
 
+    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 f0f28fd..0ec2418 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;
@@ -113,6 +114,7 @@
 import java.lang.annotation.Target;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Objects;
 
 /**
  * Bubbles are a special type of content that can "float" on top of other apps or System UI.
@@ -243,13 +245,13 @@
          * 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
          * filtered from the shade.
          */
-        void invalidateNotifications(String reason);
+        void invalidateNotifications(@NonNull String reason);
 
         /**
          * Called on a bubbled entry that has been removed when there are no longer
@@ -259,7 +261,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);
     }
 
     /**
@@ -604,6 +606,8 @@
 
             mStackView.setUnbubbleConversationCallback(notificationEntry ->
                     onUserChangedBubble(notificationEntry, false /* shouldBubble */));
+            // Lazy load overflow bubbles from disk
+            loadOverflowBubblesFromDisk();
         }
 
         addToWindowManagerMaybe();
@@ -755,10 +759,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 */);
         }
     }
 
@@ -845,7 +851,7 @@
 
     void promoteBubbleFromOverflow(Bubble bubble) {
         bubble.setInflateSynchronously(mInflateSynchronously);
-        setIsBubble(bubble.getEntry(), /* isBubble */ true);
+        setIsBubble(bubble, /* isBubble */ true);
         mBubbleData.promoteBubbleFromOverflow(bubble, mStackView, mBubbleIconFactory);
     }
 
@@ -895,10 +901,30 @@
         updateBubble(notif, false /* 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) {
         // Lazy init stack view when a bubble is created
         ensureStackViewCreated();
-
         // If this is an interruptive notif, mark that it's interrupted
         if (notif.getImportance() >= NotificationManager.IMPORTANCE_HIGH) {
             notif.setInterruption();
@@ -918,11 +944,11 @@
                             return;
                         }
                         mHandler.post(
-                                () -> removeBubble(bubble.getEntry(),
+                                () -> removeBubble(bubble.getKey(),
                                         BubbleController.DISMISS_INVALID_INTENT));
                     });
                 },
-                mContext, mStackView, mBubbleIconFactory);
+                mContext, mStackView, mBubbleIconFactory, false /* skipInflation */);
     }
 
     /**
@@ -934,7 +960,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();
@@ -973,14 +1002,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);
         }
     }
 
@@ -998,7 +1027,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 && entry.isBubble()) {
             updateBubble(entry);
         }
@@ -1012,10 +1041,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);
         }
     }
 
@@ -1037,7 +1066,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);
@@ -1045,7 +1075,8 @@
         }
     }
 
-    private void setIsBubble(NotificationEntry entry, boolean isBubble) {
+    private void setIsBubble(@NonNull final NotificationEntry entry, final boolean isBubble) {
+        Objects.requireNonNull(entry);
         if (isBubble) {
             entry.getSbn().getNotification().flags |= FLAG_BUBBLE;
         } else {
@@ -1058,6 +1089,24 @@
         }
     }
 
+    private void setIsBubble(@NonNull final Bubble b, final boolean isBubble) {
+        Objects.requireNonNull(b);
+        if (isBubble) {
+            b.enable(FLAG_BUBBLE);
+        } else {
+            b.disable(FLAG_BUBBLE);
+        }
+        if (b.getEntry() != null) {
+            setIsBubble(b.getEntry(), isBubble);
+        } else {
+            try {
+                mBarService.onNotificationBubbleChanged(b.getKey(), isBubble, 0);
+            } catch (RemoteException e) {
+                // Bad things have happened
+            }
+        }
+    }
+
     @SuppressWarnings("FieldCanBeLocal")
     private final BubbleData.Listener mBubbleDataListener = new BubbleData.Listener() {
 
@@ -1097,23 +1146,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()) {
-                            setIsBubble(bubble.getEntry(), false /* isBubble */);
+                        if (bubble.isBubble() && bubble.showInShade()) {
+                            setIsBubble(bubble, false /* isBubble */);
                         }
-                        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());
+                        }
                     }
                 }
             }
@@ -1138,7 +1191,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 35647b0..857f1dd 100644
--- a/packages/SystemUI/src/com/android/systemui/bubbles/BubbleData.java
+++ b/packages/SystemUI/src/com/android/systemui/bubbles/BubbleData.java
@@ -21,6 +21,7 @@
 import static com.android.systemui.bubbles.BubbleDebugConfig.TAG_BUBBLES;
 import static com.android.systemui.bubbles.BubbleDebugConfig.TAG_WITH_CLASS_NAME;
 
+import android.annotation.NonNull;
 import android.app.Notification;
 import android.app.PendingIntent;
 import android.content.Context;
@@ -214,7 +215,7 @@
                     notificationEntryUpdated(bubble, false /* suppressFlyout */,
                             true /* showInShade */);
                 },
-                mContext, stack, factory);
+                mContext, stack, factory, false /* skipInflation */);
     }
 
     void setShowingOverflow(boolean showingOverflow) {
@@ -268,7 +269,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
@@ -297,11 +299,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();
     }
 
@@ -349,7 +354,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);
             }
         }
@@ -447,7 +452,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) {
@@ -615,7 +622,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..a0ec7cd 100644
--- a/packages/SystemUI/src/com/android/systemui/bubbles/BubbleDataRepository.kt
+++ b/packages/SystemUI/src/com/android/systemui/bubbles/BubbleDataRepository.kt
@@ -74,7 +74,9 @@
 
     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) return@mapNotNull null
             BubbleEntity(userId, b.packageName, shortcutId)
         }
     }
@@ -108,7 +110,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 c4b4f43..494a51d 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));
             }
         }
@@ -297,10 +296,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/BubbleOverflowActivity.java b/packages/SystemUI/src/com/android/systemui/bubbles/BubbleOverflowActivity.java
index 8fec338..2109a7b 100644
--- a/packages/SystemUI/src/com/android/systemui/bubbles/BubbleOverflowActivity.java
+++ b/packages/SystemUI/src/com/android/systemui/bubbles/BubbleOverflowActivity.java
@@ -21,7 +21,6 @@
 import static com.android.systemui.bubbles.BubbleDebugConfig.TAG_WITH_CLASS_NAME;
 
 import android.app.Activity;
-import android.app.Notification;
 import android.content.Context;
 import android.content.res.Configuration;
 import android.content.res.Resources;
@@ -260,12 +259,9 @@
             mPromoteBubbleFromOverflow.accept(b);
         });
 
-        final CharSequence titleCharSeq =
-                b.getEntry().getSbn().getNotification().extras.getCharSequence(
-                        Notification.EXTRA_TITLE);
-        String titleStr = mContext.getResources().getString(R.string.notification_bubble_title);
-        if (titleCharSeq != null) {
-            titleStr = titleCharSeq.toString();
+        String titleStr = b.getTitle();
+        if (titleStr == null) {
+            titleStr = mContext.getResources().getString(R.string.notification_bubble_title);
         }
         vh.iconView.setContentDescription(mContext.getResources().getString(
                 R.string.bubble_content_description_single, titleStr, b.getAppName()));
diff --git a/packages/SystemUI/src/com/android/systemui/bubbles/BubbleStackView.java b/packages/SystemUI/src/com/android/systemui/bubbles/BubbleStackView.java
index 418cc50..6ba1aa8 100644
--- a/packages/SystemUI/src/com/android/systemui/bubbles/BubbleStackView.java
+++ b/packages/SystemUI/src/com/android/systemui/bubbles/BubbleStackView.java
@@ -31,7 +31,6 @@
 import android.animation.AnimatorListenerAdapter;
 import android.animation.ValueAnimator;
 import android.annotation.SuppressLint;
-import android.app.Notification;
 import android.content.ContentResolver;
 import android.content.Context;
 import android.content.Intent;
@@ -50,7 +49,6 @@
 import android.graphics.Region;
 import android.os.Bundle;
 import android.provider.Settings;
-import android.service.notification.StatusBarNotification;
 import android.util.Log;
 import android.view.Choreographer;
 import android.view.DisplayCutout;
@@ -941,10 +939,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);
@@ -1202,13 +1200,10 @@
         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);
-            if (titleCharSeq != null) {
-                titleStr = titleCharSeq.toString();
+            String titleStr = bubble.getTitle();
+            if (titleStr == null) {
+                titleStr = getResources().getString(R.string.notification_bubble_title);
             }
 
             if (bubble.getIconView() != null) {
@@ -1821,7 +1816,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);
         }
     }
 
@@ -2319,18 +2314,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(),
-                false /* isOngoing (unused) */,
-                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/tests/src/com/android/systemui/bubbles/BubbleControllerTest.java b/packages/SystemUI/tests/src/com/android/systemui/bubbles/BubbleControllerTest.java
index 9b377ca..c89f6c2 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/bubbles/BubbleControllerTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/bubbles/BubbleControllerTest.java
@@ -319,7 +319,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());
 
@@ -331,7 +331,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));
@@ -352,9 +352,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();
@@ -367,7 +368,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());
@@ -565,7 +566,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());
 
@@ -576,7 +578,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
@@ -702,7 +705,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();
     }
 
@@ -710,7 +713,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();
     }
 
@@ -813,7 +816,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(
@@ -834,7 +837,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(
@@ -856,12 +859,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 eca78ec..8224c88e 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/bubbles/BubbleDataTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/bubbles/BubbleDataTest.java
@@ -177,7 +177,8 @@
         mBubbleData.setListener(mListener);
 
         // Test
-        mBubbleData.notificationEntryRemoved(mEntryA1, BubbleController.DISMISS_USER_GESTURE);
+        mBubbleData.notificationEntryRemoved(
+                mEntryA1.getKey(), BubbleController.DISMISS_USER_GESTURE);
 
         // Verify
         verifyUpdateReceived();
@@ -299,12 +300,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));
     }
@@ -325,12 +328,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());
     }
@@ -410,7 +415,8 @@
         mBubbleData.setListener(mListener);
 
         // Test
-        mBubbleData.notificationEntryRemoved(mEntryA2, BubbleController.DISMISS_USER_GESTURE);
+        mBubbleData.notificationEntryRemoved(
+                mEntryA2.getKey(), BubbleController.DISMISS_USER_GESTURE);
         verifyUpdateReceived();
         // TODO: this should fail if things work as I expect them to?
         assertOrderChangedTo(mBubbleB2, mBubbleB1, mBubbleA1);
@@ -430,7 +436,8 @@
         mBubbleData.setListener(mListener);
 
         // Test
-        mBubbleData.notificationEntryRemoved(mEntryA1, BubbleController.DISMISS_USER_GESTURE);
+        mBubbleData.notificationEntryRemoved(
+                mEntryA1.getKey(), BubbleController.DISMISS_USER_GESTURE);
         verifyUpdateReceived();
         assertOrderNotChanged();
     }
@@ -449,7 +456,8 @@
         mBubbleData.setListener(mListener);
 
         // Test
-        mBubbleData.notificationEntryRemoved(mEntryA2, BubbleController.DISMISS_NOTIF_CANCEL);
+        mBubbleData.notificationEntryRemoved(
+                mEntryA2.getKey(), BubbleController.DISMISS_NOTIF_CANCEL);
         verifyUpdateReceived();
         assertSelectionChangedTo(mBubbleB2);
     }
@@ -523,7 +531,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();
@@ -623,7 +632,8 @@
         mBubbleData.setListener(mListener);
 
         // Test
-        mBubbleData.notificationEntryRemoved(mEntryB2, BubbleController.DISMISS_USER_GESTURE);
+        mBubbleData.notificationEntryRemoved(
+                mEntryB2.getKey(), BubbleController.DISMISS_USER_GESTURE);
         verifyUpdateReceived();
         assertOrderChangedTo(mBubbleA2, mBubbleB1, mBubbleA1);
     }
@@ -647,11 +657,13 @@
         mBubbleData.setListener(mListener);
 
         // Test
-        mBubbleData.notificationEntryRemoved(mEntryA2, BubbleController.DISMISS_USER_GESTURE);
+        mBubbleData.notificationEntryRemoved(
+                mEntryA2.getKey(), BubbleController.DISMISS_USER_GESTURE);
         verifyUpdateReceived();
         assertSelectionChangedTo(mBubbleB1);
 
-        mBubbleData.notificationEntryRemoved(mEntryB1, BubbleController.DISMISS_USER_GESTURE);
+        mBubbleData.notificationEntryRemoved(
+                mEntryB1.getKey(), BubbleController.DISMISS_USER_GESTURE);
         verifyUpdateReceived();
         assertSelectionChangedTo(mBubbleA1);
     }
@@ -765,7 +777,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 b18d67b..ead95ca1 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/bubbles/NewNotifPipelineBubbleControllerTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/bubbles/NewNotifPipelineBubbleControllerTest.java
@@ -287,7 +287,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());
     }
@@ -304,7 +305,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.
@@ -324,7 +326,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,
@@ -504,7 +507,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());
 
@@ -515,7 +519,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
@@ -615,7 +620,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();
     }
 
@@ -623,7 +628,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();
     }
 
@@ -691,7 +696,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
@@ -712,7 +717,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