Revert "Persists bubbles to disk (part 4)"

This reverts commit 293b83d81ee2b537d2a8b5a1f81fce6270bd0d0c.

Reason for revert: May have caused http://b/156697897
Bug: 156697897
Change-Id: I9ce37a54abd25cc6a6dde2c7728d2da1076654c5
diff --git a/packages/SystemUI/src/com/android/systemui/bubbles/Bubble.java b/packages/SystemUI/src/com/android/systemui/bubbles/Bubble.java
index f044389..f1b401e 100644
--- a/packages/SystemUI/src/com/android/systemui/bubbles/Bubble.java
+++ b/packages/SystemUI/src/com/android/systemui/bubbles/Bubble.java
@@ -16,7 +16,6 @@
 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;
 
@@ -28,7 +27,6 @@
 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;
@@ -57,11 +55,6 @@
 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;
@@ -102,56 +95,18 @@
     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();
     }
 
-    /**
-     * 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.
-     */
+    // TODO: Decouple Bubble from NotificationEntry and transform ShortcutInfo into Bubble
     Bubble(ShortcutInfo shortcutInfo) {
         mShortcutInfo = shortcutInfo;
-        mKey = key(shortcutInfo);
-        mGroupId = groupId(shortcutInfo);
-        mFlags = 0;
+        mKey = shortcutInfo.getId();
+        mGroupId = shortcutInfo.getId();
     }
 
     /** Used in tests when no UI is required. */
@@ -159,11 +114,10 @@
     Bubble(NotificationEntry e,
             BubbleController.NotificationSuppressionChangedListener listener) {
         mEntry = e;
-        mKey = key(e);
+        mKey = e.getKey();
         mLastUpdated = e.getSbn().getPostTime();
         mGroupId = groupId(e);
         mSuppressionListener = listener;
-        mFlags = e.getSbn().getNotification().flags;
     }
 
     @Override
@@ -171,26 +125,16 @@
         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 == null
-                ? mShortcutInfo == null ? null : mShortcutInfo.getPackage()
-                : mEntry.getSbn().getPackageName();
+        return mEntry.getSbn().getPackageName();
     }
 
     @Override
@@ -274,8 +218,7 @@
     void inflate(BubbleViewInfoTask.Callback callback,
             Context context,
             BubbleStackView stackView,
-            BubbleIconFactory iconFactory,
-            boolean skipInflation) {
+            BubbleIconFactory iconFactory) {
         if (isBubbleLoading()) {
             mInflationTask.cancel(true /* mayInterruptIfRunning */);
         }
@@ -283,7 +226,6 @@
                 context,
                 stackView,
                 iconFactory,
-                skipInflation,
                 callback);
         if (mInflateSynchronously) {
             mInflationTask.onPostExecute(mInflationTask.doInBackground());
@@ -403,7 +345,6 @@
      * Whether this notification should be shown in the shade.
      */
     boolean showInShade() {
-        if (mEntry == null) return false;
         return !shouldSuppressNotification() || !mEntry.isClearable();
     }
 
@@ -411,8 +352,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) {
@@ -443,7 +384,6 @@
      */
     @Override
     public boolean showDot() {
-        if (mEntry == null) return false;
         return mShowBubbleUpdateDot
                 && !mEntry.shouldSuppressNotificationDot()
                 && !shouldSuppressNotification();
@@ -453,7 +393,6 @@
      * Whether the flyout for the bubble should be shown.
      */
     boolean showFlyout() {
-        if (mEntry == null) return false;
         return !mSuppressFlyout && !mEntry.shouldSuppressPeek()
                 && !shouldSuppressNotification()
                 && !mEntry.shouldSuppressNotificationList();
@@ -477,13 +416,11 @@
      * 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) {
@@ -497,7 +434,6 @@
     }
 
     String getDesiredHeightString() {
-        if (mEntry == null) return String.valueOf(0);
         Notification.BubbleMetadata data = mEntry.getBubbleMetadata();
         boolean useRes = data.getDesiredHeightResId() != 0;
         if (useRes) {
@@ -514,13 +450,11 @@
      * To populate the icon use {@link LauncherApps#getShortcutIconDrawable(ShortcutInfo, int)}.
      */
     boolean usingShortcutInfo() {
-        return mEntry != null && mEntry.getBubbleMetadata().getShortcutId() != null
-                || mShortcutInfo != null;
+        return mEntry.getBubbleMetadata().getShortcutId() != null;
     }
 
     @Nullable
     PendingIntent getBubbleIntent() {
-        if (mEntry == null) return null;
         Notification.BubbleMetadata data = mEntry.getBubbleMetadata();
         if (data != null) {
             return data.getIntent();
@@ -528,32 +462,16 @@
         return null;
     }
 
-    Intent getSettingsIntent(final Context context) {
+    Intent getSettingsIntent() {
         final Intent intent = new Intent(Settings.ACTION_APP_NOTIFICATION_BUBBLE_SETTINGS);
         intent.putExtra(Settings.EXTRA_APP_PACKAGE, getPackageName());
-        final int uid = getUid(context);
-        if (uid != -1) {
-            intent.putExtra(Settings.EXTRA_APP_UID, uid);
-        }
+        intent.putExtra(Settings.EXTRA_APP_UID, mEntry.getSbn().getUid());
         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;
@@ -575,30 +493,15 @@
     }
 
     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 ec1216d..6a5e94c 100644
--- a/packages/SystemUI/src/com/android/systemui/bubbles/BubbleController.java
+++ b/packages/SystemUI/src/com/android/systemui/bubbles/BubbleController.java
@@ -42,7 +42,6 @@
 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;
@@ -238,7 +237,7 @@
          * This can happen when an app cancels a bubbled notification or when the user dismisses a
          * bubble.
          */
-        void removeNotification(@NonNull NotificationEntry entry, int reason);
+        void removeNotification(NotificationEntry entry, int reason);
 
         /**
          * Called when a bubbled notification has changed whether it should be
@@ -254,7 +253,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(@NonNull NotificationEntry entry);
+        void maybeCancelSummary(NotificationEntry entry);
     }
 
     /**
@@ -477,7 +476,7 @@
 
         addNotifCallback(new NotifCallback() {
             @Override
-            public void removeNotification(@NonNull final NotificationEntry entry, int reason) {
+            public void removeNotification(NotificationEntry entry, int reason) {
                 mNotificationEntryManager.performRemoveNotification(entry.getSbn(),
                         reason);
             }
@@ -488,7 +487,7 @@
             }
 
             @Override
-            public void maybeCancelSummary(@NonNull final NotificationEntry entry) {
+            public void maybeCancelSummary(NotificationEntry entry) {
                 // Check if removed bubble has an associated suppressed group summary that needs
                 // to be removed now.
                 final String groupKey = entry.getSbn().getGroupKey();
@@ -697,12 +696,10 @@
         mBubbleIconFactory = new BubbleIconFactory(mContext);
         // Reload each bubble
         for (Bubble b: mBubbleData.getBubbles()) {
-            b.inflate(null /* callback */, mContext, mStackView, mBubbleIconFactory,
-                    false /* skipInflation */);
+            b.inflate(null /* callback */, mContext, mStackView, mBubbleIconFactory);
         }
         for (Bubble b: mBubbleData.getOverflowBubbles()) {
-            b.inflate(null /* callback */, mContext, mStackView, mBubbleIconFactory,
-                    false /* skipInflation */);
+            b.inflate(null /* callback */, mContext, mStackView, mBubbleIconFactory);
         }
     }
 
@@ -794,7 +791,7 @@
             if (bubble != null) {
                 mBubbleData.promoteBubbleFromOverflow(bubble, mStackView, mBubbleIconFactory);
             }
-        } else if (bubble.isBubble()) {
+        } else if (bubble.getEntry().isBubble()){
             mBubbleData.setSelectedBubble(bubble);
         }
         mBubbleData.setExpanded(true);
@@ -823,33 +820,10 @@
         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) {
@@ -869,11 +843,11 @@
                             return;
                         }
                         mHandler.post(
-                                () -> removeBubble(bubble.getKey(),
+                                () -> removeBubble(bubble.getEntry(),
                                         BubbleController.DISMISS_INVALID_INTENT));
                     });
                 },
