Merge "RemoteInputCoordinator to update notif on reply" into main
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/coordinator/RemoteInputCoordinator.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/coordinator/RemoteInputCoordinator.kt
index 5ff5d2d..1fe32c9 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/coordinator/RemoteInputCoordinator.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/coordinator/RemoteInputCoordinator.kt
@@ -117,18 +117,12 @@
                             (entry.getSbn().getNotification().flags and
                                 FLAG_LIFETIME_EXTENDED_BY_DIRECT_REPLY) > 0
                         ) {
+                            // If we've received an update from the system and the entry is marked
+                            // as lifetime extended, that means system server has received a
+                            // cancelation in response to a direct reply, and sent an update to
+                            // let system ui know that it should rebuild the notification with
+                            // that direct reply.
                             if (
-                                mNotificationRemoteInputManager.shouldKeepForRemoteInputHistory(
-                                    entry
-                                )
-                            ) {
-                                val newSbn = mRebuilder.rebuildForRemoteInputReply(entry)
-                                entry.onRemoteInputInserted()
-                                mNotifUpdater.onInternalNotificationUpdate(
-                                    newSbn,
-                                    "Extending lifetime of notification with remote input",
-                                )
-                            } else if (
                                 mNotificationRemoteInputManager.shouldKeepForSmartReplyHistory(
                                     entry
                                 )
@@ -140,16 +134,11 @@
                                     "Extending lifetime of notification with smart reply",
                                 )
                             } else {
-                                // The app may have re-cancelled a notification after it had already
-                                // been lifetime extended.
-                                // Rebuild the notification with the replies it already had to
-                                // ensure
-                                // those replies continue to be displayed.
-                                val newSbn = mRebuilder.rebuildWithExistingReplies(entry)
+                                val newSbn = mRebuilder.rebuildForRemoteInputReply(entry)
+                                entry.onRemoteInputInserted()
                                 mNotifUpdater.onInternalNotificationUpdate(
                                     newSbn,
-                                    "Extending lifetime of notification that has already been " +
-                                        "lifetime extended.",
+                                    "Extending lifetime of notification with remote input",
                                 )
                             }
                         } else {
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/coordinator/RemoteInputCoordinatorTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/coordinator/RemoteInputCoordinatorTest.kt
index deb3fc1..a3f8452 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/coordinator/RemoteInputCoordinatorTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/coordinator/RemoteInputCoordinatorTest.kt
@@ -15,8 +15,8 @@
  */
 package com.android.systemui.statusbar.notification.collection.coordinator
 
-import android.app.Flags.lifetimeExtensionRefactor
 import android.app.Flags.FLAG_LIFETIME_EXTENSION_REFACTOR
+import android.app.Flags.lifetimeExtensionRefactor
 import android.app.Notification
 import android.app.RemoteInputHistoryItem
 import android.os.Handler
@@ -47,10 +47,10 @@
 import org.junit.Test
 import org.junit.runner.RunWith
 import org.mockito.Mock
-import org.mockito.Mockito.`when`
 import org.mockito.Mockito.never
 import org.mockito.Mockito.times
 import org.mockito.Mockito.verify
+import org.mockito.Mockito.`when`
 import org.mockito.MockitoAnnotations.initMocks
 
 @SmallTest
@@ -78,21 +78,20 @@
     @Before
     fun setUp() {
         initMocks(this)
-        coordinator = RemoteInputCoordinator(
+        coordinator =
+            RemoteInputCoordinator(
                 dumpManager,
                 rebuilder,
                 remoteInputManager,
                 mainHandler,
-                smartReplyController
-        )
+                smartReplyController,
+            )
         `when`(pipeline.addNotificationLifetimeExtender(any())).thenAnswer {
             (it.arguments[0] as NotifLifetimeExtender).setCallback(lifetimeExtensionCallback)
         }
         `when`(pipeline.getInternalNotifUpdater(any())).thenReturn(notifUpdater)
         coordinator.attach(pipeline)
-        listener = withArgCaptor {
-            verify(remoteInputManager).setRemoteInputListener(capture())
-        }
+        listener = withArgCaptor { verify(remoteInputManager).setRemoteInputListener(capture()) }
         entry1 = NotificationEntryBuilder().setId(1).build()
         entry2 = NotificationEntryBuilder().setId(2).build()
         `when`(rebuilder.rebuildForCanceledSmartReplies(any())).thenReturn(sbn)
@@ -101,13 +100,17 @@
         `when`(rebuilder.rebuildWithExistingReplies(any())).thenReturn(sbn)
     }
 
-    val remoteInputActiveExtender get() = coordinator.mRemoteInputActiveExtender
-    val remoteInputHistoryExtender get() = coordinator.mRemoteInputHistoryExtender
-    val smartReplyHistoryExtender get() = coordinator.mSmartReplyHistoryExtender
+    val remoteInputActiveExtender
+        get() = coordinator.mRemoteInputActiveExtender
 
-    val collectionListeners get() = captureMany {
-        verify(pipeline, times(1)).addCollectionListener(capture())
-    }
+    val remoteInputHistoryExtender
+        get() = coordinator.mRemoteInputHistoryExtender
+
+    val smartReplyHistoryExtender
+        get() = coordinator.mSmartReplyHistoryExtender
+
+    val collectionListeners
+        get() = captureMany { verify(pipeline, times(1)).addCollectionListener(capture()) }
 
     @Test
     fun testRemoteInputActive() {
@@ -179,7 +182,8 @@
     @EnableFlags(FLAG_LIFETIME_EXTENSION_REFACTOR)
     fun testRemoteInputLifetimeExtensionListenerTrigger() {
         // Create notification with LIFETIME_EXTENDED_BY_DIRECT_REPLY flag.
-        val entry = NotificationEntryBuilder()
+        val entry =
+            NotificationEntryBuilder()
                 .setId(3)
                 .setTag("entry")
                 .setFlag(mContext, Notification.FLAG_LIFETIME_EXTENDED_BY_DIRECT_REPLY, true)
@@ -187,9 +191,7 @@
         `when`(remoteInputManager.shouldKeepForRemoteInputHistory(entry)).thenReturn(true)
         `when`(remoteInputManager.shouldKeepForSmartReplyHistory(entry)).thenReturn(false)
 
-        collectionListeners.forEach {
-            it.onEntryUpdated(entry, true)
-        }
+        collectionListeners.forEach { it.onEntryUpdated(entry, true) }
 
         verify(rebuilder, times(1)).rebuildForRemoteInputReply(entry)
     }
@@ -198,16 +200,15 @@
     @EnableFlags(FLAG_LIFETIME_EXTENSION_REFACTOR)
     fun testSmartReplyLifetimeExtensionListenerTrigger() {
         // Create notification with LIFETIME_EXTENDED_BY_DIRECT_REPLY flag.
-        val entry = NotificationEntryBuilder()
+        val entry =
+            NotificationEntryBuilder()
                 .setId(3)
                 .setTag("entry")
                 .setFlag(mContext, Notification.FLAG_LIFETIME_EXTENDED_BY_DIRECT_REPLY, true)
                 .build()
         `when`(remoteInputManager.shouldKeepForRemoteInputHistory(entry)).thenReturn(false)
         `when`(remoteInputManager.shouldKeepForSmartReplyHistory(entry)).thenReturn(true)
-        collectionListeners.forEach {
-            it.onEntryUpdated(entry, true)
-        }
+        collectionListeners.forEach { it.onEntryUpdated(entry, true) }
 
         verify(rebuilder, times(1)).rebuildForCanceledSmartReplies(entry)
         verify(smartReplyController, times(1)).stopSending(entry)
@@ -217,25 +218,25 @@
     @EnableFlags(FLAG_LIFETIME_EXTENSION_REFACTOR)
     fun testRepeatedUpdateTriggersRebuild() {
         // Create notification with LIFETIME_EXTENDED_BY_DIRECT_REPLY flag.
-        val entry = NotificationEntryBuilder()
+        val entry =
+            NotificationEntryBuilder()
                 .setId(3)
                 .setTag("entry")
                 .setFlag(mContext, Notification.FLAG_LIFETIME_EXTENDED_BY_DIRECT_REPLY, true)
                 .build()
         `when`(remoteInputManager.shouldKeepForRemoteInputHistory(entry)).thenReturn(false)
         `when`(remoteInputManager.shouldKeepForSmartReplyHistory(entry)).thenReturn(false)
-        collectionListeners.forEach {
-            it.onEntryUpdated(entry, true)
-        }
+        collectionListeners.forEach { it.onEntryUpdated(entry, true) }
 
-        verify(rebuilder, times(1)).rebuildWithExistingReplies(entry)
+        verify(rebuilder, times(1)).rebuildForRemoteInputReply(entry)
     }
 
     @Test
     @EnableFlags(FLAG_LIFETIME_EXTENSION_REFACTOR)
     fun testLifetimeExtensionListenerClearsRemoteInputs() {
         // Create notification with LIFETIME_EXTENDED_BY_DIRECT_REPLY flag.
-        val entry = NotificationEntryBuilder()
+        val entry =
+            NotificationEntryBuilder()
                 .setId(3)
                 .setTag("entry")
                 .setFlag(mContext, Notification.FLAG_LIFETIME_EXTENDED_BY_DIRECT_REPLY, false)
@@ -245,9 +246,7 @@
         `when`(remoteInputManager.shouldKeepForRemoteInputHistory(entry)).thenReturn(false)
         `when`(remoteInputManager.shouldKeepForSmartReplyHistory(entry)).thenReturn(false)
 
-        collectionListeners.forEach {
-            it.onEntryUpdated(entry, true)
-        }
+        collectionListeners.forEach { it.onEntryUpdated(entry, true) }
 
         assertThat(entry.remoteInputs).isNull()
     }