2 files changed, 27 insertions, 17 deletions

diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/NotificationInfo.java b/packages/SystemUI/src/com/android/systemui/statusbar/NotificationInfo.java
index e52829a088e1..4f29b31fd941 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/NotificationInfo.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/NotificationInfo.java
@@ -102,8 +102,15 @@ public class NotificationInfo extends LinearLayout implements NotificationGuts.G
     };
 
     private OnClickListener mOnStopOrMinimizeNotifications = v -> {
-        mExitReason = NotificationCounters.BLOCKING_HELPER_STOP_NOTIFICATIONS;
-        swapContent(false);
+        Runnable saveImportance = () -> {
+            mExitReason = NotificationCounters.BLOCKING_HELPER_STOP_NOTIFICATIONS;
+            swapContent(false);
+        };
+        if (mCheckSaveListener != null) {
+            mCheckSaveListener.checkSave(saveImportance, mSbn);
+        } else {
+            saveImportance.run();
+        }
     };
 
     private OnClickListener mOnUndo = v -> {
@@ -300,15 +307,7 @@ public class NotificationInfo extends LinearLayout implements NotificationGuts.G
 
     private void saveImportance() {
         if (!mIsNonblockable) {
-            // Only go through the lock screen/bouncer if the user hit 'Stop notifications'.
-            // Otherwise, update the importance immediately.
-            if (mCheckSaveListener != null
-                    && NotificationCounters.BLOCKING_HELPER_STOP_NOTIFICATIONS.equals(
-                            mExitReason)) {
-                mCheckSaveListener.checkSave(this::updateImportance, mSbn);
-            } else {
-                updateImportance();
-            }
+            updateImportance();
         }
     }
 
@@ -513,6 +512,11 @@ public class NotificationInfo extends LinearLayout implements NotificationGuts.G
         return getHeight();
     }
 
+    @VisibleForTesting
+    public boolean isAnimating() {
+        return mExpandAnimation != null && mExpandAnimation.isRunning();
+    }
+
     /**
      * Runnable to either update the given channel (with a new importance value) or, if no channel
      * is provided, update notifications enabled state for the package.
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/NotificationInfoTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/NotificationInfoTest.java
index a72fed4a2b41..0017943cfc98 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/NotificationInfoTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/NotificationInfoTest.java
@@ -158,6 +158,11 @@ public class NotificationInfoTest extends SysuiTestCase {
         PollingCheck.waitFor(1000,
                 () -> VISIBLE == mNotificationInfo.findViewById(R.id.confirmation).getVisibility());
     }
+    private void ensureNoUndoButton() {
+        PollingCheck.waitFor(1000,
+                () -> GONE == mNotificationInfo.findViewById(R.id.confirmation).getVisibility()
+                        && !mNotificationInfo.isAnimating());
+    }
     private void waitForStopButton() {
         PollingCheck.waitFor(1000,
                 () -> VISIBLE == mNotificationInfo.findViewById(R.id.prompt).getVisibility());
@@ -567,9 +572,6 @@ public class NotificationInfoTest extends SysuiTestCase {
                 true /* isUserSentimentNegative */);
 
         mNotificationInfo.findViewById(R.id.block).performClick();
-        waitForUndoButton();
-        mNotificationInfo.handleCloseControls(true /* save */, false /* force */);
-
         mTestableLooper.processAllMessages();
         verify(listener).checkSave(any(Runnable.class), eq(mSbn));
     }
@@ -787,7 +789,7 @@ public class NotificationInfoTest extends SysuiTestCase {
     }
 
     @Test
-    public void testCloseControlsDoesNotUpdateIfCheckSaveListenerIsNoOp() throws Exception {
+    public void testBlockDoesNothingIfCheckSaveListenerIsNoOp() throws Exception {
         mNotificationChannel.setImportance(IMPORTANCE_LOW);
         mNotificationInfo.bindNotification(mMockPackageManager, mMockINotificationManager,
                 TEST_PACKAGE_NAME, mNotificationChannel, 1, mSbn,
@@ -795,10 +797,10 @@ public class NotificationInfoTest extends SysuiTestCase {
                 }, null, null, true);
 
         mNotificationInfo.findViewById(R.id.block).performClick();
-        waitForUndoButton();
+        mTestableLooper.processAllMessages();
+        ensureNoUndoButton();
         mNotificationInfo.handleCloseControls(true, false);
 
-        mTestableLooper.processAllMessages();
         verify(mMockINotificationManager, never()).updateNotificationChannelForPackage(
                 eq(TEST_PACKAGE_NAME), eq(TEST_UID), eq(mNotificationChannel));
     }
@@ -813,6 +815,10 @@ public class NotificationInfoTest extends SysuiTestCase {
                 }, null, null, false);
 
         mNotificationInfo.findViewById(R.id.block).performClick();
+        mTestableLooper.processAllMessages();
+        verify(mMockINotificationManager, never()).updateNotificationChannelForPackage(
+                eq(TEST_PACKAGE_NAME), eq(TEST_UID), eq(mNotificationChannel));
+
         waitForUndoButton();
         mNotificationInfo.handleCloseControls(true, false);