-                mContext, mStackView, mBubbleIconFactory, false /* skipInflation */);
+                mContext, mStackView, mBubbleIconFactory);
     }
 
     /**
@@ -885,10 +859,7 @@
      * @param entry the notification to change bubble state for.
      * @param shouldBubble whether the notification should show as a bubble or not.
      */
-    public void onUserChangedBubble(@Nullable final NotificationEntry entry, boolean shouldBubble) {
-        if (entry == null) {
-            return;
-        }
+    public void onUserChangedBubble(NotificationEntry entry, boolean shouldBubble) {
         NotificationChannel channel = entry.getChannel();
         final String appPkg = entry.getSbn().getPackageName();
         final int appUid = entry.getSbn().getUid();
@@ -926,14 +897,14 @@
     }
 
     /**
-     * Removes the bubble with the given key.
+     * Removes the bubble with the given NotificationEntry.
      * <p>
      * Must be called from the main thread.
      */
     @MainThread
-    void removeBubble(String key, int reason) {
-        if (mBubbleData.hasAnyBubbleWithKey(key)) {
-            mBubbleData.notificationEntryRemoved(key, reason);
+    void removeBubble(NotificationEntry entry, int reason) {
+        if (mBubbleData.hasAnyBubbleWithKey(entry.getKey())) {
+            mBubbleData.notificationEntryRemoved(entry, reason);
         }
     }
 
@@ -949,7 +920,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.getKey(), DISMISS_NO_LONGER_BUBBLE);
+            removeBubble(entry, DISMISS_NO_LONGER_BUBBLE);
         } else if (shouldBubble) {
             updateBubble(entry);
         }
@@ -963,10 +934,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).getKey(), DISMISS_GROUP_CANCELLED);
+                removeBubble(bubbleChildren.get(i).getEntry(), DISMISS_GROUP_CANCELLED);
             }
         } else {
-            removeBubble(entry.getKey(), DISMISS_NOTIF_CANCEL);
+            removeBubble(entry, DISMISS_NOTIF_CANCEL);
         }
     }
 
