Merge "Fix an ANR caused by the dream overlay status bar." into tm-dev
diff --git a/packages/SystemUI/src/com/android/systemui/dreams/DreamOverlayNotificationCountProvider.java b/packages/SystemUI/src/com/android/systemui/dreams/DreamOverlayNotificationCountProvider.java
new file mode 100644
index 0000000..aaa34ed
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/dreams/DreamOverlayNotificationCountProvider.java
@@ -0,0 +1,116 @@
+/*
+ * Copyright (C) 2022 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.dreams;
+
+import android.annotation.NonNull;
+import android.service.notification.NotificationListenerService;
+import android.service.notification.StatusBarNotification;
+
+import com.android.systemui.dagger.SysUISingleton;
+import com.android.systemui.statusbar.NotificationListener;
+import com.android.systemui.statusbar.NotificationListener.NotificationHandler;
+import com.android.systemui.statusbar.policy.CallbackController;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Set;
+
+import javax.inject.Inject;
+
+/***
+ * {@link DreamOverlayNotificationCountProvider} provides the current notification count to
+ * registered callbacks.
+ */
+@SysUISingleton
+public class DreamOverlayNotificationCountProvider
+        implements CallbackController<DreamOverlayNotificationCountProvider.Callback> {
+    private final Set<String> mNotificationKeys = new HashSet<>();
+    private final List<Callback> mCallbacks = new ArrayList<>();
+
+    private final NotificationHandler mNotificationHandler = new NotificationHandler() {
+        @Override
+        public void onNotificationPosted(
+                StatusBarNotification sbn, NotificationListenerService.RankingMap rankingMap) {
+            mNotificationKeys.add(sbn.getKey());
+            reportNotificationCountChanged();
+        }
+
+        @Override
+        public void onNotificationRemoved(
+                StatusBarNotification sbn, NotificationListenerService.RankingMap rankingMap) {
+            mNotificationKeys.remove(sbn.getKey());
+            reportNotificationCountChanged();
+        }
+
+        @Override
+        public void onNotificationRemoved(
+                StatusBarNotification sbn,
+                NotificationListenerService.RankingMap rankingMap,
+                int reason) {
+            mNotificationKeys.remove(sbn.getKey());
+            reportNotificationCountChanged();
+        }
+
+        @Override
+        public void onNotificationRankingUpdate(NotificationListenerService.RankingMap rankingMap) {
+        }
+
+        @Override
+        public void onNotificationsInitialized() {
+        }
+    };
+
+    @Inject
+    public DreamOverlayNotificationCountProvider(
+            NotificationListener notificationListener) {
+        notificationListener.addNotificationHandler(mNotificationHandler);
+        Arrays.stream(notificationListener.getActiveNotifications())
+                .forEach(sbn -> mNotificationKeys.add(sbn.getKey()));
+    }
+
+    @Override
+    public void addCallback(@NonNull Callback callback) {
+        if (!mCallbacks.contains(callback)) {
+            mCallbacks.add(callback);
+            callback.onNotificationCountChanged(mNotificationKeys.size());
+        }
+    }
+
+    @Override
+    public void removeCallback(@NonNull Callback callback) {
+        mCallbacks.remove(callback);
+    }
+
+    private void reportNotificationCountChanged() {
+        final int notificationCount = mNotificationKeys.size();
+        mCallbacks.forEach(callback -> callback.onNotificationCountChanged(notificationCount));
+    }
+
+    /**
+     * A callback to be registered with {@link DreamOverlayNotificationCountProvider} to receive
+     * changes to the current notification count.
+     */
+    public interface Callback {
+        /**
+         * Called when the notification count has changed.
+         * @param count The current notification count.
+         */
+        void onNotificationCountChanged(int count);
+    }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/dreams/DreamOverlayStatusBarViewController.java b/packages/SystemUI/src/com/android/systemui/dreams/DreamOverlayStatusBarViewController.java
index 761f28c..d4909c78 100644
--- a/packages/SystemUI/src/com/android/systemui/dreams/DreamOverlayStatusBarViewController.java
+++ b/packages/SystemUI/src/com/android/systemui/dreams/DreamOverlayStatusBarViewController.java
@@ -27,16 +27,12 @@
 import android.net.NetworkRequest;
 import android.os.UserHandle;
 import android.provider.Settings;
-import android.service.notification.NotificationListenerService.RankingMap;
-import android.service.notification.StatusBarNotification;
 import android.text.format.DateFormat;
 import android.util.PluralsMessageFormatter;
 
 import com.android.systemui.R;
 import com.android.systemui.dagger.qualifiers.Main;
 import com.android.systemui.dreams.dagger.DreamOverlayComponent;
-import com.android.systemui.statusbar.NotificationListener;
-import com.android.systemui.statusbar.NotificationListener.NotificationHandler;
 import com.android.systemui.statusbar.policy.IndividualSensorPrivacyController;
 import com.android.systemui.statusbar.policy.NextAlarmController;
 import com.android.systemui.statusbar.policy.ZenModeController;
@@ -62,7 +58,7 @@
     private final Resources mResources;
     private final DateFormatUtil mDateFormatUtil;
     private final IndividualSensorPrivacyController mSensorPrivacyController;
-    private final NotificationListener mNotificationListener;
+    private final DreamOverlayNotificationCountProvider mDreamOverlayNotificationCountProvider;
     private final ZenModeController mZenModeController;
     private final Executor mMainExecutor;
 
@@ -96,35 +92,6 @@
     private final NextAlarmController.NextAlarmChangeCallback mNextAlarmCallback =
             nextAlarm -> updateAlarmStatusIcon();
 
-    private final NotificationHandler mNotificationHandler = new NotificationHandler() {
-        @Override
-        public void onNotificationPosted(StatusBarNotification sbn, RankingMap rankingMap) {
-            updateNotificationsStatusIcon();
-        }
-
-        @Override
-        public void onNotificationRemoved(StatusBarNotification sbn, RankingMap rankingMap) {
-            updateNotificationsStatusIcon();
-        }
-
-        @Override
-        public void onNotificationRemoved(
-                StatusBarNotification sbn,
-                RankingMap rankingMap,
-                int reason) {
-            updateNotificationsStatusIcon();
-        }
-
-        @Override
-        public void onNotificationRankingUpdate(RankingMap rankingMap) {
-        }
-
-        @Override
-        public void onNotificationsInitialized() {
-            updateNotificationsStatusIcon();
-        }
-    };
-
     private final ZenModeController.Callback mZenModeCallback = new ZenModeController.Callback() {
         @Override
         public void onZenChanged(int zen) {
@@ -132,6 +99,14 @@
         }
     };
 
+    private final DreamOverlayNotificationCountProvider.Callback mNotificationCountCallback =
+            notificationCount -> showIcon(
+                    DreamOverlayStatusBarView.STATUS_ICON_NOTIFICATIONS,
+                    notificationCount > 0,
+                    notificationCount > 0
+                            ? buildNotificationsContentDescription(notificationCount)
+                            : null);
+
     @Inject
     public DreamOverlayStatusBarViewController(
             DreamOverlayStatusBarView view,
@@ -143,7 +118,7 @@
             NextAlarmController nextAlarmController,
             DateFormatUtil dateFormatUtil,
             IndividualSensorPrivacyController sensorPrivacyController,
-            NotificationListener notificationListener,
+            DreamOverlayNotificationCountProvider dreamOverlayNotificationCountProvider,
             ZenModeController zenModeController) {
         super(view);
         mResources = resources;
@@ -154,20 +129,14 @@
         mNextAlarmController = nextAlarmController;
         mDateFormatUtil = dateFormatUtil;
         mSensorPrivacyController = sensorPrivacyController;
-        mNotificationListener = notificationListener;
+        mDreamOverlayNotificationCountProvider = dreamOverlayNotificationCountProvider;
         mZenModeController = zenModeController;
-
-        // Handlers can be added to NotificationListener, but apparently they can't be removed. So
-        // add the handler here in the constructor rather than in onViewAttached to avoid confusion.
-        mNotificationListener.addNotificationHandler(mNotificationHandler);
     }
 
     @Override
     protected void onViewAttached() {
         mIsAttached = true;
 
-        updateNotificationsStatusIcon();
-
         mConnectivityManager.registerNetworkCallback(mNetworkRequest, mNetworkCallback);
         updateWifiUnavailableStatusIcon();
 
@@ -180,6 +149,7 @@
         mZenModeController.addCallback(mZenModeCallback);
         updatePriorityModeStatusIcon();
 
+        mDreamOverlayNotificationCountProvider.addCallback(mNotificationCountCallback);
         mTouchInsetSession.addViewToTracking(mView);
     }
 
@@ -189,6 +159,7 @@
         mSensorPrivacyController.removeCallback(mSensorCallback);
         mNextAlarmController.removeCallback(mNextAlarmCallback);
         mConnectivityManager.unregisterNetworkCallback(mNetworkCallback);
+        mDreamOverlayNotificationCountProvider.removeCallback(mNotificationCountCallback);
         mTouchInsetSession.clear();
 
         mIsAttached = false;
@@ -231,24 +202,6 @@
                 micBlocked && cameraBlocked);
     }
 
