Merge "Revert "decouple Bubble from NotificationEntry"" into rvc-dev
diff --git a/packages/SystemUI/src/com/android/systemui/bubbles/Bubble.java b/packages/SystemUI/src/com/android/systemui/bubbles/Bubble.java
index 6da7bc8..7f78ddf 100644
--- a/packages/SystemUI/src/com/android/systemui/bubbles/Bubble.java
+++ b/packages/SystemUI/src/com/android/systemui/bubbles/Bubble.java
@@ -15,6 +15,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;
 
@@ -28,19 +29,21 @@
 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;
 import android.content.res.Resources;
 import android.graphics.Bitmap;
 import android.graphics.Path;
+import android.graphics.Rect;
 import android.graphics.drawable.Drawable;
-import android.graphics.drawable.Icon;
+import android.os.Bundle;
 import android.os.UserHandle;
 import android.provider.Settings;
+import android.service.notification.StatusBarNotification;
 import android.util.Log;
 
 import com.android.internal.annotations.VisibleForTesting;
-import com.android.internal.logging.InstanceId;
 import com.android.systemui.shared.system.SysUiStatsLog;
 import com.android.systemui.statusbar.notification.collection.NotificationEntry;
 
@@ -54,12 +57,17 @@
 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 long mLastUpdated;
     private long mLastAccessed;
 
-    @Nullable
     private BubbleController.NotificationSuppressionChangedListener mSuppressionListener;
 
     /** Whether the bubble should show a dot for the notification indicating updated content. */
@@ -67,6 +75,8 @@
 
     /** Whether flyout text should be suppressed, regardless of any other flags or state. */
     private boolean mSuppressFlyout;
+    /** Whether this bubble should auto expand regardless of the normal flag, used for overflow. */
+    private boolean mShouldAutoExpand;
 
     // Items that are typically loaded later
     private String mAppName;
@@ -82,7 +92,6 @@
      * Presentational info about the flyout.
      */
     public static class FlyoutMessage {
-        @Nullable public Icon senderIcon;
         @Nullable public Drawable senderAvatar;
         @Nullable public CharSequence senderName;
         @Nullable public CharSequence message;
@@ -100,39 +109,16 @@
     private UserHandle mUser;
     @NonNull
     private String mPackageName;
-    @Nullable
-    private String mTitle;
-    @Nullable
-    private Icon mIcon;
-    private boolean mIsBubble;
-    private boolean mIsVisuallyInterruptive;
-    private boolean mIsClearable;
-    private boolean mShouldSuppressNotificationDot;
-    private boolean mShouldSuppressNotificationList;
-    private boolean mShouldSuppressPeek;
     private int mDesiredHeight;
     @DimenRes
     private int mDesiredHeightResId;
 
-    /** for logging **/
-    @Nullable
-    private InstanceId mInstanceId;
-    @Nullable
-    private String mChannelId;
-    private int mNotificationId;
-    private int mAppUid = -1;
-
-    @Nullable
-    private PendingIntent mIntent;
-    @Nullable
-    private PendingIntent mDeleteIntent;
-
     /**
      * 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(@NonNull final String key, @NonNull final ShortcutInfo shortcutInfo,
-            final int desiredHeight, final int desiredHeightResId, @Nullable final String title) {
+            final int desiredHeight, final int desiredHeightResId) {
         Objects.requireNonNull(key);
         Objects.requireNonNull(shortcutInfo);
         mShortcutInfo = shortcutInfo;
@@ -140,10 +126,8 @@
         mFlags = 0;
         mUser = shortcutInfo.getUserHandle();
         mPackageName = shortcutInfo.getPackage();
-        mIcon = shortcutInfo.getIcon();
         mDesiredHeight = desiredHeight;
         mDesiredHeightResId = desiredHeightResId;
-        mTitle = title;
     }
 
     /** Used in tests when no UI is required. */
@@ -161,6 +145,12 @@
         return mKey;
     }
 
+    @Nullable
+    public NotificationEntry getEntry() {
+        return mEntry;
+    }
+
+    @NonNull
     public UserHandle getUser() {
         return mUser;
     }
@@ -213,7 +203,14 @@
 
     @Nullable
     public String getTitle() {
-        return mTitle;
+        final CharSequence titleCharSeq;
+        if (mEntry == null) {
+            titleCharSeq = null;
+        } else {
+            titleCharSeq = mEntry.getSbn().getNotification().extras.getCharSequence(
+                    Notification.EXTRA_TITLE);
+        }
+        return titleCharSeq != null ? titleCharSeq.toString() : null;
     }
 
     /**
@@ -334,44 +331,17 @@
     void setEntry(@NonNull final NotificationEntry entry) {
         Objects.requireNonNull(entry);
         Objects.requireNonNull(entry.getSbn());
+        mEntry = entry;
         mLastUpdated = entry.getSbn().getPostTime();
-        mIsBubble = entry.getSbn().getNotification().isBubbleNotification();
+        mFlags = entry.getSbn().getNotification().flags;
         mPackageName = entry.getSbn().getPackageName();
         mUser = entry.getSbn().getUser();
-        mTitle = getTitle(entry);
-        mIsClearable = entry.isClearable();
-        mShouldSuppressNotificationDot = entry.shouldSuppressNotificationDot();
-        mShouldSuppressNotificationList = entry.shouldSuppressNotificationList();
-        mShouldSuppressPeek = entry.shouldSuppressPeek();
-        mChannelId = entry.getSbn().getNotification().getChannelId();
-        mNotificationId = entry.getSbn().getId();
-        mAppUid = entry.getSbn().getUid();
-        mInstanceId = entry.getSbn().getInstanceId();
-        mFlyoutMessage = BubbleViewInfoTask.extractFlyoutMessage(entry);
-        if (entry.getRanking() != null) {
-            mShortcutInfo = entry.getRanking().getShortcutInfo() != null
-                    ? entry.getRanking().getShortcutInfo() : mShortcutInfo;
-            mIsVisuallyInterruptive = entry.getRanking().visuallyInterruptive();
-        }
         if (entry.getBubbleMetadata() != null) {
-            mFlags = entry.getBubbleMetadata().getFlags();
             mDesiredHeight = entry.getBubbleMetadata().getDesiredHeight();
             mDesiredHeightResId = entry.getBubbleMetadata().getDesiredHeightResId();
-            mIcon = entry.getBubbleMetadata().getIcon();
-            mIntent = entry.getBubbleMetadata().getIntent();
-            mDeleteIntent = entry.getBubbleMetadata().getDeleteIntent();
         }
     }
 
-    @Nullable
-    Icon getIcon() {
-        return mIcon;
-    }
-
-    boolean isVisuallyInterruptive() {
-        return mIsVisuallyInterruptive;
-    }
-
     /**
      * @return the last time this bubble was updated or accessed, whichever is most recent.
      */
@@ -394,19 +364,6 @@
         return mExpandedView != null ? mExpandedView.getVirtualDisplayId() : INVALID_DISPLAY;
     }
 
-    public InstanceId getInstanceId() {
-        return mInstanceId;
-    }
-
-    @Nullable
-    public String getChannelId() {
-        return mChannelId;
-    }
-
-    public int getNotificationId() {
-        return mNotificationId;
-    }
-
     /**
      * Should be invoked whenever a Bubble is accessed (selected while expanded).
      */
@@ -427,19 +384,24 @@
      * Whether this notification should be shown in the shade.
      */
     boolean showInShade() {
-        return !shouldSuppressNotification() || !mIsClearable;
+        if (mEntry == null) return false;
+        return !shouldSuppressNotification() || !mEntry.isClearable();
     }
 
     /**
      * 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) {
-            mFlags |= Notification.BubbleMetadata.FLAG_SUPPRESS_NOTIFICATION;
+            flags |= Notification.BubbleMetadata.FLAG_SUPPRESS_NOTIFICATION;
         } else {
-            mFlags &= ~Notification.BubbleMetadata.FLAG_SUPPRESS_NOTIFICATION;
+            flags &= ~Notification.BubbleMetadata.FLAG_SUPPRESS_NOTIFICATION;
         }
+        data.setFlags(flags);
 
         if (showInShade() != prevShowInShade && mSuppressionListener != null) {
             mSuppressionListener.onBubbleNotificationSuppressionChange(this);
@@ -462,8 +424,9 @@
      */
     @Override
     public boolean showDot() {
+        if (mEntry == null) return false;
         return mShowBubbleUpdateDot
-                && !mShouldSuppressNotificationDot
+                && !mEntry.shouldSuppressNotificationDot()
                 && !shouldSuppressNotification();
     }
 