@@ -988,8 +959,7 @@
             rankingMap.getRanking(key, mTmpRanking);
             boolean isActiveBubble = mBubbleData.hasAnyBubbleWithKey(key);
             if (isActiveBubble && !mTmpRanking.canBubble()) {
-                mBubbleData.notificationEntryRemoved(entry.getKey(),
-                        BubbleController.DISMISS_BLOCKED);
+                mBubbleData.notificationEntryRemoved(entry, BubbleController.DISMISS_BLOCKED);
             } else if (entry != null && mTmpRanking.isBubble() && !isActiveBubble) {
                 entry.setFlagBubble(true);
                 onEntryUpdated(entry);
@@ -999,15 +969,9 @@
 
     private void setIsBubble(Bubble b, boolean isBubble) {
         if (isBubble) {
-            if (b.getEntry() != null) {
-                b.getEntry().getSbn().getNotification().flags |= FLAG_BUBBLE;
-            }
-            b.enable(FLAG_BUBBLE);
+            b.getEntry().getSbn().getNotification().flags |= FLAG_BUBBLE;
         } else {
-            if (b.getEntry() != null) {
-                b.getEntry().getSbn().getNotification().flags &= ~FLAG_BUBBLE;
-            }
-            b.disable(FLAG_BUBBLE);
+            b.getEntry().getSbn().getNotification().flags &= ~FLAG_BUBBLE;
         }
         try {
             mBarService.onNotificationBubbleChanged(b.getKey(), isBubble, 0);
@@ -1055,27 +1019,23 @@
                         // The bubble is now gone & the notification is hidden from the shade, so
                         // time to actually remove it
                         for (NotifCallback cb : mCallbacks) {
-                            if (bubble.getEntry() != null) {
-                                cb.removeNotification(bubble.getEntry(), REASON_CANCEL);
-                            }
+                            cb.removeNotification(bubble.getEntry(), REASON_CANCEL);
                         }
                     } else {
-                        if (bubble.isBubble() && bubble.showInShade()) {
+                        if (bubble.getEntry().isBubble() && bubble.showInShade()) {
                             setIsBubble(bubble, /* isBubble */ false);
                         }
-                        if (bubble.getEntry() != null && bubble.getEntry().getRow() != null) {
+                        if (bubble.getEntry().getRow() != null) {
                             bubble.getEntry().getRow().updateBubbleButton();
                         }
                     }
 
                 }
-                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());
-                        }
+                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());
                     }
                 }
             }
@@ -1100,7 +1060,7 @@
 
             if (update.selectionChanged) {
                 mStackView.setSelectedBubble(update.selectedBubble);
-                if (update.selectedBubble != null && update.selectedBubble.getEntry() != null) {
+                if (update.selectedBubble != 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 7b323ce..65d5beb 100644
--- a/packages/SystemUI/src/com/android/systemui/bubbles/BubbleData.java
+++ b/packages/SystemUI/src/com/android/systemui/bubbles/BubbleData.java
@@ -23,7 +23,6 @@
 
 import static java.util.stream.Collectors.toList;
 
-import android.annotation.NonNull;
 import android.app.Notification;
 import android.app.PendingIntent;
 import android.content.Context;
@@ -224,7 +223,7 @@
                             false, /* showInShade */ true);
                     setSelectedBubble(bubble);
                 },
-                mContext, stack, factory, false /* skipInflation */);
+                mContext, stack, factory);
         dispatchPendingChanges();
     }
 
@@ -279,8 +278,7 @@
         }
         mPendingBubbles.remove(bubble); // No longer pending once we're here
         Bubble prevBubble = getBubbleInStackWithKey(bubble.getKey());
-        suppressFlyout |= bubble.getEntry() == null
-                || !bubble.getEntry().getRanking().visuallyInterruptive();
+        suppressFlyout |= !bubble.getEntry().getRanking().visuallyInterruptive();
 
         if (prevBubble == null) {
             // Create a new bubble
@@ -309,14 +307,11 @@
         dispatchPendingChanges();
     }
 
-    /**
-     * Called when a notification associated with a bubble is removed.
-     */
-    public void notificationEntryRemoved(String key, @DismissReason int reason) {
+    public void notificationEntryRemoved(NotificationEntry entry, @DismissReason int reason) {
         if (DEBUG_BUBBLE_DATA) {
-            Log.d(TAG, "notificationEntryRemoved: key=" + key + " reason=" + reason);
+            Log.d(TAG, "notificationEntryRemoved: entry=" + entry + " reason=" + reason);
         }
-        doRemove(key, reason);
+        doRemove(entry.getKey(), reason);
         dispatchPendingChanges();
     }
 
@@ -364,7 +359,7 @@
             return bubbleChildren;
         }
         for (Bubble b : mBubbles) {
-            if (b.getEntry() != null && groupKey.equals(b.getEntry().getSbn().getGroupKey())) {
+            if (groupKey.equals(b.getEntry().getSbn().getGroupKey())) {
                 bubbleChildren.add(b);
             }
         }
@@ -475,9 +470,7 @@
             Bubble newSelected = mBubbles.get(newIndex);
             setSelectedBubbleInternal(newSelected);
         }
-        if (bubbleToRemove.getEntry() != null) {
-            maybeSendDeleteIntent(reason, bubbleToRemove.getEntry());
-        }
+        maybeSendDeleteIntent(reason, bubbleToRemove.getEntry());
     }
 
     void overflowBubble(@DismissReason int reason, Bubble bubble) {
@@ -751,8 +744,7 @@
         return true;
     }
 
