Merge "Add cached state to prevent blocking calls from SysUI->Shell" into sc-dev

6 files changed, 332 insertions, 183 deletions

diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/Bubble.java b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/Bubble.java
index afb40d1ff95c..9d65d28b21b4 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/Bubble.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/Bubble.java
@@ -284,6 +284,15 @@ public class Bubble implements BubbleViewProvider {
         return mTitle;
     }
 
+    /**
+     * @return the ShortcutInfo id if it exists, or the metadata shortcut id otherwise.
+     */
+    String getShortcutId() {
+        return getShortcutInfo() != null
+                ? getShortcutInfo().getId()
+                : getMetadataShortcutId();
+    }
+
     String getMetadataShortcutId() {
         return mMetadataShortcutId;
     }
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleController.java b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleController.java
index dfd878f63283..202d9f013cc0 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleController.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleController.java
@@ -93,6 +93,7 @@ import java.io.FileDescriptor;
 import java.io.PrintWriter;
 import java.util.ArrayList;
 import java.util.HashMap;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Objects;
 import java.util.concurrent.Executor;
@@ -265,15 +266,7 @@ public class BubbleController {
 
     public void initialize() {
         mBubbleData.setListener(mBubbleDataListener);
-        mBubbleData.setSuppressionChangedListener(bubble -> {
-            // Make sure NoMan knows suppression state so that anyone querying it can tell.
-            try {
-                mBarService.onBubbleNotificationSuppressionChanged(bubble.getKey(),
-                        !bubble.showInShade(), bubble.isSuppressed());
-            } catch (RemoteException e) {
-                // Bad things have happened
-            }
-        });
+        mBubbleData.setSuppressionChangedListener(this::onBubbleNotificationSuppressionChanged);
 
         mBubbleData.setPendingIntentCancelledListener(bubble -> {
             if (bubble.getBubbleIntent() == null) {
@@ -401,6 +394,11 @@ public class BubbleController {
         return mImpl;
     }
 
+    @VisibleForTesting
+    public BubblesImpl.CachedState getImplCachedState() {
+        return mImpl.mCachedState;
+    }
+
     public ShellExecutor getMainExecutor() {
         return mMainExecutor;
     }
@@ -500,6 +498,18 @@ public class BubbleController {
         updateStack();
     }
 
+    @VisibleForTesting
+    public void onBubbleNotificationSuppressionChanged(Bubble bubble) {
+        // Make sure NoMan knows suppression state so that anyone querying it can tell.
+        try {
+            mBarService.onBubbleNotificationSuppressionChanged(bubble.getKey(),
+                    !bubble.showInShade(), bubble.isSuppressed());
+        } catch (RemoteException e) {
+            // Bad things have happened
+        }
+        mImpl.mCachedState.updateBubbleSuppressedState(bubble);
+    }
+
     /** Called when the current user changes. */
     @VisibleForTesting
     public void onUserChanged(int newUserId) {
@@ -808,11 +818,6 @@ public class BubbleController {
         }
     }
 
-    private boolean isBubbleExpanded(String key) {
-        return isStackExpanded() && mBubbleData != null && mBubbleData.getSelectedBubble() != null
-                && mBubbleData.getSelectedBubble().getKey().equals(key);
-    }
-
     /** Promote the provided bubble from the overflow view. */
     public void promoteBubbleFromOverflow(Bubble bubble) {
         mLogger.log(bubble, BubbleLogger.Event.BUBBLE_OVERFLOW_REMOVE_BACK_TO_STACK);
@@ -1191,6 +1196,9 @@ public class BubbleController {
 
             mSysuiProxy.notifyInvalidateNotifications("BubbleData.Listener.applyUpdate");
             updateStack();
+
+            // Update the cached state for queries from SysUI
+            mImpl.mCachedState.update(update);
         }
     };
 
@@ -1364,25 +1372,124 @@ public class BubbleController {
     }
 
     private class BubblesImpl implements Bubbles {
+        // Up-to-date cached state of bubbles data for SysUI to query from the calling thread
+        @VisibleForTesting
+        public class CachedState {
+            private boolean mIsStackExpanded;
+            private String mSelectedBubbleKey;
+            private HashSet<String> mSuppressedBubbleKeys = new HashSet<>();
+            private HashMap<String, String> mSuppressedGroupToNotifKeys = new HashMap<>();
+            private HashMap<String, Bubble> mShortcutIdToBubble = new HashMap<>();
+
+            private ArrayList<Bubble> mTmpBubbles = new ArrayList<>();
+
+            /**
+             * Updates the cached state based on the last full BubbleData change.
+             */
+            synchronized void update(BubbleData.Update update) {
+                if (update.selectionChanged) {
+                    mSelectedBubbleKey = update.selectedBubble != null
+                            ? update.selectedBubble.getKey()
+                            : null;
+                }
+                if (update.expandedChanged) {
+                    mIsStackExpanded = update.expanded;
+                }
+                if (update.suppressedSummaryChanged) {
+                    String summaryKey =
+                            mBubbleData.getSummaryKey(update.suppressedSummaryGroup);
+                    if (summaryKey != null) {
+                        mSuppressedGroupToNotifKeys.put(update.suppressedSummaryGroup, summaryKey);
+                    } else {
+                        mSuppressedGroupToNotifKeys.remove(update.suppressedSummaryGroup);
+                    }
+                }
+
+                mTmpBubbles.clear();
+                mTmpBubbles.addAll(update.bubbles);
+                mTmpBubbles.addAll(update.overflowBubbles);
+
+                mSuppressedBubbleKeys.clear();
+                mShortcutIdToBubble.clear();
+                for (Bubble b : mTmpBubbles) {
+                    mShortcutIdToBubble.put(b.getShortcutId(), b);
+                    updateBubbleSuppressedState(b);
+                }
+            }
+
+            /**
+             * Updates a specific bubble suppressed state.  This is used mainly because notification
+             * suppression changes don't go through the same BubbleData update mechanism.
+             */
+            synchronized void updateBubbleSuppressedState(Bubble b) {
+                if (!b.showInShade()) {
+                    mSuppressedBubbleKeys.add(b.getKey());
+                } else {
+                    mSuppressedBubbleKeys.remove(b.getKey());
+                }
+            }
+
+            public synchronized boolean isStackExpanded() {
+                return mIsStackExpanded;
+            }
+
+            public synchronized boolean isBubbleExpanded(String key) {
+                return mIsStackExpanded && key.equals(mSelectedBubbleKey);
+            }
+
+            public synchronized boolean isBubbleNotificationSuppressedFromShade(String key,
+                    String groupKey) {
+                return mSuppressedBubbleKeys.contains(key)
+                        || (mSuppressedGroupToNotifKeys.containsKey(groupKey)
+                                && key.equals(mSuppressedGroupToNotifKeys.get(groupKey)));
+            }
+
+            @Nullable
+            public synchronized Bubble getBubbleWithShortcutId(String id) {
+                return mShortcutIdToBubble.get(id);
+            }
+
+            synchronized void dump(PrintWriter pw) {
+                pw.println("BubbleImpl.CachedState state:");
+
+                pw.println("mIsStackExpanded: " + mIsStackExpanded);
+                pw.println("mSelectedBubbleKey: " + mSelectedBubbleKey);
+
+                pw.print("mSuppressedBubbleKeys: ");
+                pw.println(mSuppressedBubbleKeys.size());
+                for (String key : mSuppressedBubbleKeys) {
+                    pw.println("   suppressing: " + key);
+                }
+
+                pw.print("mSuppressedGroupToNotifKeys: ");
+                pw.println(mSuppressedGroupToNotifKeys.size());
+                for (String key : mSuppressedGroupToNotifKeys.keySet()) {
+                    pw.println("   suppressing: " + key);
+                }
+            }
+        }
+
+        private CachedState mCachedState = new CachedState();
+
         @Override
         public boolean isBubbleNotificationSuppressedFromShade(String key, String groupKey) {
-            return mMainExecutor.executeBlockingForResult(() -> {
-                return BubbleController.this.isBubbleNotificationSuppressedFromShade(key, groupKey);
-            }, Boolean.class);
+            return mCachedState.isBubbleNotificationSuppressedFromShade(key, groupKey);
         }
 
         @Override
         public boolean isBubbleExpanded(String key) {
-            return mMainExecutor.executeBlockingForResult(() -> {
-                return BubbleController.this.isBubbleExpanded(key);
-            }, Boolean.class);
+            return mCachedState.isBubbleExpanded(key);
         }
 
         @Override
         public boolean isStackExpanded() {
-            return mMainExecutor.executeBlockingForResult(() -> {
-                return BubbleController.this.isStackExpanded();
-            }, Boolean.class);
+            return mCachedState.isStackExpanded();
+        }
+
+        @Override
+        @Nullable
+        public Bubble getBubbleWithShortcutId(String shortcutId) {
+            return mCachedState.getBubbleWithShortcutId(shortcutId);
         }
 
         @Override
@@ -1425,14 +1532,6 @@ public class BubbleController {
         }
 
         @Override
-        @Nullable
-        public Bubble getBubbleWithShortcutId(String shortcutId) {
-            return mMainExecutor.executeBlockingForResult(() -> {
-                return BubbleController.this.mBubbleData.getAnyBubbleWithShortcutId(shortcutId);
-            }, Bubble.class);
-        }
-
-        @Override
         public void onTaskbarChanged(Bundle b) {
             mMainExecutor.execute(() -> {
                 BubbleController.this.onTaskbarChanged(b);
@@ -1555,6 +1654,7 @@ public class BubbleController {
             try {
                 mMainExecutor.executeBlocking(() -> {
                     BubbleController.this.dump(fd, pw, args);
+                    mCachedState.dump(pw);
                 });
             } catch (InterruptedException e) {
                 Slog.e(TAG, "Failed to dump BubbleController in 2s");
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleData.java b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleData.java
index 6f5cfd114688..d73ce6951e6d 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleData.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleData.java
@@ -73,6 +73,7 @@ public class BubbleData {
         boolean expandedChanged;
         boolean selectionChanged;
         boolean orderChanged;
+        boolean suppressedSummaryChanged;
         boolean expanded;
         @Nullable BubbleViewProvider selectedBubble;
         @Nullable Bubble addedBubble;
@@ -81,6 +82,7 @@ public class BubbleData {
         @Nullable Bubble removedOverflowBubble;
         @Nullable Bubble suppressedBubble;
         @Nullable Bubble unsuppressedBubble;
+        @Nullable String suppressedSummaryGroup;
         // Pair with Bubble and @DismissReason Integer
         final List<Pair<Bubble, Integer>> removedBubbles = new ArrayList<>();
 
@@ -103,7 +105,9 @@ public class BubbleData {
                     || removedOverflowBubble != null
                     || orderChanged
                     || suppressedBubble != null
-                    || unsuppressedBubble != null;
+                    || unsuppressedBubble != null
+                    || suppressedSummaryChanged
+                    || suppressedSummaryGroup != null;
         }
 
         void bubbleRemoved(Bubble bubbleToRemove, @DismissReason int reason) {
@@ -380,6 +384,9 @@ public class BubbleData {
      */
     void addSummaryToSuppress(String groupKey, String notifKey) {
         mSuppressedGroupKeys.put(groupKey, notifKey);
+        mStateChange.suppressedSummaryChanged = true;
+        mStateChange.suppressedSummaryGroup = groupKey;
+        dispatchPendingChanges();
     }
 
     /**
@@ -397,6 +404,9 @@ public class BubbleData {
      */
     void removeSuppressedSummary(String groupKey) {
         mSuppressedGroupKeys.remove(groupKey);
+        mStateChange.suppressedSummaryChanged = true;
+        mStateChange.suppressedSummaryGroup = groupKey;
+        dispatchPendingChanges();
     }
 
     /**
diff --git a/packages/SystemUI/src/com/android/systemui/wmshell/BubblesManager.java b/packages/SystemUI/src/com/android/systemui/wmshell/BubblesManager.java
index 5441bd4c958d..a29a638b91a8 100644
--- a/packages/SystemUI/src/com/android/systemui/wmshell/BubblesManager.java
+++ b/packages/SystemUI/src/com/android/systemui/wmshell/BubblesManager.java
@@ -617,7 +617,7 @@ public class BubblesManager implements Dumpable {
      * cancel it (and hence the bubbles associated with it).
      *
      * @return true if we want to intercept the dismissal of the entry, else false.
-     * @see Bubbles#handleDismissalInterception(BubbleEntry, List, IntConsumer)
+     * @see Bubbles#handleDismissalInterception(BubbleEntry, List, IntConsumer, Executor)
      */
     public boolean handleDismissalInterception(NotificationEntry entry) {
         if (entry == null) {
diff --git a/packages/SystemUI/tests/src/com/android/systemui/wmshell/BubblesTest.java b/packages/SystemUI/tests/src/com/android/systemui/wmshell/BubblesTest.java
index 496976e8f55c..9fa35f8bd9f6 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/wmshell/BubblesTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/wmshell/BubblesTest.java
@@ -23,6 +23,8 @@ import static android.service.notification.NotificationListenerService.REASON_CA
 import static android.service.notification.NotificationListenerService.REASON_CANCEL_ALL;
 import static android.service.notification.NotificationListenerService.REASON_GROUP_SUMMARY_CANCELED;
 
+import static com.android.dx.mockito.inline.extended.ExtendedMockito.spyOn;
+
 import static com.google.common.truth.Truth.assertThat;
 
 import static org.junit.Assert.assertEquals;
@@ -37,6 +39,7 @@ import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.atLeastOnce;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
@@ -324,6 +327,7 @@ public class BubblesTest extends SysuiTestCase {
                 syncExecutor,
                 mock(Handler.class));
         mBubbleController.setExpandListener(mBubbleExpandListener);
+        spyOn(mBubbleController);
 
         mBubblesManager = new BubblesManager(
                 mContext,
@@ -470,7 +474,7 @@ public class BubblesTest extends SysuiTestCase {
 
     @Test
     public void testExpandCollapseStack() {
-        assertFalse(mBubbleController.isStackExpanded());
+        assertStackCollapsed();
 
         // Mark it as a bubble and add it explicitly
         mEntryListener.onPendingEntryAdded(mRow);
@@ -478,25 +482,23 @@ public class BubblesTest extends SysuiTestCase {
 
         // We should have bubbles & their notifs should not be suppressed
         assertTrue(mBubbleController.hasBubbles());
-        assertFalse(mBubbleController.isBubbleNotificationSuppressedFromShade(
-                mBubbleEntry.getKey(), mBubbleEntry.getGroupKey()));
+        assertBubbleNotificationNotSuppressedFromShade(mBubbleEntry);
 
         // Expand the stack
         BubbleStackView stackView = mBubbleController.getStackView();
         mBubbleData.setExpanded(true);
-        assertTrue(mBubbleController.isStackExpanded());
+        assertStackExpanded();
         verify(mBubbleExpandListener).onBubbleExpandChanged(true, mRow.getKey());
 
         assertTrue(mSysUiStateBubblesExpanded);
 
         // Make sure the notif is suppressed
-        assertTrue(mBubbleController.isBubbleNotificationSuppressedFromShade(
-                mBubbleEntry.getKey(), mBubbleEntry.getGroupKey()));
+        assertBubbleNotificationSuppressedFromShade(mBubbleEntry);
 
         // Collapse
         mBubbleController.collapseStack();
         verify(mBubbleExpandListener).onBubbleExpandChanged(false, mRow.getKey());
-        assertFalse(mBubbleController.isStackExpanded());
+        assertStackCollapsed();
 
         assertFalse(mSysUiStateBubblesExpanded);
     }
@@ -512,15 +514,13 @@ public class BubblesTest extends SysuiTestCase {
 
         // We should have bubbles & their notifs should not be suppressed
         assertTrue(mBubbleController.hasBubbles());
-        assertFalse(mBubbleController.isBubbleNotificationSuppressedFromShade(
-                mBubbleEntry.getKey(), mBubbleEntry.getGroupKey()));
-        assertFalse(mBubbleController.isBubbleNotificationSuppressedFromShade(
-                mBubbleEntry2.getKey(), mBubbleEntry2.getGroupKey()));
+        assertBubbleNotificationNotSuppressedFromShade(mBubbleEntry);
+        assertBubbleNotificationNotSuppressedFromShade(mBubbleEntry2);
 
         // Expand
         BubbleStackView stackView = mBubbleController.getStackView();
         mBubbleData.setExpanded(true);
-        assertTrue(mBubbleController.isStackExpanded());
+        assertStackExpanded();
         verify(mBubbleExpandListener, atLeastOnce()).onBubbleExpandChanged(
                 true, mRow2.getKey());
 
@@ -528,8 +528,7 @@ public class BubblesTest extends SysuiTestCase {
 
         // Last added is the one that is expanded
         assertEquals(mRow2.getKey(), mBubbleData.getSelectedBubble().getKey());
-        assertTrue(mBubbleController.isBubbleNotificationSuppressedFromShade(
-                mBubbleEntry2.getKey(), mBubbleEntry2.getGroupKey()));
+        assertBubbleNotificationSuppressedFromShade(mBubbleEntry2);
 
         // Switch which bubble is expanded
         mBubbleData.setSelectedBubble(mBubbleData.getBubbleInStackWithKey(
@@ -537,8 +536,7 @@ public class BubblesTest extends SysuiTestCase {
         mBubbleData.setExpanded(true);
         assertEquals(mRow.getKey(), mBubbleData.getBubbleInStackWithKey(
                 stackView.getExpandedBubble().getKey()).getKey());
-        assertTrue(mBubbleController.isBubbleNotificationSuppressedFromShade(
-                mBubbleEntry.getKey(), mBubbleEntry.getGroupKey()));
+        assertBubbleNotificationSuppressedFromShade(mBubbleEntry);
 
         // collapse for previous bubble
         verify(mBubbleExpandListener, atLeastOnce()).onBubbleExpandChanged(
@@ -549,7 +547,7 @@ public class BubblesTest extends SysuiTestCase {
 
         // Collapse
         mBubbleController.collapseStack();
-        assertFalse(mBubbleController.isStackExpanded());
+        assertStackCollapsed();
 
         assertFalse(mSysUiStateBubblesExpanded);
     }
@@ -562,22 +560,20 @@ public class BubblesTest extends SysuiTestCase {
 
         // We should have bubbles & their notifs should not be suppressed
         assertTrue(mBubbleController.hasBubbles());
-        assertFalse(mBubbleController.isBubbleNotificationSuppressedFromShade(
-                mBubbleEntry.getKey(), mBubbleEntry.getGroupKey()));
+        assertBubbleNotificationNotSuppressedFromShade(mBubbleEntry);
 
         mTestableLooper.processAllMessages();
         assertTrue(mBubbleData.getBubbleInStackWithKey(mRow.getKey()).showDot());
 
         // Expand
         mBubbleData.setExpanded(true);
-        assertTrue(mBubbleController.isStackExpanded());
+        assertStackExpanded();
         verify(mBubbleExpandListener).onBubbleExpandChanged(true, mRow.getKey());
 
         assertTrue(mSysUiStateBubblesExpanded);
 
         // Notif is suppressed after expansion
-        assertTrue(mBubbleController.isBubbleNotificationSuppressedFromShade(
-                mBubbleEntry.getKey(), mBubbleEntry.getGroupKey()));
+        assertBubbleNotificationSuppressedFromShade(mBubbleEntry);
         // Notif shouldn't show dot after expansion
         assertFalse(mBubbleData.getBubbleInStackWithKey(mRow.getKey()).showDot());
     }
@@ -590,22 +586,20 @@ public class BubblesTest extends SysuiTestCase {
 
         // We should have bubbles & their notifs should not be suppressed
         assertTrue(mBubbleController.hasBubbles());
-        assertFalse(mBubbleController.isBubbleNotificationSuppressedFromShade(
-                mBubbleEntry.getKey(), mBubbleEntry.getGroupKey()));
+        assertBubbleNotificationNotSuppressedFromShade(mBubbleEntry);
 
         mTestableLooper.processAllMessages();
         assertTrue(mBubbleData.getBubbleInStackWithKey(mRow.getKey()).showDot());
 
         // Expand
         mBubbleData.setExpanded(true);
-        assertTrue(mBubbleController.isStackExpanded());
+        assertStackExpanded();
         verify(mBubbleExpandListener).onBubbleExpandChanged(true, mRow.getKey());
 
         assertTrue(mSysUiStateBubblesExpanded);
 
         // Notif is suppressed after expansion
-        assertTrue(mBubbleController.isBubbleNotificationSuppressedFromShade(
-                mBubbleEntry.getKey(), mBubbleEntry.getGroupKey()));
+        assertBubbleNotificationSuppressedFromShade(mBubbleEntry);
         // Notif shouldn't show dot after expansion
         assertFalse(mBubbleData.getBubbleInStackWithKey(mRow.getKey()).showDot());
 
@@ -614,8 +608,7 @@ public class BubblesTest extends SysuiTestCase {
 
         // Nothing should have changed
         // Notif is suppressed after expansion
-        assertTrue(mBubbleController.isBubbleNotificationSuppressedFromShade(
-                mBubbleEntry.getKey(), mBubbleEntry.getGroupKey()));
+        assertBubbleNotificationSuppressedFromShade(mBubbleEntry);
         // Notif shouldn't show dot after expansion
         assertFalse(mBubbleData.getBubbleInStackWithKey(mRow.getKey()).showDot());
     }
@@ -634,14 +627,13 @@ public class BubblesTest extends SysuiTestCase {
 
         assertTrue(mSysUiStateBubblesExpanded);
 
-        assertTrue(mBubbleController.isStackExpanded());
+        assertStackExpanded();
         verify(mBubbleExpandListener).onBubbleExpandChanged(true, mRow2.getKey());
 
         // Last added is the one that is expanded
         assertEquals(mRow2.getKey(), mBubbleData.getBubbleInStackWithKey(
                 stackView.getExpandedBubble().getKey()).getKey());
-        assertTrue(mBubbleController.isBubbleNotificationSuppressedFromShade(
-                mBubbleEntry2.getKey(), mBubbleEntry2.getGroupKey()));
+        assertBubbleNotificationSuppressedFromShade(mBubbleEntry2);
 
         // Dismiss currently expanded
         mBubbleController.removeBubble(
@@ -679,7 +671,7 @@ public class BubblesTest extends SysuiTestCase {
         mBubbleData.setExpanded(true);
 
         assertTrue(mSysUiStateBubblesExpanded);
-        assertTrue(mBubbleController.isStackExpanded());
+        assertStackExpanded();
         verify(mBubbleExpandListener).onBubbleExpandChanged(true, mRow.getKey());
 
         // Block the bubble so it won't be in the overflow
@@ -698,7 +690,7 @@ public class BubblesTest extends SysuiTestCase {
 
     @Test
     public void testAutoExpand_fails_noFlag() {
-        assertFalse(mBubbleController.isStackExpanded());
+        assertStackCollapsed();
         setMetadataFlags(mRow,
                 Notification.BubbleMetadata.FLAG_AUTO_EXPAND_BUBBLE, false /* enableFlag */);
 
@@ -709,7 +701,7 @@ public class BubblesTest extends SysuiTestCase {
         // Expansion shouldn't change
         verify(mBubbleExpandListener, never()).onBubbleExpandChanged(false /* expanded */,
                 mRow.getKey());
-        assertFalse(mBubbleController.isStackExpanded());
+        assertStackCollapsed();
 
         assertFalse(mSysUiStateBubblesExpanded);
     }
@@ -726,7 +718,7 @@ public class BubblesTest extends SysuiTestCase {
         // Expansion should change
         verify(mBubbleExpandListener).onBubbleExpandChanged(true /* expanded */,
                 mRow.getKey());
-        assertTrue(mBubbleController.isStackExpanded());
+        assertStackExpanded();
 
         assertTrue(mSysUiStateBubblesExpanded);
     }
@@ -741,8 +733,7 @@ public class BubblesTest extends SysuiTestCase {
         mBubbleController.updateBubble(mBubbleEntry);
 
         // Notif should be suppressed because we were foreground
-        assertTrue(mBubbleController.isBubbleNotificationSuppressedFromShade(
-                mBubbleEntry.getKey(), mBubbleEntry.getGroupKey()));
+        assertBubbleNotificationSuppressedFromShade(mBubbleEntry);
         // Dot + flyout is hidden because notif is suppressed
         assertFalse(mBubbleData.getBubbleInStackWithKey(mRow.getKey()).showDot());
         assertFalse(mBubbleData.getBubbleInStackWithKey(mRow.getKey()).showFlyout());
@@ -755,8 +746,7 @@ public class BubblesTest extends SysuiTestCase {
         mBubbleController.updateBubble(mBubbleEntry);
 
         // Should not be suppressed
-        assertFalse(mBubbleController.isBubbleNotificationSuppressedFromShade(
-                mBubbleEntry.getKey(), mBubbleEntry.getGroupKey()));
+        assertBubbleNotificationNotSuppressedFromShade(mBubbleEntry);
         // Should show dot
         assertTrue(mBubbleData.getBubbleInStackWithKey(mRow.getKey()).showDot());
 
@@ -766,8 +756,7 @@ public class BubblesTest extends SysuiTestCase {
         mBubbleController.updateBubble(mBubbleEntry);
 
         // Notif should be suppressed
-        assertTrue(mBubbleController.isBubbleNotificationSuppressedFromShade(
-                mBubbleEntry.getKey(), mBubbleEntry.getGroupKey()));
+        assertBubbleNotificationSuppressedFromShade(mBubbleEntry);
         // Dot + flyout is hidden because notif is suppressed
         assertFalse(mBubbleData.getBubbleInStackWithKey(mRow.getKey()).showDot());
         assertFalse(mBubbleData.getBubbleInStackWithKey(mRow.getKey()).showFlyout());
@@ -792,8 +781,7 @@ public class BubblesTest extends SysuiTestCase {
     @Test
     public void testMarkNewNotificationAsShowInShade() {
         mEntryListener.onPendingEntryAdded(mRow);
-        assertFalse(mBubbleController.isBubbleNotificationSuppressedFromShade(
-                mBubbleEntry.getKey(), mBubbleEntry.getGroupKey()));
+        assertBubbleNotificationNotSuppressedFromShade(mBubbleEntry);
 
         mTestableLooper.processAllMessages();
         assertTrue(mBubbleData.getBubbleInStackWithKey(mRow.getKey()).showDot());
@@ -878,8 +866,7 @@ public class BubblesTest extends SysuiTestCase {
         mBubbleController.updateBubble(mBubbleEntry);
 
         assertTrue(mBubbleController.hasBubbles());
-        assertFalse(mBubbleController.isBubbleNotificationSuppressedFromShade(
-                mBubbleEntry.getKey(), mBubbleEntry.getGroupKey()));
+        assertBubbleNotificationNotSuppressedFromShade(mBubbleEntry);
 
         boolean intercepted = mRemoveInterceptor.onNotificationRemoveRequested(
                 mRow.getKey(), mRow, REASON_CANCEL_ALL);
@@ -887,8 +874,7 @@ public class BubblesTest extends SysuiTestCase {
         // Intercept!
         assertTrue(intercepted);
         // Should update show in shade state
-        assertTrue(mBubbleController.isBubbleNotificationSuppressedFromShade(
-                mBubbleEntry.getKey(), mBubbleEntry.getGroupKey()));
+        assertBubbleNotificationSuppressedFromShade(mBubbleEntry);
     }
 
     @Test
@@ -897,8 +883,7 @@ public class BubblesTest extends SysuiTestCase {
         mBubbleController.updateBubble(mBubbleEntry);
 
         assertTrue(mBubbleController.hasBubbles());
-        assertFalse(mBubbleController.isBubbleNotificationSuppressedFromShade(
-                mBubbleEntry.getKey(), mBubbleEntry.getGroupKey()));
+        assertBubbleNotificationNotSuppressedFromShade(mBubbleEntry);
 
         boolean intercepted = mRemoveInterceptor.onNotificationRemoveRequested(
                 mRow.getKey(), mRow, REASON_CANCEL);
@@ -906,8 +891,7 @@ public class BubblesTest extends SysuiTestCase {
         // Intercept!
         assertTrue(intercepted);
         // Should update show in shade state
-        assertTrue(mBubbleController.isBubbleNotificationSuppressedFromShade(
-                mBubbleEntry.getKey(), mBubbleEntry.getGroupKey()));
+        assertBubbleNotificationSuppressedFromShade(mBubbleEntry);
     }
 
     @Test
@@ -916,8 +900,7 @@ public class BubblesTest extends SysuiTestCase {
         mEntryListener.onPendingEntryAdded(mRow);
         mBubbleController.updateBubble(mBubbleEntry);
         assertTrue(mBubbleController.hasBubbles());
-        assertFalse(mBubbleController.isBubbleNotificationSuppressedFromShade(
-                mBubbleEntry.getKey(), mBubbleEntry.getGroupKey()));
+        assertBubbleNotificationNotSuppressedFromShade(mBubbleEntry);
 
         // Dismiss the bubble into overflow.
         mBubbleController.removeBubble(
@@ -938,8 +921,7 @@ public class BubblesTest extends SysuiTestCase {
         mBubbleController.updateBubble(mBubbleEntry);
 
         assertTrue(mBubbleController.hasBubbles());
-        assertFalse(mBubbleController.isBubbleNotificationSuppressedFromShade(
-                mBubbleEntry.getKey(), mBubbleEntry.getGroupKey()));
+        assertBubbleNotificationNotSuppressedFromShade(mBubbleEntry);
 
         mBubbleController.removeBubble(
                 mRow.getKey(), Bubbles.DISMISS_NO_LONGER_BUBBLE);
@@ -985,48 +967,36 @@ public class BubblesTest extends SysuiTestCase {
 
     @Test
     public void testNotifyShadeSuppressionChange_notificationDismiss() {
-        Bubbles.SuppressionChangedListener listener =
-                mock(Bubbles.SuppressionChangedListener.class);
-        mBubbleData.setSuppressionChangedListener(listener);
-
         mEntryListener.onPendingEntryAdded(mRow);
 
         assertTrue(mBubbleController.hasBubbles());
-        assertFalse(mBubbleController.isBubbleNotificationSuppressedFromShade(
-                mBubbleEntry.getKey(), mBubbleEntry.getGroupKey()));
+        assertBubbleNotificationNotSuppressedFromShade(mBubbleEntry);
 
         mRemoveInterceptor.onNotificationRemoveRequested(
                 mRow.getKey(), mRow, REASON_CANCEL);
 
         // Should update show in shade state
-        assertTrue(mBubbleController.isBubbleNotificationSuppressedFromShade(
-                mBubbleEntry.getKey(), mBubbleEntry.getGroupKey()));
+        assertBubbleNotificationSuppressedFromShade(mBubbleEntry);
 
         // Should notify delegate that shade state changed
-        verify(listener).onBubbleNotificationSuppressionChange(
+        verify(mBubbleController).onBubbleNotificationSuppressionChanged(
                 mBubbleData.getBubbleInStackWithKey(mRow.getKey()));
     }
 
     @Test
     public void testNotifyShadeSuppressionChange_bubbleExpanded() {
-        Bubbles.SuppressionChangedListener listener =
-                mock(Bubbles.SuppressionChangedListener.class);
-        mBubbleData.setSuppressionChangedListener(listener);
-
         mEntryListener.onPendingEntryAdded(mRow);
 
         assertTrue(mBubbleController.hasBubbles());
-        assertFalse(mBubbleController.isBubbleNotificationSuppressedFromShade(
-                mBubbleEntry.getKey(), mBubbleEntry.getGroupKey()));
+        assertBubbleNotificationNotSuppressedFromShade(mBubbleEntry);
 
         mBubbleData.setExpanded(true);
 
         // Once a bubble is expanded the notif is suppressed
-        assertTrue(mBubbleController.isBubbleNotificationSuppressedFromShade(
-                mBubbleEntry.getKey(), mBubbleEntry.getGroupKey()));
+        assertBubbleNotificationSuppressedFromShade(mBubbleEntry);
 
         // Should notify delegate that shade state changed
-        verify(listener).onBubbleNotificationSuppressionChange(
+        verify(mBubbleController).onBubbleNotificationSuppressionChanged(
                 mBubbleData.getBubbleInStackWithKey(mRow.getKey()));
     }
 
@@ -1046,7 +1016,11 @@ public class BubblesTest extends SysuiTestCase {
 
         // THEN the summary and bubbled child are suppressed from the shade
         assertTrue(mBubbleController.isBubbleNotificationSuppressedFromShade(
-                groupedBubble.getEntry().getKey(), groupSummary.getEntry().getSbn().getGroupKey()));
+                groupedBubble.getEntry().getKey(),
+                groupSummary.getEntry().getSbn().getGroupKey()));
+        assertTrue(mBubbleController.getImplCachedState().isBubbleNotificationSuppressedFromShade(
+                groupedBubble.getEntry().getKey(),
+                groupSummary.getEntry().getSbn().getGroupKey()));
         assertTrue(mBubbleData.isSummarySuppressed(groupSummary.getEntry().getSbn().getGroupKey()));
     }
 
@@ -1102,6 +1076,9 @@ public class BubblesTest extends SysuiTestCase {
         assertTrue(mBubbleController.isBubbleNotificationSuppressedFromShade(
                 groupedBubble.getEntry().getKey(),
                 groupedBubble.getEntry().getSbn().getGroupKey()));
+        assertTrue(mBubbleController.getImplCachedState().isBubbleNotificationSuppressedFromShade(
+                groupedBubble.getEntry().getKey(),
+                groupedBubble.getEntry().getSbn().getGroupKey()));
 
         // THEN the summary is removed from GroupManager
         verify(mNotificationGroupManager, times(1)).onEntryRemoved(groupSummary.getEntry());
@@ -1257,4 +1234,42 @@ public class BubblesTest extends SysuiTestCase {
                 Icon.createWithResource(mContext, R.drawable.bubble_ic_create_bubble))
                 .build();
     }
+
+    /**
+     * Asserts that the bubble stack is expanded and also validates the cached state is updated.
+     */
+    private void assertStackExpanded() {
+        assertTrue(mBubbleController.isStackExpanded());
+        assertTrue(mBubbleController.getImplCachedState().isStackExpanded());
+    }
+
+    /**
+     * Asserts that the bubble stack is collapsed and also validates the cached state is updated.
+     */
+    private void assertStackCollapsed() {
+        assertFalse(mBubbleController.isStackExpanded());
+        assertFalse(mBubbleController.getImplCachedState().isStackExpanded());
+    }
+
+    /**
+     * Asserts that a bubble notification is suppressed from the shade and also validates the cached
+     * state is updated.
+     */
+    private void assertBubbleNotificationSuppressedFromShade(BubbleEntry entry) {
+        assertTrue(mBubbleController.isBubbleNotificationSuppressedFromShade(
+                entry.getKey(), entry.getGroupKey()));
+        assertTrue(mBubbleController.getImplCachedState().isBubbleNotificationSuppressedFromShade(
+                entry.getKey(), entry.getGroupKey()));
+    }
+
+    /**
+     * Asserts that a bubble notification is not suppressed from the shade and also validates the
+     * cached state is updated.
+     */
+    private void assertBubbleNotificationNotSuppressedFromShade(BubbleEntry entry) {
+        assertFalse(mBubbleController.isBubbleNotificationSuppressedFromShade(
+                entry.getKey(), entry.getGroupKey()));
+        assertFalse(mBubbleController.getImplCachedState().isBubbleNotificationSuppressedFromShade(
+                entry.getKey(), entry.getGroupKey()));
+    }
 }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/wmshell/NewNotifPipelineBubblesTest.java b/packages/SystemUI/tests/src/com/android/systemui/wmshell/NewNotifPipelineBubblesTest.java
index 9114b7a35fd2..55cc8bb1780d 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/wmshell/NewNotifPipelineBubblesTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/wmshell/NewNotifPipelineBubblesTest.java
@@ -19,6 +19,8 @@ package com.android.systemui.wmshell;
 import static android.app.Notification.FLAG_BUBBLE;
 import static android.service.notification.NotificationListenerService.REASON_GROUP_SUMMARY_CANCELED;
 
+import static com.android.dx.mockito.inline.extended.ExtendedMockito.spyOn;
+
 import static com.google.common.truth.Truth.assertThat;
 
 import static org.junit.Assert.assertEquals;
@@ -268,6 +270,7 @@ public class NewNotifPipelineBubblesTest extends SysuiTestCase {
                 syncExecutor,
                 mock(Handler.class));
         mBubbleController.setExpandListener(mBubbleExpandListener);
+        spyOn(mBubbleController);
 
         mBubblesManager = new BubblesManager(
                 mContext,
@@ -328,8 +331,7 @@ public class NewNotifPipelineBubblesTest extends SysuiTestCase {
         mBubbleController.updateBubble(mBubbleEntry);
 
         assertTrue(mBubbleController.hasBubbles());
-        assertFalse(mBubbleController.isBubbleNotificationSuppressedFromShade(
-                mBubbleEntry.getKey(), mBubbleEntry.getGroupKey()));
+        assertBubbleNotificationNotSuppressedFromShade(mBubbleEntry);
 
         // Make it look like dismissed notif
         mBubbleData.getBubbleInStackWithKey(mRow.getKey()).setSuppressNotification(true);
@@ -352,8 +354,7 @@ public class NewNotifPipelineBubblesTest extends SysuiTestCase {
                 .thenReturn(mRow);
 
         assertTrue(mBubbleController.hasBubbles());
-        assertFalse(mBubbleController.isBubbleNotificationSuppressedFromShade(
-                mBubbleEntry.getKey(), mBubbleEntry.getGroupKey()));
+        assertBubbleNotificationNotSuppressedFromShade(mBubbleEntry);
 
         // Make it look like dismissed notif
         mBubbleData.getBubbleInStackWithKey(mRow.getKey()).setSuppressNotification(true);
@@ -388,7 +389,7 @@ public class NewNotifPipelineBubblesTest extends SysuiTestCase {
 
     @Test
     public void testExpandCollapseStack() {
-        assertFalse(mBubbleController.isStackExpanded());
+        assertStackCollapsed();
 
         // Mark it as a bubble and add it explicitly
         mEntryListener.onEntryAdded(mRow);
@@ -396,23 +397,20 @@ public class NewNotifPipelineBubblesTest extends SysuiTestCase {
 
         // We should have bubbles & their notifs should not be suppressed
         assertTrue(mBubbleController.hasBubbles());
-        assertFalse(mBubbleController.isBubbleNotificationSuppressedFromShade(
-                mBubbleEntry.getKey(), mBubbleEntry.getGroupKey()));
+        assertBubbleNotificationNotSuppressedFromShade(mBubbleEntry);
 
         // Expand the stack
-        BubbleStackView stackView = mBubbleController.getStackView();
         mBubbleData.setExpanded(true);
-        assertTrue(mBubbleController.isStackExpanded());
+        assertStackExpanded();
         verify(mBubbleExpandListener).onBubbleExpandChanged(true, mRow.getKey());
 
         // Make sure the notif is suppressed
-        assertTrue(mBubbleController.isBubbleNotificationSuppressedFromShade(
-                mBubbleEntry.getKey(), mBubbleEntry.getGroupKey()));
+        assertBubbleNotificationSuppressedFromShade(mBubbleEntry);
 
         // Collapse
         mBubbleController.collapseStack();
         verify(mBubbleExpandListener).onBubbleExpandChanged(false, mRow.getKey());
-        assertFalse(mBubbleController.isStackExpanded());
+        assertStackCollapsed();
     }
 
     @Test
@@ -426,22 +424,19 @@ public class NewNotifPipelineBubblesTest extends SysuiTestCase {
 
         // We should have bubbles & their notifs should not be suppressed
         assertTrue(mBubbleController.hasBubbles());
-        assertFalse(mBubbleController.isBubbleNotificationSuppressedFromShade(
-                mBubbleEntry.getKey(), mBubbleEntry.getGroupKey()));
-        assertFalse(mBubbleController.isBubbleNotificationSuppressedFromShade(
-                mBubbleEntry.getKey(), mBubbleEntry.getGroupKey()));
+        assertBubbleNotificationNotSuppressedFromShade(mBubbleEntry);
+        assertBubbleNotificationNotSuppressedFromShade(mBubbleEntry2);
 
         // Expand
         BubbleStackView stackView = mBubbleController.getStackView();
         mBubbleData.setExpanded(true);
-        assertTrue(mBubbleController.isStackExpanded());
+        assertStackExpanded();
         verify(mBubbleExpandListener, atLeastOnce()).onBubbleExpandChanged(
                 true, mRow2.getKey());
 
         // Last added is the one that is expanded
         assertEquals(mRow2.getKey(), mBubbleData.getSelectedBubble().getKey());
-        assertTrue(mBubbleController.isBubbleNotificationSuppressedFromShade(
-                mBubbleEntry2.getKey(), mBubbleEntry2.getGroupKey()));
+        assertBubbleNotificationSuppressedFromShade(mBubbleEntry2);
 
         // Switch which bubble is expanded
         mBubbleData.setSelectedBubble(mBubbleData.getBubbleInStackWithKey(
@@ -449,8 +444,7 @@ public class NewNotifPipelineBubblesTest extends SysuiTestCase {
         mBubbleData.setExpanded(true);
         assertEquals(mRow.getKey(), mBubbleData.getBubbleInStackWithKey(
                 stackView.getExpandedBubble().getKey()).getKey());
-        assertTrue(mBubbleController.isBubbleNotificationSuppressedFromShade(
-                mBubbleEntry.getKey(), mBubbleEntry.getGroupKey()));
+        assertBubbleNotificationSuppressedFromShade(mBubbleEntry);
 
         // collapse for previous bubble
         verify(mBubbleExpandListener, atLeastOnce()).onBubbleExpandChanged(
@@ -462,7 +456,7 @@ public class NewNotifPipelineBubblesTest extends SysuiTestCase {
 
         // Collapse
         mBubbleController.collapseStack();
-        assertFalse(mBubbleController.isStackExpanded());
+        assertStackCollapsed();
     }
 
     @Test
@@ -473,20 +467,18 @@ public class NewNotifPipelineBubblesTest extends SysuiTestCase {
 
         // We should have bubbles & their notifs should not be suppressed
         assertTrue(mBubbleController.hasBubbles());
-        assertFalse(mBubbleController.isBubbleNotificationSuppressedFromShade(
-                mBubbleEntry.getKey(), mBubbleEntry.getGroupKey()));
+        assertBubbleNotificationNotSuppressedFromShade(mBubbleEntry);
 
         mTestableLooper.processAllMessages();
         assertTrue(mBubbleData.getBubbleInStackWithKey(mRow.getKey()).showDot());
 
         // Expand
         mBubbleData.setExpanded(true);
-        assertTrue(mBubbleController.isStackExpanded());
+        assertStackExpanded();
         verify(mBubbleExpandListener).onBubbleExpandChanged(true, mRow.getKey());
 
         // Notif is suppressed after expansion
-        assertTrue(mBubbleController.isBubbleNotificationSuppressedFromShade(
-                mBubbleEntry.getKey(), mBubbleEntry.getGroupKey()));
+        assertBubbleNotificationSuppressedFromShade(mBubbleEntry);
         // Notif shouldn't show dot after expansion
         assertFalse(mBubbleData.getBubbleInStackWithKey(mRow.getKey()).showDot());
     }
@@ -499,20 +491,18 @@ public class NewNotifPipelineBubblesTest extends SysuiTestCase {
 
         // We should have bubbles & their notifs should not be suppressed
         assertTrue(mBubbleController.hasBubbles());
-        assertFalse(mBubbleController.isBubbleNotificationSuppressedFromShade(
-                mBubbleEntry.getKey(), mBubbleEntry.getGroupKey()));
+        assertBubbleNotificationNotSuppressedFromShade(mBubbleEntry);
 
         mTestableLooper.processAllMessages();
         assertTrue(mBubbleData.getBubbleInStackWithKey(mRow.getKey()).showDot());
 
         // Expand
         mBubbleData.setExpanded(true);
-        assertTrue(mBubbleController.isStackExpanded());
+        assertStackExpanded();
         verify(mBubbleExpandListener).onBubbleExpandChanged(true, mRow.getKey());
 
         // Notif is suppressed after expansion
-        assertTrue(mBubbleController.isBubbleNotificationSuppressedFromShade(
-                mBubbleEntry.getKey(), mBubbleEntry.getGroupKey()));
+        assertBubbleNotificationSuppressedFromShade(mBubbleEntry);
         // Notif shouldn't show dot after expansion
         assertFalse(mBubbleData.getBubbleInStackWithKey(mRow.getKey()).showDot());
 
@@ -521,8 +511,7 @@ public class NewNotifPipelineBubblesTest extends SysuiTestCase {
 
         // Nothing should have changed
         // Notif is suppressed after expansion
-        assertTrue(mBubbleController.isBubbleNotificationSuppressedFromShade(
-                mBubbleEntry.getKey(), mBubbleEntry.getGroupKey()));
+        assertBubbleNotificationSuppressedFromShade(mBubbleEntry);
         // Notif shouldn't show dot after expansion
         assertFalse(mBubbleData.getBubbleInStackWithKey(mRow.getKey()).showDot());
     }
@@ -539,14 +528,13 @@ public class NewNotifPipelineBubblesTest extends SysuiTestCase {
         BubbleStackView stackView = mBubbleController.getStackView();
         mBubbleData.setExpanded(true);
 
-        assertTrue(mBubbleController.isStackExpanded());
+        assertStackExpanded();
         verify(mBubbleExpandListener).onBubbleExpandChanged(true, mRow2.getKey());
 
         // Last added is the one that is expanded
         assertEquals(mRow2.getKey(), mBubbleData.getBubbleInStackWithKey(
                 stackView.getExpandedBubble().getKey()).getKey());
-        assertTrue(mBubbleController.isBubbleNotificationSuppressedFromShade(
-                mBubbleEntry2.getKey(), mBubbleEntry2.getGroupKey()));
+        assertBubbleNotificationSuppressedFromShade(mBubbleEntry2);
 
         // Dismiss currently expanded
         mBubbleController.removeBubble(
@@ -582,7 +570,7 @@ public class NewNotifPipelineBubblesTest extends SysuiTestCase {
         BubbleStackView stackView = mBubbleController.getStackView();
         mBubbleData.setExpanded(true);
 
-        assertTrue(mBubbleController.isStackExpanded());
+        assertStackExpanded();
         verify(mBubbleExpandListener).onBubbleExpandChanged(true, mRow.getKey());
 
         // Block the bubble so it won't be in the overflow
@@ -601,7 +589,7 @@ public class NewNotifPipelineBubblesTest extends SysuiTestCase {
 
     @Test
     public void testAutoExpand_fails_noFlag() {
-        assertFalse(mBubbleController.isStackExpanded());
+        assertStackCollapsed();
         setMetadataFlags(mRow,
                 Notification.BubbleMetadata.FLAG_AUTO_EXPAND_BUBBLE, false /* enableFlag */);
 
@@ -612,7 +600,7 @@ public class NewNotifPipelineBubblesTest extends SysuiTestCase {
         // Expansion shouldn't change
         verify(mBubbleExpandListener, never()).onBubbleExpandChanged(false /* expanded */,
                 mRow.getKey());
-        assertFalse(mBubbleController.isStackExpanded());
+        assertStackCollapsed();
     }
 
     @Test
@@ -627,7 +615,7 @@ public class NewNotifPipelineBubblesTest extends SysuiTestCase {
         // Expansion should change
         verify(mBubbleExpandListener).onBubbleExpandChanged(true /* expanded */,
                 mRow.getKey());
-        assertTrue(mBubbleController.isStackExpanded());
+        assertStackExpanded();
     }
 
     @Test
@@ -640,8 +628,7 @@ public class NewNotifPipelineBubblesTest extends SysuiTestCase {
         mBubbleController.updateBubble(mBubbleEntry);
 
         // Notif should be suppressed because we were foreground
-        assertTrue(mBubbleController.isBubbleNotificationSuppressedFromShade(
-                mBubbleEntry.getKey(), mBubbleEntry.getGroupKey()));
+        assertBubbleNotificationSuppressedFromShade(mBubbleEntry);
         // Dot + flyout is hidden because notif is suppressed
         assertFalse(mBubbleData.getBubbleInStackWithKey(mRow.getKey()).showDot());
         assertFalse(mBubbleData.getBubbleInStackWithKey(mRow.getKey()).showFlyout());
@@ -652,8 +639,7 @@ public class NewNotifPipelineBubblesTest extends SysuiTestCase {
         mBubbleController.updateBubble(mBubbleEntry);
 
         // Should not be suppressed
-        assertFalse(mBubbleController.isBubbleNotificationSuppressedFromShade(
-                mBubbleEntry.getKey(), mBubbleEntry.getGroupKey()));
+        assertBubbleNotificationNotSuppressedFromShade(mBubbleEntry);
         // Should show dot
         assertTrue(mBubbleData.getBubbleInStackWithKey(mRow.getKey()).showDot());
 
@@ -663,8 +649,7 @@ public class NewNotifPipelineBubblesTest extends SysuiTestCase {
         mBubbleController.updateBubble(mBubbleEntry);
 
         // Notif should be suppressed
-        assertTrue(mBubbleController.isBubbleNotificationSuppressedFromShade(
-                mBubbleEntry.getKey(), mBubbleEntry.getGroupKey()));
+        assertBubbleNotificationSuppressedFromShade(mBubbleEntry);
         // Dot + flyout is hidden because notif is suppressed
         assertFalse(mBubbleData.getBubbleInStackWithKey(mRow.getKey()).showDot());
         assertFalse(mBubbleData.getBubbleInStackWithKey(mRow.getKey()).showFlyout());
@@ -673,8 +658,7 @@ public class NewNotifPipelineBubblesTest extends SysuiTestCase {
     @Test
     public void testMarkNewNotificationAsShowInShade() {
         mEntryListener.onEntryAdded(mRow);
-        assertFalse(mBubbleController.isBubbleNotificationSuppressedFromShade(
-                mBubbleEntry.getKey(), mBubbleEntry.getGroupKey()));
+        assertBubbleNotificationNotSuppressedFromShade(mBubbleEntry);
 
         mTestableLooper.processAllMessages();
         assertTrue(mBubbleData.getBubbleInStackWithKey(mRow.getKey()).showDot());
@@ -745,16 +729,14 @@ public class NewNotifPipelineBubblesTest extends SysuiTestCase {
         mBubbleController.updateBubble(mBubbleEntry);
 
         assertTrue(mBubbleController.hasBubbles());
-        assertFalse(mBubbleController.isBubbleNotificationSuppressedFromShade(
-                mBubbleEntry.getKey(), mBubbleEntry.getGroupKey()));
+        assertBubbleNotificationNotSuppressedFromShade(mBubbleEntry);
 
         boolean intercepted = mBubblesManager.handleDismissalInterception(mRow);
 
         // Intercept!
         assertTrue(intercepted);
         // Should update show in shade state
-        assertTrue(mBubbleController.isBubbleNotificationSuppressedFromShade(
-                mBubbleEntry.getKey(), mBubbleEntry.getGroupKey()));
+        assertBubbleNotificationSuppressedFromShade(mBubbleEntry);
     }
 
     @Test
@@ -763,8 +745,7 @@ public class NewNotifPipelineBubblesTest extends SysuiTestCase {
         mBubbleController.updateBubble(mBubbleEntry);
 
         assertTrue(mBubbleController.hasBubbles());
-        assertFalse(mBubbleController.isBubbleNotificationSuppressedFromShade(
-                mBubbleEntry.getKey(), mBubbleEntry.getGroupKey()));
+        assertBubbleNotificationNotSuppressedFromShade(mBubbleEntry);
 
         // Dismiss the bubble
         mBubbleController.removeBubble(mRow.getKey(), Bubbles.DISMISS_USER_GESTURE);
@@ -783,8 +764,7 @@ public class NewNotifPipelineBubblesTest extends SysuiTestCase {
         mBubbleController.updateBubble(mBubbleEntry);
 
         assertTrue(mBubbleController.hasBubbles());
-        assertFalse(mBubbleController.isBubbleNotificationSuppressedFromShade(
-                mBubbleEntry.getKey(), mBubbleEntry.getGroupKey()));
+        assertBubbleNotificationNotSuppressedFromShade(mBubbleEntry);
 
         // Dismiss the bubble
         mBubbleController.removeBubble(mRow.getKey(), Bubbles.DISMISS_NOTIF_CANCEL);
@@ -799,47 +779,35 @@ public class NewNotifPipelineBubblesTest extends SysuiTestCase {
 
     @Test
     public void testNotifyShadeSuppressionChange_notificationDismiss() {
-        Bubbles.SuppressionChangedListener listener =
-                mock(Bubbles.SuppressionChangedListener.class);
-        mBubbleData.setSuppressionChangedListener(listener);
-
         mEntryListener.onEntryAdded(mRow);
 
         assertTrue(mBubbleController.hasBubbles());
-        assertFalse(mBubbleController.isBubbleNotificationSuppressedFromShade(
-                mBubbleEntry.getKey(), mBubbleEntry.getGroupKey()));
+        assertBubbleNotificationNotSuppressedFromShade(mBubbleEntry);
 
         mBubblesManager.handleDismissalInterception(mRow);
 
         // Should update show in shade state
-        assertTrue(mBubbleController.isBubbleNotificationSuppressedFromShade(
-                mBubbleEntry.getKey(), mBubbleEntry.getGroupKey()));
+        assertBubbleNotificationSuppressedFromShade(mBubbleEntry);
 
         // Should notify delegate that shade state changed
-        verify(listener).onBubbleNotificationSuppressionChange(
+        verify(mBubbleController).onBubbleNotificationSuppressionChanged(
                 mBubbleData.getBubbleInStackWithKey(mRow.getKey()));
     }
 
     @Test
     public void testNotifyShadeSuppressionChange_bubbleExpanded() {
-        Bubbles.SuppressionChangedListener listener =
-                mock(Bubbles.SuppressionChangedListener.class);
-        mBubbleData.setSuppressionChangedListener(listener);
-
         mEntryListener.onEntryAdded(mRow);
 
         assertTrue(mBubbleController.hasBubbles());
-        assertFalse(mBubbleController.isBubbleNotificationSuppressedFromShade(
-                mBubbleEntry.getKey(), mBubbleEntry.getGroupKey()));
+        assertBubbleNotificationNotSuppressedFromShade(mBubbleEntry);
 
         mBubbleData.setExpanded(true);
 
         // Once a bubble is expanded the notif is suppressed
-        assertTrue(mBubbleController.isBubbleNotificationSuppressedFromShade(
-                mBubbleEntry.getKey(), mBubbleEntry.getGroupKey()));
+        assertBubbleNotificationSuppressedFromShade(mBubbleEntry);
 
         // Should notify delegate that shade state changed
-        verify(listener).onBubbleNotificationSuppressionChange(
+        verify(mBubbleController).onBubbleNotificationSuppressionChanged(
                 mBubbleData.getBubbleInStackWithKey(mRow.getKey()));
     }
 
@@ -861,6 +829,9 @@ public class NewNotifPipelineBubblesTest extends SysuiTestCase {
         assertTrue(mBubbleController.isBubbleNotificationSuppressedFromShade(
                 groupedBubble.getEntry().getKey(),
                 groupedBubble.getEntry().getSbn().getGroupKey()));
+        assertTrue(mBubbleController.getImplCachedState().isBubbleNotificationSuppressedFromShade(
+                groupedBubble.getEntry().getKey(),
+                groupedBubble.getEntry().getSbn().getGroupKey()));
         assertTrue(mBubbleData.isSummarySuppressed(groupSummary.getEntry().getSbn().getGroupKey()));
     }
 
@@ -915,11 +886,17 @@ public class NewNotifPipelineBubblesTest extends SysuiTestCase {
         assertTrue(mBubbleController.isBubbleNotificationSuppressedFromShade(
                 groupedBubble.getEntry().getKey(),
                 groupedBubble.getEntry().getSbn().getGroupKey()));
+        assertTrue(mBubbleController.getImplCachedState().isBubbleNotificationSuppressedFromShade(
+                groupedBubble.getEntry().getKey(),
+                groupedBubble.getEntry().getSbn().getGroupKey()));
 
         // THEN the summary is also suppressed from the shade
         assertTrue(mBubbleController.isBubbleNotificationSuppressedFromShade(
                 groupSummary.getEntry().getKey(),
                 groupSummary.getEntry().getSbn().getGroupKey()));
+        assertTrue(mBubbleController.getImplCachedState().isBubbleNotificationSuppressedFromShade(
+                groupSummary.getEntry().getKey(),
+                groupSummary.getEntry().getSbn().getGroupKey()));
     }
 
     /**
@@ -938,4 +915,42 @@ public class NewNotifPipelineBubblesTest extends SysuiTestCase {
         }
         bubbleMetadata.setFlags(flags);
     }
+
+    /**
+     * Asserts that the bubble stack is expanded and also validates the cached state is updated.
+     */
+    private void assertStackExpanded() {
+        assertTrue(mBubbleController.isStackExpanded());
+        assertTrue(mBubbleController.getImplCachedState().isStackExpanded());
+    }
+
+    /**
+     * Asserts that the bubble stack is collapsed and also validates the cached state is updated.
+     */
+    private void assertStackCollapsed() {
+        assertFalse(mBubbleController.isStackExpanded());
+        assertFalse(mBubbleController.getImplCachedState().isStackExpanded());
+    }
+
+    /**
+     * Asserts that a bubble notification is suppressed from the shade and also validates the cached
+     * state is updated.
+     */
+    private void assertBubbleNotificationSuppressedFromShade(BubbleEntry entry) {
+        assertTrue(mBubbleController.isBubbleNotificationSuppressedFromShade(
+                entry.getKey(), entry.getGroupKey()));
+        assertTrue(mBubbleController.getImplCachedState().isBubbleNotificationSuppressedFromShade(
+                entry.getKey(), entry.getGroupKey()));
+    }
+
+    /**
+     * Asserts that a bubble notification is not suppressed from the shade and also validates the
+     * cached state is updated.
+     */
+    private void assertBubbleNotificationNotSuppressedFromShade(BubbleEntry entry) {
+        assertFalse(mBubbleController.isBubbleNotificationSuppressedFromShade(
+                entry.getKey(), entry.getGroupKey()));
+        assertFalse(mBubbleController.getImplCachedState().isBubbleNotificationSuppressedFromShade(
+                entry.getKey(), entry.getGroupKey()));
+    }
 }