-    private void updateNotificationsStatusIcon() {
-        if (mView == null) {
-            // It is possible for this method to be called before the view is attached, which makes
-            // null-checking necessary.
-            return;
-        }
-
-        final StatusBarNotification[] notifications =
-                mNotificationListener.getActiveNotifications();
-        final int notificationCount = notifications != null ? notifications.length : 0;
-        showIcon(
-                DreamOverlayStatusBarView.STATUS_ICON_NOTIFICATIONS,
-                notificationCount > 0,
-                notificationCount > 0
-                        ? buildNotificationsContentDescription(notificationCount)
-                        : null);
-    }
-
     private String buildNotificationsContentDescription(int notificationCount) {
         return PluralsMessageFormatter.format(
                 mResources,
diff --git a/packages/SystemUI/tests/src/com/android/systemui/dreams/DreamOverlayNotificationCountProviderTest.java b/packages/SystemUI/tests/src/com/android/systemui/dreams/DreamOverlayNotificationCountProviderTest.java
new file mode 100644
index 0000000..c861221
--- /dev/null
+++ b/packages/SystemUI/tests/src/com/android/systemui/dreams/DreamOverlayNotificationCountProviderTest.java
@@ -0,0 +1,85 @@
+/*
+ * Copyright (C) 2022 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.dreams;
+
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+import android.service.notification.NotificationListenerService;
+import android.service.notification.StatusBarNotification;
+import android.testing.AndroidTestingRunner;
+
+import androidx.test.filters.SmallTest;
+
+import com.android.systemui.SysuiTestCase;
+import com.android.systemui.statusbar.NotificationListener;
+import com.android.systemui.statusbar.NotificationListener.NotificationHandler;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.ArgumentCaptor;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+
+@SmallTest
+@RunWith(AndroidTestingRunner.class)
+public class DreamOverlayNotificationCountProviderTest extends SysuiTestCase {
+    @Mock
+    NotificationListener mNotificationListener;
+    @Mock
+    DreamOverlayNotificationCountProvider.Callback mCallback;
+    @Mock
+    StatusBarNotification mNotification1;
+    @Mock
+    StatusBarNotification mNotification2;
+    @Mock
+    NotificationListenerService.RankingMap mRankingMap;
+
+    private DreamOverlayNotificationCountProvider mProvider;
+
+    @Before
+    public void setup() {
+        MockitoAnnotations.initMocks(this);
+
+        when(mNotification1.getKey()).thenReturn("key1");
+        when(mNotification2.getKey()).thenReturn("key2");
+
+        final StatusBarNotification[] notifications = {mNotification1};
+        when(mNotificationListener.getActiveNotifications()).thenReturn(notifications);
+        mProvider = new DreamOverlayNotificationCountProvider(mNotificationListener);
+        mProvider.addCallback(mCallback);
+    }
+
+    @Test
+    public void testPostingNotificationCallsCallbackWithNotificationCount() {
+        final ArgumentCaptor<NotificationHandler> handlerArgumentCaptor =
+                ArgumentCaptor.forClass(NotificationHandler.class);
+        verify(mNotificationListener).addNotificationHandler(handlerArgumentCaptor.capture());
+        handlerArgumentCaptor.getValue().onNotificationPosted(mNotification2, mRankingMap);
+        verify(mCallback).onNotificationCountChanged(2);
+    }
+
+    @Test
+    public void testRemovingNotificationCallsCallbackWithZeroNotificationCount() {
+        final ArgumentCaptor<NotificationHandler> handlerArgumentCaptor =
+                ArgumentCaptor.forClass(NotificationHandler.class);
+        verify(mNotificationListener).addNotificationHandler(handlerArgumentCaptor.capture());
+        handlerArgumentCaptor.getValue().onNotificationRemoved(mNotification1, mRankingMap);
+        verify(mCallback).onNotificationCountChanged(0);
+    }
+}
diff --git a/packages/SystemUI/tests/src/com/android/systemui/dreams/DreamOverlayStatusBarViewControllerTest.java b/packages/SystemUI/tests/src/com/android/systemui/dreams/DreamOverlayStatusBarViewControllerTest.java
index a6921b4..4915ded 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/dreams/DreamOverlayStatusBarViewControllerTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/dreams/DreamOverlayStatusBarViewControllerTest.java
@@ -31,15 +31,12 @@
 import android.net.NetworkCapabilities;
 import android.net.NetworkRequest;
 import android.provider.Settings;
-import android.service.notification.NotificationListenerService;
-import android.service.notification.StatusBarNotification;
 import android.testing.AndroidTestingRunner;
 
 import androidx.test.filters.SmallTest;
 
 import com.android.systemui.R;
 import com.android.systemui.SysuiTestCase;
-import com.android.systemui.statusbar.NotificationListener;
 import com.android.systemui.statusbar.policy.IndividualSensorPrivacyController;
 import com.android.systemui.statusbar.policy.NextAlarmController;
 import com.android.systemui.statusbar.policy.ZenModeController;
@@ -84,13 +81,9 @@
     @Mock
     IndividualSensorPrivacyController mSensorPrivacyController;
     @Mock
-    StatusBarNotification mStatusBarNotification;
-    @Mock
-    NotificationListenerService.RankingMap mRankingMap;
-    @Mock
-    NotificationListener mNotificationListener;
-    @Mock
     ZenModeController mZenModeController;
+    @Mock
+    DreamOverlayNotificationCountProvider mDreamOverlayNotificationCountProvider;
 
     private final Executor mMainExecutor = Runnable::run;
 
@@ -113,7 +106,7 @@
                 mNextAlarmController,
                 mDateFormatUtil,
                 mSensorPrivacyController,
-                mNotificationListener,
+                mDreamOverlayNotificationCountProvider,
                 mZenModeController);
     }
 
@@ -123,6 +116,7 @@
         verify(mNextAlarmController).addCallback(any());
         verify(mSensorPrivacyController).addCallback(any());
         verify(mZenModeController).addCallback(any());
+        verify(mDreamOverlayNotificationCountProvider).addCallback(any());
     }
 
     @Test
@@ -202,17 +196,26 @@
 
     @Test
     public void testOnViewAttachedShowsNotificationsIconWhenNotificationsExist() {
-        StatusBarNotification[] notifications = { mStatusBarNotification };
-        when(mNotificationListener.getActiveNotifications()).thenReturn(notifications);
         mController.onViewAttached();
+
+        final ArgumentCaptor<DreamOverlayNotificationCountProvider.Callback> callbackCapture =
+                ArgumentCaptor.forClass(DreamOverlayNotificationCountProvider.Callback.class);
+        verify(mDreamOverlayNotificationCountProvider).addCallback(callbackCapture.capture());
+        callbackCapture.getValue().onNotificationCountChanged(1);
+
         verify(mView).showIcon(
                 eq(DreamOverlayStatusBarView.STATUS_ICON_NOTIFICATIONS), eq(true), any());
     }
 
     @Test
     public void testOnViewAttachedHidesNotificationsIconWhenNoNotificationsExist() {
-        when(mNotificationListener.getActiveNotifications()).thenReturn(null);
         mController.onViewAttached();
+
+        final ArgumentCaptor<DreamOverlayNotificationCountProvider.Callback> callbackCapture =
+                ArgumentCaptor.forClass(DreamOverlayNotificationCountProvider.Callback.class);
+        verify(mDreamOverlayNotificationCountProvider).addCallback(callbackCapture.capture());
+        callbackCapture.getValue().onNotificationCountChanged(0);
+
         verify(mView).showIcon(
                 eq(DreamOverlayStatusBarView.STATUS_ICON_NOTIFICATIONS), eq(false), isNull());
     }
@@ -248,6 +251,7 @@
         verify(mNextAlarmController).removeCallback(any());
         verify(mSensorPrivacyController).removeCallback(any());
         verify(mZenModeController).removeCallback(any());
+        verify(mDreamOverlayNotificationCountProvider).removeCallback(any());
     }
 
     @Test
@@ -309,13 +313,10 @@
     public void testNotificationsIconShownWhenNotificationAdded() {
         mController.onViewAttached();
 
-        StatusBarNotification[] notifications = { mStatusBarNotification };
-        when(mNotificationListener.getActiveNotifications()).thenReturn(notifications);
-
-        final ArgumentCaptor<NotificationListener.NotificationHandler> callbackCapture =
-                ArgumentCaptor.forClass(NotificationListener.NotificationHandler.class);
-        verify(mNotificationListener).addNotificationHandler(callbackCapture.capture());
-        callbackCapture.getValue().onNotificationPosted(mStatusBarNotification, mRankingMap);
+        final ArgumentCaptor<DreamOverlayNotificationCountProvider.Callback> callbackCapture =
+                ArgumentCaptor.forClass(DreamOverlayNotificationCountProvider.Callback.class);
+        verify(mDreamOverlayNotificationCountProvider).addCallback(callbackCapture.capture());
+        callbackCapture.getValue().onNotificationCountChanged(1);
 
         verify(mView).showIcon(
                 eq(DreamOverlayStatusBarView.STATUS_ICON_NOTIFICATIONS), eq(true), any());
@@ -323,15 +324,12 @@
 
     @Test
     public void testNotificationsIconHiddenWhenLastNotificationRemoved() {
-        StatusBarNotification[] notifications = { mStatusBarNotification };
-        when(mNotificationListener.getActiveNotifications()).thenReturn(notifications)
-                .thenReturn(null);
         mController.onViewAttached();
 
-        final ArgumentCaptor<NotificationListener.NotificationHandler> callbackCapture =
-                ArgumentCaptor.forClass(NotificationListener.NotificationHandler.class);
-        verify(mNotificationListener).addNotificationHandler(callbackCapture.capture());
-        callbackCapture.getValue().onNotificationPosted(mStatusBarNotification, mRankingMap);
+        final ArgumentCaptor<DreamOverlayNotificationCountProvider.Callback> callbackCapture =
+                ArgumentCaptor.forClass(DreamOverlayNotificationCountProvider.Callback.class);
+        verify(mDreamOverlayNotificationCountProvider).addCallback(callbackCapture.capture());
+        callbackCapture.getValue().onNotificationCountChanged(0);
 
         verify(mView).showIcon(
                 eq(DreamOverlayStatusBarView.STATUS_ICON_NOTIFICATIONS), eq(false), any());