-    private void maybeSendDeleteIntent(@DismissReason int reason,
-            @NonNull final NotificationEntry entry) {
+    private void maybeSendDeleteIntent(@DismissReason int reason, 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 6df3132..ba93f41 100644
--- a/packages/SystemUI/src/com/android/systemui/bubbles/BubbleDataRepository.kt
+++ b/packages/SystemUI/src/com/android/systemui/bubbles/BubbleDataRepository.kt
@@ -74,10 +74,7 @@
 
     private fun transform(userId: Int, bubbles: List<Bubble>): List<BubbleEntity> {
         return bubbles.mapNotNull { b ->
-            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
+            val shortcutId = b.shortcutInfo?.id ?: return@mapNotNull null
             BubbleEntity(userId, b.packageName, shortcutId)
         }
     }
@@ -111,6 +108,7 @@
     /**
      * 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 0a1a0f7..baf92dc 100644
--- a/packages/SystemUI/src/com/android/systemui/bubbles/BubbleExpandedView.java
+++ b/packages/SystemUI/src/com/android/systemui/bubbles/BubbleExpandedView.java
@@ -65,6 +65,7 @@
 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.
@@ -160,7 +161,7 @@
                             // the bubble again so we'll just remove it.
                             Log.w(TAG, "Exception while displaying bubble: " + getBubbleKey()
                                     + ", " + e.getMessage() + "; removing bubble");
-                            mBubbleController.removeBubble(getBubbleKey(),
+                            mBubbleController.removeBubble(getBubbleEntry(),
                                     BubbleController.DISMISS_INVALID_INTENT);
                         }
                     });
@@ -204,7 +205,7 @@
             }
             if (mBubble != null) {
                 // Must post because this is called from a binder thread.
-                post(() -> mBubbleController.removeBubble(mBubble.getKey(),
+                post(() -> mBubbleController.removeBubble(mBubble.getEntry(),
                         BubbleController.DISMISS_TASK_FINISHED));
             }
         }
@@ -291,6 +292,10 @@
         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 c06063b..2cb097f 100644
--- a/packages/SystemUI/src/com/android/systemui/bubbles/BubbleStackView.java
+++ b/packages/SystemUI/src/com/android/systemui/bubbles/BubbleStackView.java
@@ -49,6 +49,7 @@
 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;
@@ -915,10 +916,10 @@
                     showManageMenu(false /* show */);
                     final Bubble bubble = mBubbleData.getSelectedBubble();
                     if (bubble != null && mBubbleData.hasBubbleInStackWithKey(bubble.getKey())) {
-                        final Intent intent = bubble.getSettingsIntent(mContext);
+                        final Intent intent = bubble.getSettingsIntent();
                         collapseStack(() -> {
-
-                            mContext.startActivityAsUser(intent, bubble.getUser());
+                            mContext.startActivityAsUser(
+                                    intent, bubble.getEntry().getSbn().getUser());
                             logBubbleClickEvent(
                                     bubble,
                                     SysUiStatsLog.BUBBLE_UICHANGED__ACTION__HEADER_GO_TO_SETTINGS);
@@ -1152,19 +1153,13 @@
         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);
 
-            final CharSequence titleCharSeq;
-            if (bubble.getEntry() == null) {
-                titleCharSeq = null;
-            } else {
-                titleCharSeq = bubble.getEntry().getSbn().getNotification().extras.getCharSequence(
-                        Notification.EXTRA_TITLE);
-            }
-            final String titleStr;
+            String titleStr = getResources().getString(R.string.notification_bubble_title);
             if (titleCharSeq != null) {
                 titleStr = titleCharSeq.toString();
-            } else {
-                titleStr = getResources().getString(R.string.notification_bubble_title);
             }
 
             if (bubble.getIconView() != null) {
@@ -1782,7 +1777,7 @@
     private void dismissBubbleIfExists(@Nullable Bubble bubble) {
         if (bubble != null && mBubbleData.hasBubbleInStackWithKey(bubble.getKey())) {
             mBubbleData.notificationEntryRemoved(
-                    bubble.getKey(), BubbleController.DISMISS_USER_GESTURE);
+                    bubble.getEntry(), BubbleController.DISMISS_USER_GESTURE);
         }
     }
 
@@ -2280,12 +2275,18 @@
      * @param action the user interaction enum.
      */
     private void logBubbleClickEvent(Bubble bubble, int action) {
-        bubble.logUIEvent(
+        StatusBarNotification notification = bubble.getEntry().getSbn();
+        SysUiStatsLog.write(SysUiStatsLog.BUBBLE_UI_CHANGED,
+                notification.getPackageName(),
+                notification.getNotification().getChannelId(),
+                notification.getId(),
+                getBubbleIndex(getExpandedBubble()),
                 getBubbleCount(),
                 action,
                 getNormalizedXPosition(),
                 getNormalizedYPosition(),
-                getBubbleIndex(getExpandedBubble())
-        );
+                bubble.showInShade(),
+                bubble.isOngoing(),
+                false /* isAppForeground (unused) */);
     }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/bubbles/BubbleViewInfoTask.java b/packages/SystemUI/src/com/android/systemui/bubbles/BubbleViewInfoTask.java
index 525d5b5..8a57a73 100644
--- a/packages/SystemUI/src/com/android/systemui/bubbles/BubbleViewInfoTask.java
+++ b/packages/SystemUI/src/com/android/systemui/bubbles/BubbleViewInfoTask.java
@@ -37,7 +37,6 @@
 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;
@@ -75,7 +74,6 @@
     private WeakReference<Context> mContext;
     private WeakReference<BubbleStackView> mStackView;
     private BubbleIconFactory mIconFactory;
-    private boolean mSkipInflation;
     private Callback mCallback;
 
     /**
@@ -86,20 +84,17 @@
             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,
-                mSkipInflation);
+        return BubbleViewInfo.populate(mContext.get(), mStackView.get(), mIconFactory, mBubble);
     }
 
     @Override
@@ -128,36 +123,11 @@
 
         @Nullable
         static BubbleViewInfo populate(Context c, BubbleStackView stackView,
-                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) {
+                BubbleIconFactory iconFactory, Bubble b) {
             BubbleViewInfo info = new BubbleViewInfo();
 
             // View inflation: only should do this once per bubble
-            if (!isInflated) {
+            if (!b.isInflated()) {
                 LayoutInflater inflater = LayoutInflater.from(c);
                 info.imageView = (BadgedImageView) inflater.inflate(
                         R.layout.bubble_view, stackView, false /* attachToRoot */);
@@ -167,8 +137,12 @@
                 info.expandedView.setStackView(stackView);
             }
 
-            if (shortcutInfo != null) {
-                info.shortcutInfo = shortcutInfo;
+            StatusBarNotification sbn = b.getEntry().getSbn();
+            String packageName = sbn.getPackageName();
+
+            String bubbleShortcutId =  b.getEntry().getBubbleMetadata().getShortcutId();
+            if (bubbleShortcutId != null) {
+                info.shortcutInfo = b.getEntry().getRanking().getShortcutInfo();
             }
 
             // App name & app icon
@@ -187,7 +161,7 @@
                     info.appName = String.valueOf(pm.getApplicationLabel(appInfo));
                 }
                 appIcon = pm.getApplicationIcon(packageName);
-                badgedIcon = pm.getUserBadgedIcon(appIcon, user);
+                badgedIcon = pm.getUserBadgedIcon(appIcon, sbn.getUser());
             } 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);