@@ -471,9 +434,10 @@
      * Whether the flyout for the bubble should be shown.
      */
     boolean showFlyout() {
-        return !mSuppressFlyout && !mShouldSuppressPeek
+        if (mEntry == null) return false;
+        return !mSuppressFlyout && !mEntry.shouldSuppressPeek()
                 && !shouldSuppressNotification()
-                && !mShouldSuppressNotificationList;
+                && !mEntry.shouldSuppressNotificationList();
     }
 
     /**
@@ -516,14 +480,25 @@
         }
     }
 
-    @Nullable
-    PendingIntent getBubbleIntent() {
-        return mIntent;
+    /**
+     * Whether shortcut information should be used to populate the bubble.
+     * <p>
+     * To populate the activity use {@link LauncherApps#startShortcut(ShortcutInfo, Rect, Bundle)}.
+     * To populate the icon use {@link LauncherApps#getShortcutIconDrawable(ShortcutInfo, int)}.
+     */
+    boolean usingShortcutInfo() {
+        return mEntry != null && mEntry.getBubbleMetadata().getShortcutId() != null
+                || mShortcutInfo != null;
     }
 
     @Nullable
-    PendingIntent getDeleteIntent() {
-        return mDeleteIntent;
+    PendingIntent getBubbleIntent() {
+        if (mEntry == null) return null;
+        Notification.BubbleMetadata data = mEntry.getBubbleMetadata();
+        if (data != null) {
+            return data.getIntent();
+        }
+        return null;
     }
 
     Intent getSettingsIntent(final Context context) {
@@ -539,12 +514,8 @@
         return intent;
     }
 
-    public int getAppUid() {
-        return mAppUid;
-    }
-
     private int getUid(final Context context) {
-        if (mAppUid != -1) return mAppUid;
+        if (mEntry != null) return mEntry.getSbn().getUid();
         final PackageManager pm = context.getPackageManager();
         if (pm == null) return -1;
         try {
@@ -577,27 +548,24 @@
     }
 
     private boolean shouldSuppressNotification() {
-        return isEnabled(Notification.BubbleMetadata.FLAG_SUPPRESS_NOTIFICATION);
+        if (mEntry == null) return true;
+        return mEntry.getBubbleMetadata() != null
+                && mEntry.getBubbleMetadata().isNotificationSuppressed();
     }
 
-    public boolean shouldAutoExpand() {
-        return isEnabled(Notification.BubbleMetadata.FLAG_AUTO_EXPAND_BUBBLE);
+    boolean shouldAutoExpand() {
+        if (mEntry == null) return mShouldAutoExpand;
+        Notification.BubbleMetadata metadata = mEntry.getBubbleMetadata();
+        return (metadata != null && metadata.getAutoExpandBubble()) ||  mShouldAutoExpand;
     }
 
     void setShouldAutoExpand(boolean shouldAutoExpand) {
-        if (shouldAutoExpand) {
-            enable(Notification.BubbleMetadata.FLAG_AUTO_EXPAND_BUBBLE);
-        } else {
-            disable(Notification.BubbleMetadata.FLAG_AUTO_EXPAND_BUBBLE);
-        }
-    }
-
-    public void setIsBubble(final boolean isBubble) {
-        mIsBubble = isBubble;
+        mShouldAutoExpand = shouldAutoExpand;
     }
 
     public boolean isBubble() {
-        return mIsBubble;
+        if (mEntry == null) return (mFlags & FLAG_BUBBLE) != 0;
+        return (mEntry.getSbn().getNotification().flags & FLAG_BUBBLE) != 0;
     }
 
     public void enable(int option) {
@@ -608,10 +576,6 @@
         mFlags &= ~option;
     }
 
-    public boolean isEnabled(int option) {
-        return (mFlags & option) != 0;
-    }
-
     @Override
     public String toString() {
         return "Bubble{" + mKey + '}';
@@ -646,24 +610,34 @@
 
     @Override
     public void logUIEvent(int bubbleCount, int action, float normalX, float normalY, int index) {
-        SysUiStatsLog.write(SysUiStatsLog.BUBBLE_UI_CHANGED,
-                mPackageName,
-                mChannelId,
-                mNotificationId,
-                index,
-                bubbleCount,
-                action,
-                normalX,
-                normalY,
-                showInShade(),
-                false /* isOngoing (unused) */,
-                false /* isAppForeground (unused) */);
-    }
-
-    @Nullable
-    private static String getTitle(@NonNull final NotificationEntry e) {
-        final CharSequence titleCharSeq = e.getSbn().getNotification().extras.getCharSequence(
-                Notification.EXTRA_TITLE);
-        return titleCharSeq == null ? null : titleCharSeq.toString();
+        if (this.getEntry() == null
+                || this.getEntry().getSbn() == null) {
+            SysUiStatsLog.write(SysUiStatsLog.BUBBLE_UI_CHANGED,
+                    null /* package name */,
+                    null /* notification channel */,
+                    0 /* notification ID */,
+                    0 /* bubble position */,
+                    bubbleCount,
+                    action,
+                    normalX,
+                    normalY,
+                    false /* unread bubble */,
+                    false /* on-going bubble */,
+                    false /* isAppForeground (unused) */);
+        } else {
+            StatusBarNotification notification = this.getEntry().getSbn();
+            SysUiStatsLog.write(SysUiStatsLog.BUBBLE_UI_CHANGED,
+                    notification.getPackageName(),
+                    notification.getNotification().getChannelId(),
+                    notification.getId(),
+                    index,
+                    bubbleCount,
+                    action,
+                    normalX,
+                    normalY,
+                    this.showInShade(),
+                    false /* isOngoing (unused) */,
+                    false /* isAppForeground (unused) */);
+        }
     }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/bubbles/BubbleController.java b/packages/SystemUI/src/com/android/systemui/bubbles/BubbleController.java
index b2c5402..c4c5da4 100644
--- a/packages/SystemUI/src/com/android/systemui/bubbles/BubbleController.java
+++ b/packages/SystemUI/src/com/android/systemui/bubbles/BubbleController.java
@@ -505,7 +505,8 @@
         addNotifCallback(new NotifCallback() {
             @Override
             public void removeNotification(NotificationEntry entry, int reason) {
-                mNotificationEntryManager.performRemoveNotification(entry.getSbn(), reason);
+                mNotificationEntryManager.performRemoveNotification(entry.getSbn(),
+                        reason);
             }
 
             @Override
@@ -636,13 +637,8 @@
                 mStackView.setExpandListener(mExpandListener);
             }
 
-            mStackView.setUnbubbleConversationCallback(key -> {
-                final NotificationEntry entry =
-                        mNotificationEntryManager.getPendingOrActiveNotif(key);
-                if (entry != null) {
-                    onUserChangedBubble(entry, false /* shouldBubble */);
-                }
-            });
+            mStackView.setUnbubbleConversationCallback(notificationEntry ->
+                    onUserChangedBubble(notificationEntry, false /* shouldBubble */));
         }
 
         addToWindowManagerMaybe();
@@ -1028,7 +1024,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(@NonNull final 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();
@@ -1104,8 +1103,7 @@
             mBubbleData.removeSuppressedSummary(groupKey);
 
             // Remove any associated bubble children with the summary
-            final List<Bubble> bubbleChildren = mBubbleData.getBubblesInGroup(
-                    groupKey, mNotificationEntryManager);
+            final List<Bubble> bubbleChildren = mBubbleData.getBubblesInGroup(groupKey);
             for (int i = 0; i < bubbleChildren.size(); i++) {
                 removeBubble(bubbleChildren.get(i).getKey(), DISMISS_GROUP_CANCELLED);
             }
@@ -1163,18 +1161,21 @@
 
     private void setIsBubble(@NonNull final Bubble b, final boolean isBubble) {
         Objects.requireNonNull(b);
-        b.setIsBubble(isBubble);
-        final NotificationEntry entry = mNotificationEntryManager
-                .getPendingOrActiveNotif(b.getKey());
-        if (entry != null) {
+        if (isBubble) {
+            b.enable(FLAG_BUBBLE);
+        } else {
+            b.disable(FLAG_BUBBLE);
+        }
+        if (b.getEntry() != null) {
             // Updating the entry to be a bubble will trigger our normal update flow
-            setIsBubble(entry, isBubble, b.shouldAutoExpand());
+            setIsBubble(b.getEntry(), isBubble, b.shouldAutoExpand());
         } else if (isBubble) {
-            // If bubble doesn't exist, it's a persisted bubble so we need to add it to the
-            // stack ourselves
+            // If we have no entry to update, it's a persisted bubble so
+            // we need to add it to the stack ourselves
             Bubble bubble = mBubbleData.getOrCreateBubble(null, b /* persistedBubble */);
             inflateAndAdd(bubble, bubble.shouldAutoExpand() /* suppressFlyout */,
                     !bubble.shouldAutoExpand() /* showInShade */);
+
         }
     }
 
@@ -1213,8 +1214,6 @@
                 if (reason == DISMISS_NOTIF_CANCEL) {
                     bubblesToBeRemovedFromRepository.add(bubble);
                 }
-                final NotificationEntry entry = mNotificationEntryManager.getPendingOrActiveNotif(
-                        bubble.getKey());
                 if (!mBubbleData.hasBubbleInStackWithKey(bubble.getKey())) {
                     if (!mBubbleData.hasOverflowBubbleWithKey(bubble.getKey())
                         && (!bubble.showInShade()
@@ -1223,27 +1222,26 @@
                         // The bubble is now gone & the notification is hidden from the shade, so
                         // time to actually remove it
                         for (NotifCallback cb : mCallbacks) {
-                            if (entry != null) {
-                                cb.removeNotification(entry, REASON_CANCEL);
+                            if (bubble.getEntry() != null) {
+                                cb.removeNotification(bubble.getEntry(), REASON_CANCEL);
                             }
                         }
                     } else {
                         if (bubble.isBubble()) {
                             setIsBubble(bubble, false /* isBubble */);
                         }
-                        if (entry != null && entry.getRow() != null) {
-                            entry.getRow().updateBubbleButton();
+                        if (bubble.getEntry() != null && bubble.getEntry().getRow() != null) {
+                            bubble.getEntry().getRow().updateBubbleButton();
                         }
                     }
 
                 }
-                if (entry != null) {
-                    final String groupKey = entry.getSbn().getGroupKey();
-                    if (mBubbleData.getBubblesInGroup(
-                            groupKey, mNotificationEntryManager).isEmpty()) {
+                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(entry);
+                            cb.maybeCancelSummary(bubble.getEntry());
                         }
                     }
                 }
@@ -1268,12 +1266,9 @@
 
             if (update.selectionChanged && mStackView != null) {
                 mStackView.setSelectedBubble(update.selectedBubble);
-                if (update.selectedBubble != null) {
-                    final NotificationEntry entry = mNotificationEntryManager
-                            .getPendingOrActiveNotif(update.selectedBubble.getKey());
-                    if (entry != null) {
-                        mNotificationGroupManager.updateSuppression(entry);
-                    }
+                if (update.selectedBubble != null && update.selectedBubble.getEntry() != null) {
+                    mNotificationGroupManager.updateSuppression(
+                            update.selectedBubble.getEntry());
                 }
             }
 
@@ -1346,8 +1341,7 @@
         }
 
         String groupKey = entry.getSbn().getGroupKey();
-        ArrayList<Bubble> bubbleChildren = mBubbleData.getBubblesInGroup(
-                groupKey, mNotificationEntryManager);
+        ArrayList<Bubble> bubbleChildren = mBubbleData.getBubblesInGroup(groupKey);
         boolean isSuppressedSummary = (mBubbleData.isSummarySuppressed(groupKey)
                 && mBubbleData.getSummaryKey(groupKey).equals(entry.getKey()));
         boolean isSummary = entry.getSbn().getNotification().isGroupSummary();
@@ -1367,15 +1361,9 @@
                     // As far as group manager is concerned, once a child is no longer shown
                     // in the shade, it is essentially removed.
                     Bubble bubbleChild = mBubbleData.getAnyBubbleWithkey(child.getKey());
-                    if (bubbleChild != null) {
-                        final NotificationEntry entry = mNotificationEntryManager
-                                .getPendingOrActiveNotif(bubbleChild.getKey());
-                        if (entry != null) {
-                            mNotificationGroupManager.onEntryRemoved(entry);
-                        }
-                        bubbleChild.setSuppressNotification(true);
-                        bubbleChild.setShowDot(false /* show */);
-                    }
+                    mNotificationGroupManager.onEntryRemoved(bubbleChild.getEntry());
+                    bubbleChild.setSuppressNotification(true);
+                    bubbleChild.setShowDot(false /* show */);
                 } else {
                     // non-bubbled children can be removed
                     for (NotifCallback cb : mCallbacks) {
diff --git a/packages/SystemUI/src/com/android/systemui/bubbles/BubbleData.java b/packages/SystemUI/src/com/android/systemui/bubbles/BubbleData.java
index c870612..20a9a8c 100644
--- a/packages/SystemUI/src/com/android/systemui/bubbles/BubbleData.java
+++ b/packages/SystemUI/src/com/android/systemui/bubbles/BubbleData.java
@@ -22,6 +22,7 @@
 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;
 import android.util.Log;
@@ -33,7 +34,6 @@
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.systemui.R;
 import com.android.systemui.bubbles.BubbleController.DismissReason;
-import com.android.systemui.statusbar.notification.NotificationEntryManager;
 import com.android.systemui.statusbar.notification.collection.NotificationEntry;
 
 import java.io.FileDescriptor;
@@ -256,7 +256,8 @@
         }
         mPendingBubbles.remove(bubble.getKey()); // No longer pending once we're here
         Bubble prevBubble = getBubbleInStackWithKey(bubble.getKey());
-        suppressFlyout |= !bubble.isVisuallyInterruptive();
+        suppressFlyout |= bubble.getEntry() == null
+                || !bubble.getEntry().getRanking().visuallyInterruptive();
 
         if (prevBubble == null) {
             // Create a new bubble
@@ -334,15 +335,13 @@
      * Retrieves any bubbles that are part of the notification group represented by the provided
      * group key.
      */
-    ArrayList<Bubble> getBubblesInGroup(@Nullable String groupKey, @NonNull
-            NotificationEntryManager nem) {
+    ArrayList<Bubble> getBubblesInGroup(@Nullable String groupKey) {
         ArrayList<Bubble> bubbleChildren = new ArrayList<>();
         if (groupKey == null) {
             return bubbleChildren;
         }
         for (Bubble b : mBubbles) {
-            final NotificationEntry entry = nem.getPendingOrActiveNotif(b.getKey());
-            if (entry != null && groupKey.equals(entry.getSbn().getGroupKey())) {
+            if (b.getEntry() != null && groupKey.equals(b.getEntry().getSbn().getGroupKey())) {
                 bubbleChildren.add(b);
             }
         }
@@ -440,7 +439,9 @@
             Bubble newSelected = mBubbles.get(newIndex);
             setSelectedBubbleInternal(newSelected);
         }
-        maybeSendDeleteIntent(reason, bubbleToRemove);
+        if (bubbleToRemove.getEntry() != null) {
+            maybeSendDeleteIntent(reason, bubbleToRemove.getEntry());
+        }
     }
 
     void overflowBubble(@DismissReason int reason, Bubble bubble) {
@@ -610,14 +611,21 @@
         return true;
     }
 
-    private void maybeSendDeleteIntent(@DismissReason int reason, @NonNull final Bubble bubble) {
-        if (reason != BubbleController.DISMISS_USER_GESTURE) return;
-        PendingIntent deleteIntent = bubble.getDeleteIntent();
-        if (deleteIntent == null) return;
-        try {
-            deleteIntent.send();
-        } catch (PendingIntent.CanceledException e) {
-            Log.w(TAG, "Failed to send delete intent for bubble with key: " + bubble.getKey());
+    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
+                    ? bubbleMetadata.getDeleteIntent()
+                    : null;
+            if (deleteIntent != null) {
+                try {
+                    deleteIntent.send();
+                } catch (PendingIntent.CanceledException e) {
+                    Log.w(TAG, "Failed to send delete intent for bubble with key: "
+                            + entry.getKey());
+                }
+            }
         }
     }
 
diff --git a/packages/SystemUI/src/com/android/systemui/bubbles/BubbleDataRepository.kt b/packages/SystemUI/src/com/android/systemui/bubbles/BubbleDataRepository.kt
index 0c25d14..d20f405 100644
--- a/packages/SystemUI/src/com/android/systemui/bubbles/BubbleDataRepository.kt
+++ b/packages/SystemUI/src/com/android/systemui/bubbles/BubbleDataRepository.kt
@@ -74,15 +74,11 @@
 
     private fun transform(userId: Int, bubbles: List<Bubble>): List<BubbleEntity> {
         return bubbles.mapNotNull { b ->
-            BubbleEntity(
-                    userId,
-                    b.packageName,
-                    b.shortcutInfo?.id ?: return@mapNotNull null,
-                    b.key,
-                    b.rawDesiredHeight,
-                    b.rawDesiredHeightResId,
-                    b.title
-            )
+            var shortcutId = b.shortcutInfo?.id
+            if (shortcutId == null) shortcutId = b.entry?.bubbleMetadata?.shortcutId
+            if (shortcutId == null) return@mapNotNull null
+            BubbleEntity(userId, b.packageName, shortcutId, b.key, b.rawDesiredHeight,
+                    b.rawDesiredHeightResId)
         }
     }
 
@@ -163,13 +159,8 @@
         val bubbles = entities.mapNotNull { entity ->
             shortcutMap[ShortcutKey(entity.userId, entity.packageName)]
                     ?.first { shortcutInfo -> entity.shortcutId == shortcutInfo.id }
-                    ?.let { shortcutInfo -> Bubble(
-                            entity.key,
-                            shortcutInfo,
-                            entity.desiredHeight,
-                            entity.desiredHeightResId,
-                            entity.title
-                    ) }
+                    ?.let { shortcutInfo -> Bubble(entity.key, shortcutInfo, entity.desiredHeight,
+                            entity.desiredHeightResId) }
         }
         uiScope.launch { cb(bubbles) }
     }
diff --git a/packages/SystemUI/src/com/android/systemui/bubbles/BubbleExpandedView.java b/packages/SystemUI/src/com/android/systemui/bubbles/BubbleExpandedView.java
index 2f7ffde..ca359b8 100644
--- a/packages/SystemUI/src/com/android/systemui/bubbles/BubbleExpandedView.java
+++ b/packages/SystemUI/src/com/android/systemui/bubbles/BubbleExpandedView.java
@@ -169,7 +169,7 @@
                             return;
                         }
                         try {
-                            if (!mIsOverflow && mBubble.getShortcutInfo() != null) {
+                            if (!mIsOverflow && mBubble.usingShortcutInfo()) {
                                 options.setApplyActivityFlagsForBubbles(true);
                                 mActivityView.startShortcutActivity(mBubble.getShortcutInfo(),
                                         options, null /* sourceBounds */);
diff --git a/packages/SystemUI/src/com/android/systemui/bubbles/BubbleFlyoutView.java b/packages/SystemUI/src/com/android/systemui/bubbles/BubbleFlyoutView.java
index 1fa3aaa..8c76cda 100644
--- a/packages/SystemUI/src/com/android/systemui/bubbles/BubbleFlyoutView.java
+++ b/packages/SystemUI/src/com/android/systemui/bubbles/BubbleFlyoutView.java
@@ -31,7 +31,6 @@
 import android.graphics.Path;
 import android.graphics.PointF;
 import android.graphics.RectF;
-import android.graphics.drawable.Drawable;
 import android.graphics.drawable.ShapeDrawable;
 import android.text.TextUtils;
 import android.view.LayoutInflater;
@@ -224,10 +223,9 @@
             float[] dotCenter,
             boolean hideDot) {
 
-        final Drawable senderAvatar = flyoutMessage.senderAvatar;
-        if (senderAvatar != null && flyoutMessage.isGroupChat) {
+        if (flyoutMessage.senderAvatar != null && flyoutMessage.isGroupChat) {
             mSenderAvatar.setVisibility(VISIBLE);
-            mSenderAvatar.setImageDrawable(senderAvatar);
+            mSenderAvatar.setImageDrawable(flyoutMessage.senderAvatar);
         } else {
             mSenderAvatar.setVisibility(GONE);
             mSenderAvatar.setTranslationX(0);
diff --git a/packages/SystemUI/src/com/android/systemui/bubbles/BubbleIconFactory.java b/packages/SystemUI/src/com/android/systemui/bubbles/BubbleIconFactory.java
index a799f2d..74231c6 100644
--- a/packages/SystemUI/src/com/android/systemui/bubbles/BubbleIconFactory.java
+++ b/packages/SystemUI/src/com/android/systemui/bubbles/BubbleIconFactory.java
@@ -15,8 +15,7 @@
  */
 package com.android.systemui.bubbles;
 
-import android.annotation.NonNull;
-import android.annotation.Nullable;
+import android.app.Notification;
 import android.content.Context;
 import android.content.Intent;
 import android.content.pm.LauncherApps;
@@ -51,14 +50,15 @@
     /**
      * Returns the drawable that the developer has provided to display in the bubble.
      */
-    Drawable getBubbleDrawable(@NonNull final Context context,
-            @Nullable final ShortcutInfo shortcutInfo, @Nullable final Icon ic) {
+    Drawable getBubbleDrawable(Context context, ShortcutInfo shortcutInfo,
+            Notification.BubbleMetadata metadata) {
         if (shortcutInfo != null) {
             LauncherApps launcherApps =
                     (LauncherApps) context.getSystemService(Context.LAUNCHER_APPS_SERVICE);
             int density = context.getResources().getConfiguration().densityDpi;
             return launcherApps.getShortcutIconDrawable(shortcutInfo, density);
         } else {
+            Icon ic = metadata.getIcon();
             if (ic != null) {
                 if (ic.getType() == Icon.TYPE_URI
                         || ic.getType() == Icon.TYPE_URI_ADAPTIVE_BITMAP) {
diff --git a/packages/SystemUI/src/com/android/systemui/bubbles/BubbleLoggerImpl.java b/packages/SystemUI/src/com/android/systemui/bubbles/BubbleLoggerImpl.java
index c1dd8c3..c5faae0 100644
--- a/packages/SystemUI/src/com/android/systemui/bubbles/BubbleLoggerImpl.java
+++ b/packages/SystemUI/src/com/android/systemui/bubbles/BubbleLoggerImpl.java
@@ -16,6 +16,8 @@
 
 package com.android.systemui.bubbles;
 
+import android.service.notification.StatusBarNotification;
+
 import com.android.internal.logging.UiEventLoggerImpl;
 
 /**
@@ -30,11 +32,12 @@
      * @param e UI event
      */
     public void log(Bubble b, UiEventEnum e) {
-        if (b.getInstanceId() == null) {
+        if (b.getEntry() == null) {
             // Added from persistence -- TODO log this with specific event?
             return;
         }
-        logWithInstanceId(e, b.getAppUid(), b.getPackageName(), b.getInstanceId());
+        StatusBarNotification sbn = b.getEntry().getSbn();
+        logWithInstanceId(e, sbn.getUid(), sbn.getPackageName(), sbn.getInstanceId());
     }
 
     /**
diff --git a/packages/SystemUI/src/com/android/systemui/bubbles/BubbleOverflowActivity.java b/packages/SystemUI/src/com/android/systemui/bubbles/BubbleOverflowActivity.java
index 0b25c44..b644079 100644
--- a/packages/SystemUI/src/com/android/systemui/bubbles/BubbleOverflowActivity.java
+++ b/packages/SystemUI/src/com/android/systemui/bubbles/BubbleOverflowActivity.java
@@ -291,7 +291,9 @@
                 });
 
         // If the bubble was persisted, the entry is null but it should have shortcut info
-        ShortcutInfo info = b.getShortcutInfo();
+        ShortcutInfo info = b.getEntry() == null
+                ? b.getShortcutInfo()
+                : b.getEntry().getRanking().getShortcutInfo();
         if (info == null) {
             Log.d(TAG, "ShortcutInfo required to bubble but none found for " + b);
         } else {
diff --git a/packages/SystemUI/src/com/android/systemui/bubbles/BubbleStackView.java b/packages/SystemUI/src/com/android/systemui/bubbles/BubbleStackView.java
index 41dc081..23bd174 100644
--- a/packages/SystemUI/src/com/android/systemui/bubbles/BubbleStackView.java
+++ b/packages/SystemUI/src/com/android/systemui/bubbles/BubbleStackView.java
@@ -92,6 +92,7 @@
 import com.android.systemui.model.SysUiState;
 import com.android.systemui.shared.system.QuickStepContract;
 import com.android.systemui.shared.system.SysUiStatsLog;
+import com.android.systemui.statusbar.notification.collection.NotificationEntry;
 import com.android.systemui.statusbar.phone.CollapsedStatusBarFragment;
 import com.android.systemui.util.DismissCircleView;
 import com.android.systemui.util.FloatingContentCoordinator;
@@ -293,7 +294,7 @@
     private BubbleController.BubbleExpandListener mExpandListener;
 
     /** Callback to run when we want to unbubble the given notification's conversation. */
-    private Consumer<String> mUnbubbleConversationCallback;
+    private Consumer<NotificationEntry> mUnbubbleConversationCallback;
 
     private SysUiState mSysUiState;
 
@@ -1013,7 +1014,10 @@
         mManageMenu.findViewById(R.id.bubble_manage_menu_dont_bubble_container).setOnClickListener(
                 view -> {
                     showManageMenu(false /* show */);
-                    mUnbubbleConversationCallback.accept(mBubbleData.getSelectedBubble().getKey());
+                    final Bubble bubble = mBubbleData.getSelectedBubble();
+                    if (bubble != null && mBubbleData.hasBubbleInStackWithKey(bubble.getKey())) {
+                        mUnbubbleConversationCallback.accept(bubble.getEntry());
+                    }
                 });
 
         mManageMenu.findViewById(R.id.bubble_manage_menu_settings_container).setOnClickListener(
@@ -1365,7 +1369,7 @@
 
     /** Sets the function to call to un-bubble the given conversation. */
     public void setUnbubbleConversationCallback(
-            Consumer<String> unbubbleConversationCallback) {
+            Consumer<NotificationEntry> unbubbleConversationCallback) {
         mUnbubbleConversationCallback = unbubbleConversationCallback;
     }
 
diff --git a/packages/SystemUI/src/com/android/systemui/bubbles/BubbleViewInfoTask.java b/packages/SystemUI/src/com/android/systemui/bubbles/BubbleViewInfoTask.java
index 3e4ff52..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,8 @@
 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;
 import android.util.PathParser;
@@ -51,7 +53,6 @@
 
 import java.lang.ref.WeakReference;
 import java.util.List;
-import java.util.Objects;
 
 /**
  * Simple task to inflate views & load necessary info to display a bubble.
@@ -128,10 +129,35 @@
         @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) {
             BubbleViewInfo info = new BubbleViewInfo();
 
             // View inflation: only should do this once per bubble
-            if (!skipInflation && !b.isInflated()) {
+            if (!isInflated) {
                 LayoutInflater inflater = LayoutInflater.from(c);
                 info.imageView = (BadgedImageView) inflater.inflate(
                         R.layout.bubble_view, stackView, false /* attachToRoot */);
@@ -141,8 +167,8 @@
                 info.expandedView.setStackView(stackView);
             }
 
-            if (b.getShortcutInfo() != null) {
-                info.shortcutInfo = b.getShortcutInfo();
+            if (shortcutInfo != null) {
+                info.shortcutInfo = shortcutInfo;
             }
 
             // App name & app icon
@@ -152,7 +178,7 @@
             Drawable appIcon;
             try {
                 appInfo = pm.getApplicationInfo(
-                        b.getPackageName(),
+                        packageName,
                         PackageManager.MATCH_UNINSTALLED_PACKAGES
                                 | PackageManager.MATCH_DISABLED_COMPONENTS
                                 | PackageManager.MATCH_DIRECT_BOOT_UNAWARE
@@ -160,17 +186,17 @@
                 if (appInfo != null) {
                     info.appName = String.valueOf(pm.getApplicationLabel(appInfo));
                 }
-                appIcon = pm.getApplicationIcon(b.getPackageName());
-                badgedIcon = pm.getUserBadgedIcon(appIcon, b.getUser());
+                appIcon = pm.getApplicationIcon(packageName);
+                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: " + b.getPackageName());
+                Log.w(TAG, "Unable to find package: " + packageName);
                 return null;
             }
 
             // Badged bubble image
             Drawable bubbleDrawable = iconFactory.getBubbleDrawable(c, info.shortcutInfo,
-                    b.getIcon());
+                    entry == null ? null : entry.getBubbleMetadata());
             if (bubbleDrawable == null) {
                 // Default to app icon
                 bubbleDrawable = appIcon;
@@ -196,10 +222,8 @@
                     Color.WHITE, WHITE_SCRIM_ALPHA);
 
             // Flyout
-            info.flyoutMessage = b.getFlyoutMessage();
-            if (info.flyoutMessage != null) {
-                info.flyoutMessage.senderAvatar =
-                        loadSenderAvatar(c, info.flyoutMessage.senderIcon);
+            if (entry != null) {
+                info.flyoutMessage = extractFlyoutMessage(c, entry);
             }
             return info;
         }
@@ -211,8 +235,8 @@
      * notification, based on its type. Returns null if there should not be an update message.
      */
     @NonNull
-    static Bubble.FlyoutMessage extractFlyoutMessage(NotificationEntry entry) {
-        Objects.requireNonNull(entry);
+    static Bubble.FlyoutMessage extractFlyoutMessage(Context context,
+            NotificationEntry entry) {
         final Notification underlyingNotif = entry.getSbn().getNotification();
         final Class<? extends Notification.Style> style = underlyingNotif.getNotificationStyle();
 
@@ -240,9 +264,20 @@
                 if (latestMessage != null) {
                     bubbleMessage.message = latestMessage.getText();
                     Person sender = latestMessage.getSenderPerson();
-                    bubbleMessage.senderName = sender != null ? sender.getName() : null;
+                    bubbleMessage.senderName = sender != null
+                            ? sender.getName()
+                            : null;
+
                     bubbleMessage.senderAvatar = null;
-                    bubbleMessage.senderIcon = sender != null ? sender.getIcon() : null;
+                    if (sender != null && sender.getIcon() != null) {
+                        if (sender.getIcon().getType() == Icon.TYPE_URI
+                                || sender.getIcon().getType() == Icon.TYPE_URI_ADAPTIVE_BITMAP) {
+                            context.grantUriPermission(context.getPackageName(),
+                                    sender.getIcon().getUri(),
+                                    Intent.FLAG_GRANT_READ_URI_PERMISSION);
+                        }
+                        bubbleMessage.senderAvatar = sender.getIcon().loadDrawable(context);
+                    }
                     return bubbleMessage;
                 }
             } else if (Notification.InboxStyle.class.equals(style)) {
@@ -271,15 +306,4 @@
 
         return bubbleMessage;
     }
-
-    @Nullable
-    static Drawable loadSenderAvatar(@NonNull final Context context, @Nullable final Icon icon) {
-        Objects.requireNonNull(context);
-        if (icon == null) return null;
-        if (icon.getType() == Icon.TYPE_URI || icon.getType() == Icon.TYPE_URI_ADAPTIVE_BITMAP) {
-            context.grantUriPermission(context.getPackageName(),
-                    icon.getUri(), Intent.FLAG_GRANT_READ_URI_PERMISSION);
-        }
-        return icon.loadDrawable(context);
-    }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/bubbles/storage/BubbleEntity.kt b/packages/SystemUI/src/com/android/systemui/bubbles/storage/BubbleEntity.kt
index 24768cd..355c4b1 100644
--- a/packages/SystemUI/src/com/android/systemui/bubbles/storage/BubbleEntity.kt
+++ b/packages/SystemUI/src/com/android/systemui/bubbles/storage/BubbleEntity.kt
@@ -24,6 +24,5 @@
     val shortcutId: String,
     val key: String,
     val desiredHeight: Int,
-    @DimenRes val desiredHeightResId: Int,
-    val title: String? = null
+    @DimenRes val desiredHeightResId: Int
 )
diff --git a/packages/SystemUI/src/com/android/systemui/bubbles/storage/BubbleXmlHelper.kt b/packages/SystemUI/src/com/android/systemui/bubbles/storage/BubbleXmlHelper.kt
index 66fff338..a8faf25 100644
--- a/packages/SystemUI/src/com/android/systemui/bubbles/storage/BubbleXmlHelper.kt
+++ b/packages/SystemUI/src/com/android/systemui/bubbles/storage/BubbleXmlHelper.kt
@@ -33,7 +33,6 @@
 private const val ATTR_KEY = "key"
 private const val ATTR_DESIRED_HEIGHT = "h"
 private const val ATTR_DESIRED_HEIGHT_RES_ID = "hid"
-private const val ATTR_TITLE = "t"
 
 /**
  * Writes the bubbles in xml format into given output stream.
@@ -64,7 +63,6 @@
         serializer.attribute(null, ATTR_KEY, bubble.key)
         serializer.attribute(null, ATTR_DESIRED_HEIGHT, bubble.desiredHeight.toString())
         serializer.attribute(null, ATTR_DESIRED_HEIGHT_RES_ID, bubble.desiredHeightResId.toString())
-        bubble.title?.let { serializer.attribute(null, ATTR_TITLE, it) }
         serializer.endTag(null, TAG_BUBBLE)
     } catch (e: IOException) {
         throw RuntimeException(e)
@@ -94,8 +92,7 @@
             parser.getAttributeWithName(ATTR_SHORTCUT_ID) ?: return null,
             parser.getAttributeWithName(ATTR_KEY) ?: return null,
             parser.getAttributeWithName(ATTR_DESIRED_HEIGHT)?.toInt() ?: return null,
-            parser.getAttributeWithName(ATTR_DESIRED_HEIGHT_RES_ID)?.toInt() ?: return null,
-            parser.getAttributeWithName(ATTR_TITLE)
+            parser.getAttributeWithName(ATTR_DESIRED_HEIGHT_RES_ID)?.toInt() ?: return null
     )
 }
 
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 36398a6..59f8c4e 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/bubbles/BubbleControllerTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/bubbles/BubbleControllerTest.java
@@ -304,10 +304,6 @@
     public void testPromoteBubble_autoExpand() throws Exception {
         mBubbleController.updateBubble(mRow2.getEntry());
         mBubbleController.updateBubble(mRow.getEntry());
-        when(mNotificationEntryManager.getPendingOrActiveNotif(mRow.getEntry().getKey()))
-                .thenReturn(mRow.getEntry());
-        when(mNotificationEntryManager.getPendingOrActiveNotif(mRow2.getEntry().getKey()))
-                .thenReturn(mRow2.getEntry());
         mBubbleController.removeBubble(
                 mRow.getEntry().getKey(), BubbleController.DISMISS_USER_GESTURE);
 
@@ -335,10 +331,6 @@
         mBubbleController.updateBubble(mRow2.getEntry());
         mBubbleController.updateBubble(mRow.getEntry(), /* suppressFlyout */
                 false, /* showInShade */ true);
-        when(mNotificationEntryManager.getPendingOrActiveNotif(mRow.getEntry().getKey()))
-                .thenReturn(mRow.getEntry());
-        when(mNotificationEntryManager.getPendingOrActiveNotif(mRow2.getEntry().getKey()))
-                .thenReturn(mRow2.getEntry());
         mBubbleController.removeBubble(
                 mRow.getEntry().getKey(), BubbleController.DISMISS_USER_GESTURE);
 
@@ -441,16 +433,15 @@
         assertTrue(mSysUiStateBubblesExpanded);
 
         // Last added is the one that is expanded
-        assertEquals(mRow2.getEntry().getKey(), mBubbleData.getSelectedBubble().getKey());
+        assertEquals(mRow2.getEntry(), mBubbleData.getSelectedBubble().getEntry());
         assertTrue(mBubbleController.isBubbleNotificationSuppressedFromShade(
                 mRow2.getEntry()));
 
         // Switch which bubble is expanded
-        mBubbleData.setSelectedBubble(mBubbleData.getBubbleInStackWithKey(
-                mRow.getEntry().getKey()));
+        mBubbleData.setSelectedBubble(mBubbleData.getBubbleInStackWithKey(mRow.getEntry().getKey()));
         mBubbleData.setExpanded(true);
-        assertEquals(mRow.getEntry().getKey(), mBubbleData.getBubbleInStackWithKey(
-                stackView.getExpandedBubble().getKey()).getKey());
+        assertEquals(mRow.getEntry(),
+                mBubbleData.getBubbleInStackWithKey(stackView.getExpandedBubble().getKey()).getEntry());
         assertTrue(mBubbleController.isBubbleNotificationSuppressedFromShade(
                 mRow.getEntry()));
 
@@ -552,27 +543,27 @@
         verify(mBubbleExpandListener).onBubbleExpandChanged(true, mRow2.getEntry().getKey());
 
         // Last added is the one that is expanded
-        assertEquals(mRow2.getEntry().getKey(), mBubbleData.getBubbleInStackWithKey(
-                stackView.getExpandedBubble().getKey()).getKey());
+        assertEquals(mRow2.getEntry(),
+                mBubbleData.getBubbleInStackWithKey(stackView.getExpandedBubble().getKey()).getEntry());
         assertTrue(mBubbleController.isBubbleNotificationSuppressedFromShade(
                 mRow2.getEntry()));
 
         // Dismiss currently expanded
         mBubbleController.removeBubble(
-                mBubbleData.getBubbleInStackWithKey(
-                        stackView.getExpandedBubble().getKey()).getKey(),
+                mBubbleData.getBubbleInStackWithKey(stackView.getExpandedBubble().getKey())
+                        .getEntry().getKey(),
                 BubbleController.DISMISS_USER_GESTURE);
         verify(mBubbleExpandListener).onBubbleExpandChanged(false, mRow2.getEntry().getKey());
 
         // Make sure first bubble is selected
-        assertEquals(mRow.getEntry().getKey(), mBubbleData.getBubbleInStackWithKey(
-                stackView.getExpandedBubble().getKey()).getKey());
+        assertEquals(mRow.getEntry(),
+                mBubbleData.getBubbleInStackWithKey(stackView.getExpandedBubble().getKey()).getEntry());
         verify(mBubbleExpandListener).onBubbleExpandChanged(true, mRow.getEntry().getKey());
 
         // Dismiss that one
         mBubbleController.removeBubble(
-                mBubbleData.getBubbleInStackWithKey(
-                        stackView.getExpandedBubble().getKey()).getKey(),
+                mBubbleData.getBubbleInStackWithKey(stackView.getExpandedBubble().getKey())
+                        .getEntry().getKey(),
                 BubbleController.DISMISS_USER_GESTURE);
 
         // Make sure state changes and collapse happens
@@ -848,12 +839,6 @@
                 mRow2.getEntry(), /* suppressFlyout */ false, /* showInShade */ false);
         mBubbleController.updateBubble(
                 mRow3.getEntry(), /* suppressFlyout */ false, /* showInShade */ false);
-        when(mNotificationEntryManager.getPendingOrActiveNotif(mRow.getEntry().getKey()))
-                .thenReturn(mRow.getEntry());
-        when(mNotificationEntryManager.getPendingOrActiveNotif(mRow2.getEntry().getKey()))
-                .thenReturn(mRow2.getEntry());
-        when(mNotificationEntryManager.getPendingOrActiveNotif(mRow3.getEntry().getKey()))
-                .thenReturn(mRow3.getEntry());
         assertEquals(mBubbleData.getBubbles().size(), 3);
 
         mBubbleData.setMaxOverflowBubbles(1);
@@ -923,8 +908,6 @@
         // GIVEN a group summary with a bubble child
         ExpandableNotificationRow groupSummary = mNotificationTestHelper.createGroup(0);
         ExpandableNotificationRow groupedBubble = mNotificationTestHelper.createBubbleInGroup();
-        when(mNotificationEntryManager.getPendingOrActiveNotif(groupedBubble.getEntry().getKey()))
-                .thenReturn(groupedBubble.getEntry());
         mEntryListener.onPendingEntryAdded(groupedBubble.getEntry());
         groupSummary.addChildNotification(groupedBubble);
         assertTrue(mBubbleData.hasBubbleInStackWithKey(groupedBubble.getEntry().getKey()));
@@ -944,8 +927,6 @@
         ExpandableNotificationRow groupSummary = mNotificationTestHelper.createGroup(0);
         ExpandableNotificationRow groupedBubble = mNotificationTestHelper.createBubbleInGroup();
         mEntryListener.onPendingEntryAdded(groupedBubble.getEntry());
-        when(mNotificationEntryManager.getPendingOrActiveNotif(groupedBubble.getEntry().getKey()))
-                .thenReturn(groupedBubble.getEntry());
         groupSummary.addChildNotification(groupedBubble);
         assertTrue(mBubbleData.hasBubbleInStackWithKey(groupedBubble.getEntry().getKey()));
 
@@ -967,8 +948,6 @@
         // GIVEN a group summary with two (non-bubble) children and one bubble child
         ExpandableNotificationRow groupSummary = mNotificationTestHelper.createGroup(2);
         ExpandableNotificationRow groupedBubble = mNotificationTestHelper.createBubbleInGroup();
-        when(mNotificationEntryManager.getPendingOrActiveNotif(groupedBubble.getEntry().getKey()))
-                .thenReturn(groupedBubble.getEntry());
         mEntryListener.onPendingEntryAdded(groupedBubble.getEntry());
         groupSummary.addChildNotification(groupedBubble);
 
diff --git a/packages/SystemUI/tests/src/com/android/systemui/bubbles/BubbleTest.java b/packages/SystemUI/tests/src/com/android/systemui/bubbles/BubbleTest.java
index be03923..72f816f 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/bubbles/BubbleTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/bubbles/BubbleTest.java
@@ -86,7 +86,8 @@
         final String msg = "Hello there!";
         doReturn(Notification.Style.class).when(mNotif).getNotificationStyle();
         mExtras.putCharSequence(Notification.EXTRA_TEXT, msg);
-        assertEquals(msg, BubbleViewInfoTask.extractFlyoutMessage(mEntry).message);
+        assertEquals(msg, BubbleViewInfoTask.extractFlyoutMessage(mContext,
+                mEntry).message);
     }
 
     @Test
@@ -97,7 +98,8 @@
         mExtras.putCharSequence(Notification.EXTRA_BIG_TEXT, msg);
 
         // Should be big text, not the small text.
-        assertEquals(msg, BubbleViewInfoTask.extractFlyoutMessage(mEntry).message);
+        assertEquals(msg, BubbleViewInfoTask.extractFlyoutMessage(mContext,
+                mEntry).message);
     }
 
     @Test
@@ -105,7 +107,8 @@
         doReturn(Notification.MediaStyle.class).when(mNotif).getNotificationStyle();
 
         // Media notifs don't get update messages.
-        assertNull(BubbleViewInfoTask.extractFlyoutMessage(mEntry).message);
+        assertNull(BubbleViewInfoTask.extractFlyoutMessage(mContext,
+                mEntry).message);
     }
 
     @Test
@@ -121,7 +124,7 @@
 
         // Should be the last one only.
         assertEquals("Really? I prefer them that way.",
-                BubbleViewInfoTask.extractFlyoutMessage(mEntry).message);
+                BubbleViewInfoTask.extractFlyoutMessage(mContext, mEntry).message);
     }
 
     @Test
@@ -136,8 +139,11 @@
                                 "Oh, hello!", 0, "Mady").toBundle()});
 
         // Should be the last one only.
-        assertEquals("Oh, hello!", BubbleViewInfoTask.extractFlyoutMessage(mEntry).message);
-        assertEquals("Mady", BubbleViewInfoTask.extractFlyoutMessage(mEntry).senderName);
+        assertEquals("Oh, hello!",
+                BubbleViewInfoTask.extractFlyoutMessage(mContext, mEntry).message);
+        assertEquals("Mady",
+                BubbleViewInfoTask.extractFlyoutMessage(mContext,
+                        mEntry).senderName);
     }
 
     @Test
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 1c70db3..58e06b5 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/bubbles/NewNotifPipelineBubbleControllerTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/bubbles/NewNotifPipelineBubbleControllerTest.java
@@ -302,8 +302,6 @@
     public void testRemoveBubble_withDismissedNotif_notInOverflow() {
         mEntryListener.onEntryAdded(mRow.getEntry());
         mBubbleController.updateBubble(mRow.getEntry());
-        when(mNotificationEntryManager.getPendingOrActiveNotif(mRow.getEntry().getKey()))
-                .thenReturn(mRow.getEntry());
 
         assertTrue(mBubbleController.hasBubbles());
         assertFalse(mBubbleController.isBubbleNotificationSuppressedFromShade(mRow.getEntry()));
@@ -390,14 +388,14 @@
                 true, mRow.getEntry().getKey());
 
         // Last added is the one that is expanded
-        assertEquals(mRow2.getEntry().getKey(), mBubbleData.getSelectedBubble().getKey());
+        assertEquals(mRow2.getEntry(), mBubbleData.getSelectedBubble().getEntry());
         assertTrue(mBubbleController.isBubbleNotificationSuppressedFromShade(mRow2.getEntry()));
 
         // Switch which bubble is expanded
         mBubbleData.setSelectedBubble(mBubbleData.getBubbleInStackWithKey(mRow.getEntry().getKey()));
         mBubbleData.setExpanded(true);
-        assertEquals(mRow.getEntry().getKey(), mBubbleData.getBubbleInStackWithKey(
-                stackView.getExpandedBubble().getKey()).getKey());
+        assertEquals(mRow.getEntry(),
+                mBubbleData.getBubbleInStackWithKey(stackView.getExpandedBubble().getKey()).getEntry());
         assertTrue(mBubbleController.isBubbleNotificationSuppressedFromShade(
                 mRow.getEntry()));
 
@@ -490,27 +488,27 @@
         verify(mBubbleExpandListener).onBubbleExpandChanged(true, mRow2.getEntry().getKey());
 
         // Last added is the one that is expanded
-        assertEquals(mRow2.getEntry().getKey(), mBubbleData.getBubbleInStackWithKey(
-                stackView.getExpandedBubble().getKey()).getKey());
+        assertEquals(mRow2.getEntry(),
+                mBubbleData.getBubbleInStackWithKey(stackView.getExpandedBubble().getKey()).getEntry());
         assertTrue(mBubbleController.isBubbleNotificationSuppressedFromShade(
                 mRow2.getEntry()));
 
         // Dismiss currently expanded
         mBubbleController.removeBubble(
                 mBubbleData.getBubbleInStackWithKey(
-                        stackView.getExpandedBubble().getKey()).getKey(),
+                        stackView.getExpandedBubble().getKey()).getEntry().getKey(),
                 BubbleController.DISMISS_USER_GESTURE);
         verify(mBubbleExpandListener).onBubbleExpandChanged(false, mRow2.getEntry().getKey());
 
         // Make sure first bubble is selected
-        assertEquals(mRow.getEntry().getKey(), mBubbleData.getBubbleInStackWithKey(
-                stackView.getExpandedBubble().getKey()).getKey());
+        assertEquals(mRow.getEntry(),
+                mBubbleData.getBubbleInStackWithKey(stackView.getExpandedBubble().getKey()).getEntry());
         verify(mBubbleExpandListener).onBubbleExpandChanged(true, mRow.getEntry().getKey());
 
         // Dismiss that one
         mBubbleController.removeBubble(
                 mBubbleData.getBubbleInStackWithKey(
-                        stackView.getExpandedBubble().getKey()).getKey(),
+                        stackView.getExpandedBubble().getKey()).getEntry().getKey(),
                 BubbleController.DISMISS_USER_GESTURE);
 
         // Make sure state changes and collapse happens
@@ -769,8 +767,6 @@
         ExpandableNotificationRow groupSummary = mNotificationTestHelper.createGroup(0);
         ExpandableNotificationRow groupedBubble = mNotificationTestHelper.createBubbleInGroup();
         mEntryListener.onEntryAdded(groupedBubble.getEntry());
-        when(mNotificationEntryManager.getPendingOrActiveNotif(groupedBubble.getEntry().getKey()))
-                .thenReturn(groupedBubble.getEntry());
         groupSummary.addChildNotification(groupedBubble);
         assertTrue(mBubbleData.hasBubbleInStackWithKey(groupedBubble.getEntry().getKey()));
 
@@ -789,8 +785,6 @@
         ExpandableNotificationRow groupSummary = mNotificationTestHelper.createGroup(0);
         ExpandableNotificationRow groupedBubble = mNotificationTestHelper.createBubbleInGroup();
         mEntryListener.onEntryAdded(groupedBubble.getEntry());
-        when(mNotificationEntryManager.getPendingOrActiveNotif(groupedBubble.getEntry().getKey()))
-                .thenReturn(groupedBubble.getEntry());
         groupSummary.addChildNotification(groupedBubble);
         assertTrue(mBubbleData.hasBubbleInStackWithKey(groupedBubble.getEntry().getKey()));
 
@@ -813,8 +807,6 @@
         ExpandableNotificationRow groupSummary = mNotificationTestHelper.createGroup(2);
         ExpandableNotificationRow groupedBubble = mNotificationTestHelper.createBubbleInGroup();
         mEntryListener.onEntryAdded(groupedBubble.getEntry());
-        when(mNotificationEntryManager.getPendingOrActiveNotif(groupedBubble.getEntry().getKey()))
-                .thenReturn(groupedBubble.getEntry());
         groupSummary.addChildNotification(groupedBubble);
 
         // WHEN the summary is dismissed
diff --git a/packages/SystemUI/tests/src/com/android/systemui/bubbles/storage/BubblePersistentRepositoryTest.kt b/packages/SystemUI/tests/src/com/android/systemui/bubbles/storage/BubblePersistentRepositoryTest.kt
index 9b8fd11..1d02b8d 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/bubbles/storage/BubblePersistentRepositoryTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/bubbles/storage/BubblePersistentRepositoryTest.kt
@@ -32,7 +32,7 @@
 
     private val bubbles = listOf(
             BubbleEntity(0, "com.example.messenger", "shortcut-1", "key-1", 120, 0),
-            BubbleEntity(10, "com.example.chat", "alice and bob", "key-2", 0, 16537428, "title"),
+            BubbleEntity(10, "com.example.chat", "alice and bob", "key-2", 0, 16537428),
             BubbleEntity(0, "com.example.messenger", "shortcut-2", "key-3", 120, 0)
     )
     private lateinit var repository: BubblePersistentRepository
diff --git a/packages/SystemUI/tests/src/com/android/systemui/bubbles/storage/BubbleVolatileRepositoryTest.kt b/packages/SystemUI/tests/src/com/android/systemui/bubbles/storage/BubbleVolatileRepositoryTest.kt
index 76c5833..f9d611c 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/bubbles/storage/BubbleVolatileRepositoryTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/bubbles/storage/BubbleVolatileRepositoryTest.kt
@@ -37,10 +37,9 @@
     private val user0 = UserHandle.of(0)
     private val user10 = UserHandle.of(10)
 
-    private val bubble1 = BubbleEntity(0, "com.example.messenger", "shortcut-1", "key-1", 120, 0)
-    private val bubble2 = BubbleEntity(10, "com.example.chat", "alice and bob",
-            "key-2", 0, 16537428, "title")
-    private val bubble3 = BubbleEntity(0, "com.example.messenger", "shortcut-2", "key-3", 120, 0)
+    private val bubble1 = BubbleEntity(0, PKG_MESSENGER, "shortcut-1", "k1", 120, 0)
+    private val bubble2 = BubbleEntity(10, PKG_CHAT, "alice and bob", "k2", 0, 16537428)
+    private val bubble3 = BubbleEntity(0, PKG_MESSENGER, "shortcut-2", "k3", 120, 0)
 
     private val bubbles = listOf(bubble1, bubble2, bubble3)
 
diff --git a/packages/SystemUI/tests/src/com/android/systemui/bubbles/storage/BubbleXmlHelperTest.kt b/packages/SystemUI/tests/src/com/android/systemui/bubbles/storage/BubbleXmlHelperTest.kt
index 81687c7..4946787 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/bubbles/storage/BubbleXmlHelperTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/bubbles/storage/BubbleXmlHelperTest.kt
@@ -31,17 +31,17 @@
 class BubbleXmlHelperTest : SysuiTestCase() {
 
     private val bubbles = listOf(
-            BubbleEntity(0, "com.example.messenger", "shortcut-1", "k1", 120, 0),
-            BubbleEntity(10, "com.example.chat", "alice and bob", "k2", 0, 16537428, "title"),
-            BubbleEntity(0, "com.example.messenger", "shortcut-2", "k3", 120, 0)
+        BubbleEntity(0, "com.example.messenger", "shortcut-1", "k1", 120, 0),
+        BubbleEntity(10, "com.example.chat", "alice and bob", "k2", 0, 16537428),
+        BubbleEntity(0, "com.example.messenger", "shortcut-2", "k3", 120, 0)
     )
 
     @Test
     fun testWriteXml() {
         val expectedEntries = """
-<bb uid="0" pkg="com.example.messenger" sid="shortcut-1" key="k1" h="120" hid="0" />
-<bb uid="10" pkg="com.example.chat" sid="alice and bob" key="k2" h="0" hid="16537428" t="title" />
-<bb uid="0" pkg="com.example.messenger" sid="shortcut-2" key="k3" h="120" hid="0" />
+            <bb uid="0" pkg="com.example.messenger" sid="shortcut-1" key="k1" h="120" hid="0" />
+            <bb uid="10" pkg="com.example.chat" sid="alice and bob" key="k2" h="0" hid="16537428" />
+            <bb uid="0" pkg="com.example.messenger" sid="shortcut-2" key="k3" h="120" hid="0" />
         """.trimIndent()
         ByteArrayOutputStream().use {
             writeXml(it, bubbles)
@@ -54,12 +54,12 @@
     @Test
     fun testReadXml() {
         val src = """
-<?xml version='1.0' encoding='utf-8' standalone='yes' ?>
-<bs>
-<bb uid="0" pkg="com.example.messenger" sid="shortcut-1" key="k1" h="120" hid="0" />
-<bb uid="10" pkg="com.example.chat" sid="alice and bob" key="k2" h="0" hid="16537428" t="title" />
-<bb uid="0" pkg="com.example.messenger" sid="shortcut-2" key="k3" h="120" hid="0" />
-</bs>
+            <?xml version='1.0' encoding='utf-8' standalone='yes' ?>
+            <bs>
+            <bb uid="0" pkg="com.example.messenger" sid="shortcut-1" key="k1" h="120" hid="0" />
+            <bb uid="10" pkg="com.example.chat" sid="alice and bob" key="k2" h="0" hid="16537428" />
+            <bb uid="0" pkg="com.example.messenger" sid="shortcut-2" key="k3" h="120" hid="0" />
+            </bs>
         """.trimIndent()
         val actual = readXml(ByteArrayInputStream(src.toByteArray(Charsets.UTF_8)))
         assertEquals("failed parsing bubbles from xml\n$src", bubbles, actual)