Merge "Check avalanche entry list when needed in BaseHeadsUpManager" into main

4 files changed, 84 insertions, 10 deletions

diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/policy/BaseHeadsUpManagerTest.java b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/policy/BaseHeadsUpManagerTest.java
index 0f66a93bcbec..3bfc046e46b4 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/policy/BaseHeadsUpManagerTest.java
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/policy/BaseHeadsUpManagerTest.java
@@ -38,6 +38,7 @@ import static org.mockito.Mockito.when;
 import android.app.Notification;
 import android.app.PendingIntent;
 import android.app.Person;
+import android.platform.test.annotations.EnableFlags;
 import android.platform.test.flag.junit.FlagsParameterization;
 import android.testing.TestableLooper;
 
@@ -150,6 +151,62 @@ public class BaseHeadsUpManagerTest extends SysuiTestCase {
     }
 
     @Test
+    public void testHasNotifications_headsUpManagerMapNotEmpty_true() {
+        final BaseHeadsUpManager bhum = createHeadsUpManager();
+        final NotificationEntry entry = HeadsUpManagerTestUtil.createEntry(/* id = */ 0, mContext);
+        bhum.showNotification(entry);
+
+        assertThat(bhum.mHeadsUpEntryMap).isNotEmpty();
+        assertThat(bhum.hasNotifications()).isTrue();
+    }
+
+    @Test
+    @EnableFlags(NotificationThrottleHun.FLAG_NAME)
+    public void testHasNotifications_avalancheMapNotEmpty_true() {
+        final BaseHeadsUpManager bhum = createHeadsUpManager();
+        final NotificationEntry notifEntry = HeadsUpManagerTestUtil.createEntry(/* id = */ 0,
+                mContext);
+        final BaseHeadsUpManager.HeadsUpEntry headsUpEntry = bhum.createHeadsUpEntry(notifEntry);
+        mAvalancheController.addToNext(headsUpEntry, () -> {});
+
+        assertThat(mAvalancheController.getWaitingEntryList()).isNotEmpty();
+        assertThat(bhum.hasNotifications()).isTrue();
+    }
+
+    @Test
+    @EnableFlags(NotificationThrottleHun.FLAG_NAME)
+    public void testHasNotifications_false() {
+        final BaseHeadsUpManager bhum = createHeadsUpManager();
+        assertThat(bhum.mHeadsUpEntryMap).isEmpty();
+        assertThat(mAvalancheController.getWaitingEntryList()).isEmpty();
+        assertThat(bhum.hasNotifications()).isFalse();
+    }
+
+    @Test
+    @EnableFlags(NotificationThrottleHun.FLAG_NAME)
+    public void testGetHeadsUpEntryList_includesAvalancheEntryList() {
+        final BaseHeadsUpManager bhum = createHeadsUpManager();
+        final NotificationEntry notifEntry = HeadsUpManagerTestUtil.createEntry(/* id = */ 0,
+                mContext);
+        final BaseHeadsUpManager.HeadsUpEntry headsUpEntry = bhum.createHeadsUpEntry(notifEntry);
+        mAvalancheController.addToNext(headsUpEntry, () -> {});
+
+        assertThat(bhum.getHeadsUpEntryList()).contains(headsUpEntry);
+    }
+
+    @Test
+    @EnableFlags(NotificationThrottleHun.FLAG_NAME)
+    public void testGetHeadsUpEntry_returnsAvalancheEntry() {
+        final BaseHeadsUpManager bhum = createHeadsUpManager();
+        final NotificationEntry notifEntry = HeadsUpManagerTestUtil.createEntry(/* id = */ 0,
+                mContext);
+        final BaseHeadsUpManager.HeadsUpEntry headsUpEntry = bhum.createHeadsUpEntry(notifEntry);
+        mAvalancheController.addToNext(headsUpEntry, () -> {});
+
+        assertThat(bhum.getHeadsUpEntry(notifEntry.getKey())).isEqualTo(headsUpEntry);
+    }
+
+    @Test
     public void testShowNotification_addsEntry() {
         final BaseHeadsUpManager alm = createHeadsUpManager();
         final NotificationEntry entry = HeadsUpManagerTestUtil.createEntry(/* id = */ 0, mContext);
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/phone/HeadsUpManagerPhone.java b/packages/SystemUI/src/com/android/systemui/statusbar/phone/HeadsUpManagerPhone.java
index ffc859eb9197..4bf122dd3b6a 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/phone/HeadsUpManagerPhone.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/phone/HeadsUpManagerPhone.java
@@ -86,6 +86,8 @@ public class HeadsUpManagerPhone extends BaseHeadsUpManager implements
     private final List<OnHeadsUpPhoneListenerChange> mHeadsUpPhoneListeners = new ArrayList<>();
     private final VisualStabilityProvider mVisualStabilityProvider;
 
+    private final AvalancheController mAvalancheController;
+
     // TODO(b/328393698) move the topHeadsUpRow logic to an interactor
     private final MutableStateFlow<HeadsUpRowRepository> mTopHeadsUpRow =
             StateFlowKt.MutableStateFlow(null);
@@ -155,6 +157,7 @@ public class HeadsUpManagerPhone extends BaseHeadsUpManager implements
         mBypassController = bypassController;
         mGroupMembershipManager = groupMembershipManager;
         mVisualStabilityProvider = visualStabilityProvider;
+        mAvalancheController = avalancheController;
 
         updateResources();
         configurationController.addCallback(new ConfigurationController.ConfigurationListener() {
@@ -653,9 +656,10 @@ public class HeadsUpManagerPhone extends BaseHeadsUpManager implements
             boolean wasKeyguard = mStatusBarState == StatusBarState.KEYGUARD;
             boolean isKeyguard = newState == StatusBarState.KEYGUARD;
             mStatusBarState = newState;
+
             if (wasKeyguard && !isKeyguard && mBypassController.getBypassEnabled()) {
                 ArrayList<String> keysToRemove = new ArrayList<>();
-                for (HeadsUpEntry entry : mHeadsUpEntryMap.values()) {
+                for (HeadsUpEntry entry : getHeadsUpEntryList()) {
                     if (entry.mEntry != null && entry.mEntry.isBubble() && !entry.isSticky()) {
                         keysToRemove.add(entry.mEntry.getKey());
                     }
@@ -671,7 +675,7 @@ public class HeadsUpManagerPhone extends BaseHeadsUpManager implements
             if (!isDozing) {
                 // Let's make sure all huns we got while dozing time out within the normal timeout
                 // duration. Otherwise they could get stuck for a very long time
-                for (HeadsUpEntry entry : mHeadsUpEntryMap.values()) {
+                for (HeadsUpEntry entry : getHeadsUpEntryList()) {
                     entry.updateEntry(true /* updatePostTime */, "onDozingChanged(false)");
                 }
             }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/policy/AvalancheController.kt b/packages/SystemUI/src/com/android/systemui/statusbar/policy/AvalancheController.kt
index 21691540fc56..eb09e6ef39cb 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/policy/AvalancheController.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/policy/AvalancheController.kt
@@ -254,6 +254,13 @@ constructor(
         return null
     }
 
+    fun getWaitingEntryList(): List<HeadsUpEntry> {
+        if (!NotificationThrottleHun.isEnabled) {
+            return mutableListOf()
+        }
+        return nextMap.keys.toList()
+    }
+
     private fun isShowing(entry: HeadsUpEntry): Boolean {
         return headsUpEntryShowing != null && entry.mEntry?.key == headsUpEntryShowing?.mEntry?.key
     }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/policy/BaseHeadsUpManager.java b/packages/SystemUI/src/com/android/systemui/statusbar/policy/BaseHeadsUpManager.java
index 2ee98bb61d80..4bd868179faf 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/policy/BaseHeadsUpManager.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/policy/BaseHeadsUpManager.java
@@ -279,7 +279,6 @@ public abstract class BaseHeadsUpManager implements HeadsUpManager {
      * Returns the entry if it is managed by this manager.
      * @param key key of notification
      * @return the entry
-     * TODO(b/315362456) See if caller needs to check AvalancheController waiting entries
      */
     @Nullable
     public NotificationEntry getEntry(@NonNull String key) {
@@ -294,8 +293,13 @@ public abstract class BaseHeadsUpManager implements HeadsUpManager {
     @NonNull
     @Override
     public Stream<NotificationEntry> getAllEntries() {
-        // TODO(b/315362456) See if callers need to check AvalancheController
-        return mHeadsUpEntryMap.values().stream().map(headsUpEntry -> headsUpEntry.mEntry);
+        return getHeadsUpEntryList().stream().map(headsUpEntry -> headsUpEntry.mEntry);
+    }
+
+    public List<HeadsUpEntry> getHeadsUpEntryList() {
+        List<HeadsUpEntry> entryList = new ArrayList<>(mHeadsUpEntryMap.values());
+        entryList.addAll(mAvalancheController.getWaitingEntryList());
+        return entryList;
     }
 
     /**
@@ -304,7 +308,8 @@ public abstract class BaseHeadsUpManager implements HeadsUpManager {
      */
     @Override
     public boolean hasNotifications() {
-        return !mHeadsUpEntryMap.isEmpty();
+        return !mHeadsUpEntryMap.isEmpty()
+                || !mAvalancheController.getWaitingEntryList().isEmpty();
     }
 
     /**
@@ -507,8 +512,10 @@ public abstract class BaseHeadsUpManager implements HeadsUpManager {
 
     @Nullable
     protected HeadsUpEntry getHeadsUpEntry(@NonNull String key) {
-        // TODO(b/315362456) See if callers need to check AvalancheController
-        return mHeadsUpEntryMap.get(key);
+        if (mHeadsUpEntryMap.containsKey(key)) {
+            return mHeadsUpEntryMap.get(key);
+        }
+        return mAvalancheController.getWaitingEntry(key);
     }
 
     /**
@@ -688,8 +695,7 @@ public abstract class BaseHeadsUpManager implements HeadsUpManager {
      */
     @Override
     public boolean isSticky(String key) {
-        // TODO(b/315362456) See if callers need to check AvalancheController
-        HeadsUpEntry headsUpEntry = mHeadsUpEntryMap.get(key);
+        HeadsUpEntry headsUpEntry = getHeadsUpEntry(key);
         if (headsUpEntry != null) {
             return headsUpEntry.isSticky();
         }