@@ -196,7 +170,7 @@
 
             // Badged bubble image
             Drawable bubbleDrawable = iconFactory.getBubbleDrawable(c, info.shortcutInfo,
-                    entry == null ? null : entry.getBubbleMetadata());
+                    b.getEntry().getBubbleMetadata());
             if (bubbleDrawable == null) {
                 // Default to app icon
                 bubbleDrawable = appIcon;
@@ -222,9 +196,7 @@
                     Color.WHITE, WHITE_SCRIM_ALPHA);
 
             // Flyout
-            if (entry != null) {
-                info.flyoutMessage = extractFlyoutMessage(c, entry);
-            }
+            info.flyoutMessage = extractFlyoutMessage(c, b.getEntry());
             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 57be1a5..92426e5 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,8 +19,6 @@
 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;
@@ -123,7 +121,7 @@
     private final BubbleController.NotifCallback mNotifCallback =
             new BubbleController.NotifCallback() {
         @Override
-        public void removeNotification(@NonNull final NotificationEntry entry, int reason) {
+        public void removeNotification(NotificationEntry entry, int reason) {
             if (isInterceptingDismissal(entry)) {
                 mInterceptedDismissalEntries.remove(entry.getKey());
                 mOnEndDismissInterception.onEndDismissInterception(mDismissInterceptor, entry,
@@ -143,7 +141,7 @@
         }
 
         @Override
-        public void maybeCancelSummary(@NonNull final NotificationEntry entry) {
+        public void maybeCancelSummary(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 d0a872e..80785db 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/phone/NotificationGroupManager.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/phone/NotificationGroupManager.java
@@ -16,7 +16,6 @@
 
 package com.android.systemui.statusbar.phone;
 
-import android.annotation.NonNull;
 import android.annotation.Nullable;
 import android.service.notification.StatusBarNotification;
 import android.util.ArraySet;
@@ -455,7 +454,7 @@
      * If there is a {@link NotificationGroup} associated with the provided entry, this method
      * will update the suppression of that group.
      */
-    public void updateSuppression(@NonNull final NotificationEntry entry) {
+    public void updateSuppression(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 52923ab..96e868d 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().getKey(), BubbleController.DISMISS_USER_GESTURE);
+                mRow.getEntry(), 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().getKey(), BubbleController.DISMISS_USER_GESTURE);
+                mRow.getEntry(), BubbleController.DISMISS_USER_GESTURE);
 
         Bubble b = mBubbleData.getOverflowBubbleWithKey(mRow.getEntry().getKey());
         assertThat(mBubbleData.getOverflowBubbles()).isEqualTo(ImmutableList.of(b));
@@ -350,10 +350,9 @@
         mBubbleController.updateBubble(mRow.getEntry(), /* suppressFlyout */
                 false, /* showInShade */ true);
         mBubbleController.removeBubble(
-                mRow.getEntry().getKey(), BubbleController.DISMISS_USER_GESTURE);
+                mRow.getEntry(), BubbleController.DISMISS_USER_GESTURE);
 
-        mBubbleController.removeBubble(
-                mRow.getEntry().getKey(), BubbleController.DISMISS_NOTIF_CANCEL);
+        mBubbleController.removeBubble(mRow.getEntry(), BubbleController.DISMISS_NOTIF_CANCEL);
         verify(mNotificationEntryManager, times(1)).performRemoveNotification(
                 eq(mRow.getEntry().getSbn()), anyInt());
         assertThat(mBubbleData.getOverflowBubbles()).isEmpty();
@@ -366,7 +365,7 @@
         assertTrue(mBubbleController.hasBubbles());
 
         mBubbleController.removeBubble(
-                mRow.getEntry().getKey(), BubbleController.DISMISS_USER_CHANGED);
+                mRow.getEntry(), BubbleController.DISMISS_USER_CHANGED);
         verify(mNotificationEntryManager, never()).performRemoveNotification(
                 eq(mRow.getEntry().getSbn()), anyInt());
         assertFalse(mBubbleController.hasBubbles());
@@ -564,8 +563,7 @@
 
         // Dismiss currently expanded
         mBubbleController.removeBubble(
-                mBubbleData.getBubbleInStackWithKey(stackView.getExpandedBubble().getKey())
-                        .getEntry().getKey(),
+                mBubbleData.getBubbleInStackWithKey(stackView.getExpandedBubble().getKey()).getEntry(),
                 BubbleController.DISMISS_USER_GESTURE);
         verify(mBubbleExpandListener).onBubbleExpandChanged(false, mRow2.getEntry().getKey());
 
@@ -576,8 +574,7 @@
 
         // Dismiss that one
         mBubbleController.removeBubble(
-                mBubbleData.getBubbleInStackWithKey(stackView.getExpandedBubble().getKey())
-                        .getEntry().getKey(),
+                mBubbleData.getBubbleInStackWithKey(stackView.getExpandedBubble().getKey()).getEntry(),
                 BubbleController.DISMISS_USER_GESTURE);
 
         // Make sure state changes and collapse happens
@@ -703,7 +700,7 @@
     @Test
     public void testDeleteIntent_removeBubble_aged() throws PendingIntent.CanceledException {
         mBubbleController.updateBubble(mRow.getEntry());
-        mBubbleController.removeBubble(mRow.getEntry().getKey(), BubbleController.DISMISS_AGED);
+        mBubbleController.removeBubble(mRow.getEntry(), BubbleController.DISMISS_AGED);
         verify(mDeleteIntent, never()).send();
     }
 
@@ -711,7 +708,7 @@
     public void testDeleteIntent_removeBubble_user() throws PendingIntent.CanceledException {
         mBubbleController.updateBubble(mRow.getEntry());
         mBubbleController.removeBubble(
-                mRow.getEntry().getKey(), BubbleController.DISMISS_USER_GESTURE);
+                mRow.getEntry(), BubbleController.DISMISS_USER_GESTURE);
         verify(mDeleteIntent, times(1)).send();
     }
 
@@ -811,7 +808,7 @@
 
         // Dismiss the bubble into overflow.
         mBubbleController.removeBubble(
-                mRow.getEntry().getKey(), BubbleController.DISMISS_USER_GESTURE);
+                mRow.getEntry(), BubbleController.DISMISS_USER_GESTURE);
         assertFalse(mBubbleController.hasBubbles());
 
         boolean intercepted = mRemoveInterceptor.onNotificationRemoveRequested(
@@ -832,7 +829,7 @@
                 mRow.getEntry()));
 
         mBubbleController.removeBubble(
-                mRow.getEntry().getKey(), BubbleController.DISMISS_NO_LONGER_BUBBLE);
+                mRow.getEntry(), BubbleController.DISMISS_NO_LONGER_BUBBLE);
         assertFalse(mBubbleController.hasBubbles());
 
         boolean intercepted = mRemoveInterceptor.onNotificationRemoveRequested(
@@ -854,12 +851,12 @@
 
         mBubbleData.setMaxOverflowBubbles(1);
         mBubbleController.removeBubble(
-                mRow.getEntry().getKey(), BubbleController.DISMISS_USER_GESTURE);
+                mRow.getEntry(), BubbleController.DISMISS_USER_GESTURE);
         assertEquals(mBubbleData.getBubbles().size(), 2);
         assertEquals(mBubbleData.getOverflowBubbles().size(), 1);
 
         mBubbleController.removeBubble(
-                mRow2.getEntry().getKey(), BubbleController.DISMISS_USER_GESTURE);
+                mRow2.getEntry(), 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 882504b..66f119a 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/bubbles/BubbleDataTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/bubbles/BubbleDataTest.java
@@ -20,8 +20,11 @@
 
 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;
@@ -177,8 +180,7 @@
         mBubbleData.setListener(mListener);
 
         // Test
-        mBubbleData.notificationEntryRemoved(
-                mEntryA1.getKey(), BubbleController.DISMISS_USER_GESTURE);
+        mBubbleData.notificationEntryRemoved(mEntryA1, BubbleController.DISMISS_USER_GESTURE);
 
         // Verify
         verifyUpdateReceived();
@@ -295,14 +297,12 @@
         mBubbleData.setListener(mListener);
 
         mBubbleData.setMaxOverflowBubbles(1);
-        mBubbleData.notificationEntryRemoved(
-                mEntryA1.getKey(), BubbleController.DISMISS_USER_GESTURE);
+        mBubbleData.notificationEntryRemoved(mEntryA1, 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.getKey(), BubbleController.DISMISS_USER_GESTURE);
+        mBubbleData.notificationEntryRemoved(mEntryA2, BubbleController.DISMISS_USER_GESTURE);
         verifyUpdateReceived();
         assertOverflowChangedTo(ImmutableList.of(mBubbleA2));
     }
@@ -449,8 +449,7 @@
         mBubbleData.setListener(mListener);
 
         // Test
-        mBubbleData.notificationEntryRemoved(
-                mEntryA2.getKey(), BubbleController.DISMISS_USER_GESTURE);
+        mBubbleData.notificationEntryRemoved(mEntryA2, BubbleController.DISMISS_USER_GESTURE);
         verifyUpdateReceived();
         assertOrderChangedTo(mBubbleB2, mBubbleB1, mBubbleA1);
     }
@@ -470,8 +469,7 @@
         mBubbleData.setListener(mListener);
 
         // Test
-        mBubbleData.notificationEntryRemoved(
-                mEntryB1.getKey(), BubbleController.DISMISS_USER_GESTURE);
+        mBubbleData.notificationEntryRemoved(mEntryB1, BubbleController.DISMISS_USER_GESTURE);
         verifyUpdateReceived();
         assertOrderNotChanged();
     }
@@ -491,8 +489,7 @@
         mBubbleData.setListener(mListener);
 
         // Test
-        mBubbleData.notificationEntryRemoved(
-                mEntryA1.getKey(), BubbleController.DISMISS_NOTIF_CANCEL);
+        mBubbleData.notificationEntryRemoved(mEntryA1, BubbleController.DISMISS_NOTIF_CANCEL);
         verifyUpdateReceived();
         assertOrderChangedTo(mBubbleB2, mBubbleB1, mBubbleA2);
     }
@@ -513,14 +510,12 @@
         mBubbleData.setListener(mListener);
 
         // Test
-        mBubbleData.notificationEntryRemoved(
-                mEntryA1.getKey(), BubbleController.DISMISS_NOTIF_CANCEL);
+        mBubbleData.notificationEntryRemoved(mEntryA1, BubbleController.DISMISS_NOTIF_CANCEL);
         verifyUpdateReceived();
         assertOverflowChangedTo(ImmutableList.of(mBubbleA2));
 
         // Test
-        mBubbleData.notificationEntryRemoved(
-                mEntryA2.getKey(), BubbleController.DISMISS_GROUP_CANCELLED);
+        mBubbleData.notificationEntryRemoved(mEntryA2, BubbleController.DISMISS_GROUP_CANCELLED);
         verifyUpdateReceived();
         assertOverflowChangedTo(ImmutableList.of());
     }
@@ -539,8 +534,7 @@
         mBubbleData.setListener(mListener);
 
         // Test
-        mBubbleData.notificationEntryRemoved(
-                mEntryA2.getKey(), BubbleController.DISMISS_NOTIF_CANCEL);
+        mBubbleData.notificationEntryRemoved(mEntryA2, BubbleController.DISMISS_NOTIF_CANCEL);
         verifyUpdateReceived();
         assertSelectionChangedTo(mBubbleB2);
     }
@@ -631,8 +625,7 @@
         mBubbleData.setListener(mListener);
 
         // Test
-        mBubbleData.notificationEntryRemoved(
-                mEntryA1.getKey(), BubbleController.DISMISS_USER_GESTURE);
+        mBubbleData.notificationEntryRemoved(mEntryA1, BubbleController.DISMISS_USER_GESTURE);
 
         // Verify the selection was cleared.
         verifyUpdateReceived();
@@ -786,8 +779,7 @@
         mBubbleData.setListener(mListener);
 
         // Test
-        mBubbleData.notificationEntryRemoved(
-                mEntryB2.getKey(), BubbleController.DISMISS_USER_GESTURE);
+        mBubbleData.notificationEntryRemoved(mEntryB2, BubbleController.DISMISS_USER_GESTURE);
         verifyUpdateReceived();
         assertOrderChangedTo(mBubbleB1, mBubbleA2, mBubbleA1);
     }
@@ -812,13 +804,11 @@
         mBubbleData.setListener(mListener);
 
         // Test
-        mBubbleData.notificationEntryRemoved(
-                mEntryA2.getKey(), BubbleController.DISMISS_USER_GESTURE);
+        mBubbleData.notificationEntryRemoved(mEntryA2, BubbleController.DISMISS_USER_GESTURE);
         verifyUpdateReceived();
         assertSelectionChangedTo(mBubbleA1);
 
-        mBubbleData.notificationEntryRemoved(
-                mEntryA1.getKey(), BubbleController.DISMISS_USER_GESTURE);
+        mBubbleData.notificationEntryRemoved(mEntryA1, BubbleController.DISMISS_USER_GESTURE);
         verifyUpdateReceived();
         assertSelectionChangedTo(mBubbleB1);
     }
@@ -933,8 +923,7 @@
         mBubbleData.setListener(mListener);
 
         // Test
-        mBubbleData.notificationEntryRemoved(
-                mEntryA1.getKey(), BubbleController.DISMISS_USER_GESTURE);
+        mBubbleData.notificationEntryRemoved(mEntryA1, 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 c16801c..73b8760 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/bubbles/NewNotifPipelineBubbleControllerTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/bubbles/NewNotifPipelineBubbleControllerTest.java
@@ -285,8 +285,7 @@
         assertTrue(mBubbleController.hasBubbles());
         verify(mNotifCallback, times(1)).invalidateNotifications(anyString());
 
-        mBubbleController.removeBubble(
-                mRow.getEntry().getKey(), BubbleController.DISMISS_USER_GESTURE);
+        mBubbleController.removeBubble(mRow.getEntry(), BubbleController.DISMISS_USER_GESTURE);
         assertNull(mBubbleData.getBubbleInStackWithKey(mRow.getEntry().getKey()));
         verify(mNotifCallback, times(2)).invalidateNotifications(anyString());
     }
@@ -303,8 +302,7 @@
         mBubbleData.getBubbleInStackWithKey(mRow.getEntry().getKey()).setSuppressNotification(true);
 
         // Now remove the bubble
-        mBubbleController.removeBubble(
-                mRow.getEntry().getKey(), BubbleController.DISMISS_USER_GESTURE);
+        mBubbleController.removeBubble(mRow.getEntry(), BubbleController.DISMISS_USER_GESTURE);
         assertTrue(mBubbleData.hasOverflowBubbleWithKey(mRow.getEntry().getKey()));
 
         // We don't remove the notification since the bubble is still in overflow.
@@ -324,8 +322,7 @@
         mBubbleData.getBubbleInStackWithKey(mRow.getEntry().getKey()).setSuppressNotification(true);
 
         // Now remove the bubble
-        mBubbleController.removeBubble(
-                mRow.getEntry().getKey(), BubbleController.DISMISS_NOTIF_CANCEL);
+        mBubbleController.removeBubble(mRow.getEntry(), BubbleController.DISMISS_NOTIF_CANCEL);
         assertFalse(mBubbleData.hasOverflowBubbleWithKey(mRow.getEntry().getKey()));
 
         // Since the notif is dismissed and not in overflow, once the bubble is removed,
@@ -505,8 +502,7 @@
 
         // Dismiss currently expanded
         mBubbleController.removeBubble(
-                mBubbleData.getBubbleInStackWithKey(
-                        stackView.getExpandedBubble().getKey()).getEntry().getKey(),
+                mBubbleData.getBubbleInStackWithKey(stackView.getExpandedBubble().getKey()).getEntry(),
                 BubbleController.DISMISS_USER_GESTURE);
         verify(mBubbleExpandListener).onBubbleExpandChanged(false, mRow2.getEntry().getKey());
 
@@ -517,8 +513,7 @@
 
         // Dismiss that one
         mBubbleController.removeBubble(
-                mBubbleData.getBubbleInStackWithKey(
-                        stackView.getExpandedBubble().getKey()).getEntry().getKey(),
+                mBubbleData.getBubbleInStackWithKey(stackView.getExpandedBubble().getKey()).getEntry(),
                 BubbleController.DISMISS_USER_GESTURE);
 
         // Make sure state changes and collapse happens
@@ -618,7 +613,7 @@
     @Test
     public void testDeleteIntent_removeBubble_aged() throws PendingIntent.CanceledException {
         mBubbleController.updateBubble(mRow.getEntry());
-        mBubbleController.removeBubble(mRow.getEntry().getKey(), BubbleController.DISMISS_AGED);
+        mBubbleController.removeBubble(mRow.getEntry(), BubbleController.DISMISS_AGED);
         verify(mDeleteIntent, never()).send();
     }
 
@@ -626,7 +621,7 @@
     public void testDeleteIntent_removeBubble_user() throws PendingIntent.CanceledException {
         mBubbleController.updateBubble(mRow.getEntry());
         mBubbleController.removeBubble(
-                mRow.getEntry().getKey(), BubbleController.DISMISS_USER_GESTURE);
+                mRow.getEntry(), BubbleController.DISMISS_USER_GESTURE);
         verify(mDeleteIntent, times(1)).send();
     }
 
@@ -691,7 +686,7 @@
 
         // Dismiss the bubble
         mBubbleController.removeBubble(
-                mRow.getEntry().getKey(), BubbleController.DISMISS_USER_GESTURE);
+                mRow.getEntry(), BubbleController.DISMISS_USER_GESTURE);
         assertFalse(mBubbleController.hasBubbles());
 
         // Dismiss the notification
@@ -712,7 +707,7 @@
 
         // Dismiss the bubble
         mBubbleController.removeBubble(
-                mRow.getEntry().getKey(), BubbleController.DISMISS_NOTIF_CANCEL);
+                mRow.getEntry(), BubbleController.DISMISS_NOTIF_CANCEL);
         assertFalse(mBubbleController.hasBubbles());
 
         // Dismiss the notification