Merge changes from topic "b204127880_pipeline_backport_3" into sc-v2-dev am: a509832d9c

Original change: https://googleplex-android-review.googlesource.com/c/platform/frameworks/base/+/16167382

Change-Id: I31db4435cca95ceefc959d922ff40841b6cf893b

9 files changed, 766 insertions, 15 deletions

diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/NotificationRemoteInputManager.java b/packages/SystemUI/src/com/android/systemui/statusbar/NotificationRemoteInputManager.java
index 7151cbb3ec23..1ce7f0350019 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/NotificationRemoteInputManager.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/NotificationRemoteInputManager.java
@@ -58,7 +58,6 @@ import com.android.systemui.plugins.statusbar.StatusBarStateController;
 import com.android.systemui.statusbar.dagger.StatusBarDependenciesModule;
 import com.android.systemui.statusbar.notification.NotificationEntryListener;
 import com.android.systemui.statusbar.notification.NotificationEntryManager;
-import com.android.systemui.statusbar.notification.collection.NotifPipeline;
 import com.android.systemui.statusbar.notification.collection.NotificationEntry;
 import com.android.systemui.statusbar.notification.collection.NotificationEntry.EditedSuggestionInfo;
 import com.android.systemui.statusbar.notification.logging.NotificationLogger;
@@ -203,7 +202,7 @@ public class NotificationRemoteInputManager implements Dumpable {
                 ViewGroup actionGroup = (ViewGroup) parent;
                 buttonIndex = actionGroup.indexOfChild(view);
             }
-            // FIXME: get this for the new pipeline!
+            // TODO(b/204183781): get this from the current pipeline
             final int count = mEntryManager.getActiveNotificationsCount();
             final int rank = entry.getRanking().getRank();
 
@@ -265,7 +264,7 @@ public class NotificationRemoteInputManager implements Dumpable {
             NotificationLockscreenUserManager lockscreenUserManager,
             SmartReplyController smartReplyController,
             NotificationEntryManager notificationEntryManager,
-            NotifPipeline notifPipeline,
+            RemoteInputNotificationRebuilder rebuilder,
             Lazy<Optional<StatusBar>> statusBarOptionalLazy,
             StatusBarStateController statusBarStateController,
             @Main Handler mainHandler,
@@ -284,7 +283,7 @@ public class NotificationRemoteInputManager implements Dumpable {
         mBarService = IStatusBarService.Stub.asInterface(
                 ServiceManager.getService(Context.STATUS_BAR_SERVICE));
         mUserManager = (UserManager) mContext.getSystemService(Context.USER_SERVICE);
-        mRebuilder = new RemoteInputNotificationRebuilder(context);  // TODO: inject?
+        mRebuilder = rebuilder;
         if (!featureFlags.isNewNotifPipelineRenderingEnabled()) {
             mRemoteInputListener = createLegacyRemoteInputLifetimeExtender(mainHandler,
                     notificationEntryManager, smartReplyController);
@@ -320,6 +319,19 @@ public class NotificationRemoteInputManager implements Dumpable {
         });
     }
 
+    /** Add a listener for various remote input events.  Works with NEW pipeline only. */
+    public void setRemoteInputListener(@NonNull RemoteInputListener remoteInputListener) {
+        if (mFeatureFlags.isNewNotifPipelineRenderingEnabled()) {
+            if (mRemoteInputListener != null) {
+                throw new IllegalStateException("mRemoteInputListener is already set");
+            }
+            mRemoteInputListener = remoteInputListener;
+            if (mRemoteInputController != null) {
+                mRemoteInputListener.setRemoteInputController(mRemoteInputController);
+            }
+        }
+    }
+
     @NonNull
     @VisibleForTesting
     protected LegacyRemoteInputLifetimeExtender createLegacyRemoteInputLifetimeExtender(
@@ -333,7 +345,9 @@ public class NotificationRemoteInputManager implements Dumpable {
     public void setUpWithCallback(Callback callback, RemoteInputController.Delegate delegate) {
         mCallback = callback;
         mRemoteInputController = new RemoteInputController(delegate, mRemoteInputUriController);
-        mRemoteInputListener.setRemoteInputController(mRemoteInputController);
+        if (mRemoteInputListener != null) {
+            mRemoteInputListener.setRemoteInputController(mRemoteInputController);
+        }
         // Register all stored callbacks from before the Controller was initialized.
         for (RemoteInputController.Callback cb : mControllerCallbacks) {
             mRemoteInputController.addCallback(cb);
@@ -366,7 +380,6 @@ public class NotificationRemoteInputManager implements Dumpable {
             }
         });
         if (!mFeatureFlags.isNewNotifPipelineRenderingEnabled()) {
-            // FIXME: Don't forget to implement this in the coordinator!
             mSmartReplyController.setCallback((entry, reply) -> {
                 StatusBarNotification newSbn = mRebuilder.rebuildForSendingSmartReply(entry, reply);
                 mEntryManager.updateNotification(newSbn, null /* ranking */);
@@ -572,6 +585,14 @@ public class NotificationRemoteInputManager implements Dumpable {
         // OLD pipeline code ONLY; can assume implementation
         ((LegacyRemoteInputLifetimeExtender) mRemoteInputListener)
                 .mKeysKeptForRemoteInputHistory.remove(key);
+        cleanUpRemoteInputForUserRemoval(entry);
+    }
+
+    /**
+     * Disable remote input on the entry and remove the remote input view.
+     * This should be called when a user dismisses a notification that won't be lifetime extended.
+     */
+    public void cleanUpRemoteInputForUserRemoval(NotificationEntry entry) {
         if (isRemoteInputActive(entry)) {
             entry.mRemoteEditImeVisible = false;
             mRemoteInputController.removeRemoteInput(entry, null);
@@ -762,15 +783,21 @@ public class NotificationRemoteInputManager implements Dumpable {
         boolean showBouncerIfNecessary();
     }
 
+    /** An interface for listening to remote input events that relate to notification lifetime */
     public interface RemoteInputListener {
-        void onRemoteInputSent(NotificationEntry entry);
+        /** Called when remote input pending intent has been sent */
+        void onRemoteInputSent(@NonNull NotificationEntry entry);
 
+        /** Called when the notification shade becomes fully closed */
         void onPanelCollapsed();
 
-        boolean isNotificationKeptForRemoteInputHistory(String key);
+        /** @return whether lifetime of a notification is being extended by the listener */
+        boolean isNotificationKeptForRemoteInputHistory(@NonNull String key);
 
+        /** Called on user interaction to end lifetime extension for history */
         void releaseNotificationIfKeptForRemoteInputHistory(@NonNull NotificationEntry entry);
 
+        /** Called when the RemoteInputController is attached to the manager */
         void setRemoteInputController(@NonNull RemoteInputController remoteInputController);
     }
 
@@ -826,7 +853,7 @@ public class NotificationRemoteInputManager implements Dumpable {
         }
 
         @Override
-        public void onRemoteInputSent(NotificationEntry entry) {
+        public void onRemoteInputSent(@NonNull NotificationEntry entry) {
             if (FORCE_REMOTE_INPUT_HISTORY
                     && isNotificationKeptForRemoteInputHistory(entry.getKey())) {
                 mNotificationLifetimeFinishedCallback.onSafeToRemove(entry.getKey());
@@ -858,7 +885,7 @@ public class NotificationRemoteInputManager implements Dumpable {
         }
 
         @Override
-        public boolean isNotificationKeptForRemoteInputHistory(String key) {
+        public boolean isNotificationKeptForRemoteInputHistory(@NonNull String key) {
             return mKeysKeptForRemoteInputHistory.contains(key);
         }
 
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/dagger/StatusBarDependenciesModule.java b/packages/SystemUI/src/com/android/systemui/statusbar/dagger/StatusBarDependenciesModule.java
index e9071f075e5e..bb697c3b0851 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/dagger/StatusBarDependenciesModule.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/dagger/StatusBarDependenciesModule.java
@@ -43,6 +43,7 @@ import com.android.systemui.statusbar.NotificationMediaManager;
 import com.android.systemui.statusbar.NotificationRemoteInputManager;
 import com.android.systemui.statusbar.NotificationShadeWindowController;
 import com.android.systemui.statusbar.NotificationViewHierarchyManager;
+import com.android.systemui.statusbar.RemoteInputNotificationRebuilder;
 import com.android.systemui.statusbar.SmartReplyController;
 import com.android.systemui.statusbar.StatusBarStateControllerImpl;
 import com.android.systemui.statusbar.SysuiStatusBarStateController;
@@ -100,7 +101,7 @@ public interface StatusBarDependenciesModule {
             NotificationLockscreenUserManager lockscreenUserManager,
             SmartReplyController smartReplyController,
             NotificationEntryManager notificationEntryManager,
-            NotifPipeline notifPipeline,
+            RemoteInputNotificationRebuilder rebuilder,
             Lazy<Optional<StatusBar>> statusBarOptionalLazy,
             StatusBarStateController statusBarStateController,
             Handler mainHandler,
@@ -114,7 +115,7 @@ public interface StatusBarDependenciesModule {
                 lockscreenUserManager,
                 smartReplyController,
                 notificationEntryManager,
-                notifPipeline,
+                rebuilder,
                 statusBarOptionalLazy,
                 statusBarStateController,
                 mainHandler,
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/coordinator/NotifCoordinators.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/coordinator/NotifCoordinators.kt
index 66290bb3aba6..39b1ec4ff80e 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/coordinator/NotifCoordinators.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/coordinator/NotifCoordinators.kt
@@ -48,6 +48,7 @@ class NotifCoordinatorsImpl @Inject constructor(
     conversationCoordinator: ConversationCoordinator,
     preparationCoordinator: PreparationCoordinator,
     mediaCoordinator: MediaCoordinator,
+    remoteInputCoordinator: RemoteInputCoordinator,
     shadeEventCoordinator: ShadeEventCoordinator,
     smartspaceDedupingCoordinator: SmartspaceDedupingCoordinator,
     viewConfigCoordinator: ViewConfigCoordinator,
@@ -72,6 +73,7 @@ class NotifCoordinatorsImpl @Inject constructor(
         mCoordinators.add(bubbleCoordinator)
         mCoordinators.add(conversationCoordinator)
         mCoordinators.add(mediaCoordinator)
+        mCoordinators.add(remoteInputCoordinator)
         mCoordinators.add(shadeEventCoordinator)
         mCoordinators.add(viewConfigCoordinator)
         mCoordinators.add(visualStabilityCoordinator)
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
new file mode 100644
index 000000000000..3397815f008f
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/coordinator/RemoteInputCoordinator.kt
@@ -0,0 +1,225 @@
+/*
+ * Copyright (C) 2021 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.statusbar.notification.collection.coordinator
+
+import android.os.Handler
+import android.service.notification.NotificationListenerService.REASON_CANCEL
+import android.service.notification.NotificationListenerService.REASON_CLICK
+import android.util.Log
+import androidx.annotation.VisibleForTesting
+import com.android.systemui.Dumpable
+import com.android.systemui.dagger.SysUISingleton
+import com.android.systemui.dagger.qualifiers.Main
+import com.android.systemui.dump.DumpManager
+import com.android.systemui.statusbar.NotificationRemoteInputManager
+import com.android.systemui.statusbar.NotificationRemoteInputManager.RemoteInputListener
+import com.android.systemui.statusbar.RemoteInputController
+import com.android.systemui.statusbar.RemoteInputNotificationRebuilder
+import com.android.systemui.statusbar.SmartReplyController
+import com.android.systemui.statusbar.notification.collection.NotifPipeline
+import com.android.systemui.statusbar.notification.collection.NotificationEntry
+import com.android.systemui.statusbar.notification.collection.notifcollection.SelfTrackingLifetimeExtender
+import com.android.systemui.statusbar.notification.collection.notifcollection.InternalNotifUpdater
+import com.android.systemui.statusbar.notification.collection.notifcollection.NotifCollectionListener
+import com.android.systemui.statusbar.notification.collection.notifcollection.NotifLifetimeExtender
+import java.io.FileDescriptor
+import java.io.PrintWriter
+import javax.inject.Inject
+
+private const val TAG = "RemoteInputCoordinator"
+
+/**
+ * How long to wait before auto-dismissing a notification that was kept for active remote input, and
+ * has now sent a remote input. We auto-dismiss, because the app may not cannot cancel
+ * these given that they technically don't exist anymore. We wait a bit in case the app issues
+ * an update, and to also give the other lifetime extenders a beat to decide they want it.
+ */
+private const val REMOTE_INPUT_ACTIVE_EXTENDER_AUTO_CANCEL_DELAY: Long = 500
+
+/**
+ * How long to wait before releasing a lifetime extension when requested to do so due to a user
+ * interaction (such as tapping another action).
+ * We wait a bit in case the app issues an update in response to the action, but not too long or we
+ * risk appearing unresponsive to the user.
+ */
+private const val REMOTE_INPUT_EXTENDER_RELEASE_DELAY: Long = 200
+
+/** Whether this class should print spammy debug logs */
+private val DEBUG: Boolean by lazy { Log.isLoggable(TAG, Log.DEBUG) }
+
+@SysUISingleton
+class RemoteInputCoordinator @Inject constructor(
+    dumpManager: DumpManager,
+    private val mRebuilder: RemoteInputNotificationRebuilder,
+    private val mNotificationRemoteInputManager: NotificationRemoteInputManager,
+    @Main private val mMainHandler: Handler,
+    private val mSmartReplyController: SmartReplyController
+) : Coordinator, RemoteInputListener, Dumpable {
+
+    @VisibleForTesting val mRemoteInputHistoryExtender = RemoteInputHistoryExtender()
+    @VisibleForTesting val mSmartReplyHistoryExtender = SmartReplyHistoryExtender()
+    @VisibleForTesting val mRemoteInputActiveExtender = RemoteInputActiveExtender()
+    private val mRemoteInputLifetimeExtenders = listOf(
+            mRemoteInputHistoryExtender,
+            mSmartReplyHistoryExtender,
+            mRemoteInputActiveExtender
+    )
+
+    private lateinit var mNotifUpdater: InternalNotifUpdater
+
+    init {
+        dumpManager.registerDumpable(this)
+    }
+
+    fun getLifetimeExtenders(): List<NotifLifetimeExtender> = mRemoteInputLifetimeExtenders
+
+    override fun attach(pipeline: NotifPipeline) {
+        mNotificationRemoteInputManager.setRemoteInputListener(this)
+        mRemoteInputLifetimeExtenders.forEach { pipeline.addNotificationLifetimeExtender(it) }
+        mNotifUpdater = pipeline.getInternalNotifUpdater(TAG)
+        pipeline.addCollectionListener(mCollectionListener)
+    }
+
+    val mCollectionListener = object : NotifCollectionListener {
+        override fun onEntryUpdated(entry: NotificationEntry, fromSystem: Boolean) {
+            if (DEBUG) {
+                Log.d(TAG, "mCollectionListener.onEntryUpdated(entry=${entry.key}," +
+                        " fromSystem=$fromSystem)")
+            }
+            if (fromSystem) {
+                // Mark smart replies as sent whenever a notification is updated by the app,
+                // otherwise the smart replies are never marked as sent.
+                mSmartReplyController.stopSending(entry)
+            }
+        }
+
+        override fun onEntryRemoved(entry: NotificationEntry, reason: Int) {
+            if (DEBUG) Log.d(TAG, "mCollectionListener.onEntryRemoved(entry=${entry.key})")
+            // We're removing the notification, the smart reply controller can forget about it.
+            // TODO(b/145659174): track 'sending' state on the entry to avoid having to clear it.
+            mSmartReplyController.stopSending(entry)
+
+            // When we know the entry will not be lifetime extended, clean up the remote input view
+            // TODO: Share code with NotifCollection.cannotBeLifetimeExtended
+            if (reason == REASON_CANCEL || reason == REASON_CLICK) {
+                mNotificationRemoteInputManager.cleanUpRemoteInputForUserRemoval(entry)
+            }
+        }
+    }
+
+    override fun dump(fd: FileDescriptor, pw: PrintWriter, args: Array<out String>) {
+        mRemoteInputLifetimeExtenders.forEach { it.dump(fd, pw, args) }
+    }
+
+    override fun onRemoteInputSent(entry: NotificationEntry) {
+        if (DEBUG) Log.d(TAG, "onRemoteInputSent(entry=${entry.key})")
+        // These calls effectively ensure the freshness of the lifetime extensions.
+        // NOTE: This is some trickery! By removing the lifetime extensions when we know they should
+        // be immediately re-upped, we ensure that the side-effects of the lifetime extenders get to
+        // fire again, thus ensuring that we add subsequent replies to the notification.
+        mRemoteInputHistoryExtender.endLifetimeExtension(entry.key)
+        mSmartReplyHistoryExtender.endLifetimeExtension(entry.key)
+
+        // If we're extending for remote input being active, then from the apps point of
+        // view it is already canceled, so we'll need to cancel it on the apps behalf
+        // now that a reply has been sent. However, delay so that the app has time to posts an
+        // update in the mean time, and to give another lifetime extender time to pick it up.
+        mRemoteInputActiveExtender.endLifetimeExtensionAfterDelay(entry.key,
+                REMOTE_INPUT_ACTIVE_EXTENDER_AUTO_CANCEL_DELAY)
+    }
+
+    private fun onSmartReplySent(entry: NotificationEntry, reply: CharSequence) {
+        if (DEBUG) Log.d(TAG, "onSmartReplySent(entry=${entry.key})")
+        val newSbn = mRebuilder.rebuildForSendingSmartReply(entry, reply)
+        mNotifUpdater.onInternalNotificationUpdate(newSbn,
+                "Adding smart reply spinner for sent")
+
+        // If we're extending for remote input being active, then from the apps point of
+        // view it is already canceled, so we'll need to cancel it on the apps behalf
+        // now that a reply has been sent. However, delay so that the app has time to posts an
+        // update in the mean time, and to give another lifetime extender time to pick it up.
+        mRemoteInputActiveExtender.endLifetimeExtensionAfterDelay(entry.key,
+                REMOTE_INPUT_ACTIVE_EXTENDER_AUTO_CANCEL_DELAY)
+    }
+
+    override fun onPanelCollapsed() {
+        mRemoteInputActiveExtender.endAllLifetimeExtensions()
+    }
+
+    override fun isNotificationKeptForRemoteInputHistory(key: String) =
+            mRemoteInputHistoryExtender.isExtending(key) ||
+                    mSmartReplyHistoryExtender.isExtending(key)
+
+    override fun releaseNotificationIfKeptForRemoteInputHistory(entry: NotificationEntry) {
+        if (DEBUG) Log.d(TAG, "releaseNotificationIfKeptForRemoteInputHistory(entry=${entry.key})")
+        mRemoteInputHistoryExtender.endLifetimeExtensionAfterDelay(entry.key,
+                REMOTE_INPUT_EXTENDER_RELEASE_DELAY)
+        mSmartReplyHistoryExtender.endLifetimeExtensionAfterDelay(entry.key,
+                REMOTE_INPUT_EXTENDER_RELEASE_DELAY)
+        mRemoteInputActiveExtender.endLifetimeExtensionAfterDelay(entry.key,
+                REMOTE_INPUT_EXTENDER_RELEASE_DELAY)
+    }
+
+    override fun setRemoteInputController(remoteInputController: RemoteInputController) {
+        mSmartReplyController.setCallback(this::onSmartReplySent)
+    }
+
+    @VisibleForTesting
+    inner class RemoteInputHistoryExtender :
+            SelfTrackingLifetimeExtender(TAG, "RemoteInputHistory", DEBUG, mMainHandler) {
+
+        override fun queryShouldExtendLifetime(entry: NotificationEntry): Boolean =
+                mNotificationRemoteInputManager.shouldKeepForRemoteInputHistory(entry)
+
+        override fun onStartedLifetimeExtension(entry: NotificationEntry) {
+            val newSbn = mRebuilder.rebuildForRemoteInputReply(entry)
+            entry.onRemoteInputInserted()
+            mNotifUpdater.onInternalNotificationUpdate(newSbn,
+                    "Extending lifetime of notification with remote input")
+            // TODO: Check if the entry was removed due perhaps to an inflation exception?
+        }
+    }
+
+    @VisibleForTesting
+    inner class SmartReplyHistoryExtender :
+            SelfTrackingLifetimeExtender(TAG, "SmartReplyHistory", DEBUG, mMainHandler) {
+
+        override fun queryShouldExtendLifetime(entry: NotificationEntry): Boolean =
+                mNotificationRemoteInputManager.shouldKeepForSmartReplyHistory(entry)
+
+        override fun onStartedLifetimeExtension(entry: NotificationEntry) {
+            val newSbn = mRebuilder.rebuildForCanceledSmartReplies(entry)
+            mSmartReplyController.stopSending(entry)
+            mNotifUpdater.onInternalNotificationUpdate(newSbn,
+                    "Extending lifetime of notification with smart reply")
+            // TODO: Check if the entry was removed due perhaps to an inflation exception?
+        }
+
+        override fun onCanceledLifetimeExtension(entry: NotificationEntry) {
+            // TODO(b/145659174): track 'sending' state on the entry to avoid having to clear it.
+            mSmartReplyController.stopSending(entry)
+        }
+    }
+
+    @VisibleForTesting
+    inner class RemoteInputActiveExtender :
+            SelfTrackingLifetimeExtender(TAG, "RemoteInputActive", DEBUG, mMainHandler) {
+
+        override fun queryShouldExtendLifetime(entry: NotificationEntry): Boolean =
+                mNotificationRemoteInputManager.isRemoteInputActive(entry)
+    }
+}
\ No newline at end of file
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/notifcollection/SelfTrackingLifetimeExtender.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/notifcollection/SelfTrackingLifetimeExtender.kt
new file mode 100644
index 000000000000..145c1e54d732
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/notifcollection/SelfTrackingLifetimeExtender.kt
@@ -0,0 +1,113 @@
+package com.android.systemui.statusbar.notification.collection.notifcollection
+
+import android.os.Handler
+import android.util.ArrayMap
+import android.util.Log
+import com.android.systemui.Dumpable
+import com.android.systemui.statusbar.notification.collection.NotificationEntry
+import java.io.FileDescriptor
+import java.io.PrintWriter
+
+/**
+ * A helpful class that implements the core contract of the lifetime extender internally,
+ * making it easier for coordinators to interact with them
+ */
+abstract class SelfTrackingLifetimeExtender(
+    private val tag: String,
+    private val name: String,
+    private val debug: Boolean,
+    private val mainHandler: Handler
+) : NotifLifetimeExtender, Dumpable {
+    private lateinit var mCallback: NotifLifetimeExtender.OnEndLifetimeExtensionCallback
+    protected val mEntriesExtended = ArrayMap<String, NotificationEntry>()
+    private var mEnding = false
+
+    /**
+     * When debugging, warn if the call is happening during and "end lifetime extension" call.
+     *
+     * Note: this will warn a lot! The pipeline explicitly re-invokes all lifetime extenders
+     * whenever one ends, giving all of them a chance to re-up their lifetime extension.
+     */
+    private fun warnIfEnding() {
+        if (debug && mEnding) Log.w(tag, "reentrant code while ending a lifetime extension")
+    }
+
+    fun endAllLifetimeExtensions() {
+        // clear the map before iterating over a copy of the items, because the pipeline will
+        // always give us another chance to extend the lifetime again, and we don't want
+        // concurrent modification
+        val entries = mEntriesExtended.values.toList()
+        if (debug) Log.d(tag, "$name.endAllLifetimeExtensions() entries=$entries")
+        mEntriesExtended.clear()
+        warnIfEnding()
+        mEnding = true
+        entries.forEach { mCallback.onEndLifetimeExtension(this, it) }
+        mEnding = false
+    }
+
+    fun endLifetimeExtensionAfterDelay(key: String, delayMillis: Long) {
+        if (debug) {
+            Log.d(tag, "$name.endLifetimeExtensionAfterDelay" +
+                    "(key=$key, delayMillis=$delayMillis)" +
+                    " isExtending=${isExtending(key)}")
+        }
+        if (isExtending(key)) {
+            mainHandler.postDelayed({ endLifetimeExtension(key) }, delayMillis)
+        }
+    }
+
+    fun endLifetimeExtension(key: String) {
+        if (debug) {
+            Log.d(tag, "$name.endLifetimeExtension(key=$key)" +
+                    " isExtending=${isExtending(key)}")
+        }
+        warnIfEnding()
+        mEnding = true
+        mEntriesExtended.remove(key)?.let { removedEntry ->
+            mCallback.onEndLifetimeExtension(this, removedEntry)
+        }
+        mEnding = false
+    }
+
+    fun isExtending(key: String) = mEntriesExtended.contains(key)
+
+    final override fun getName(): String = name
+
+    final override fun shouldExtendLifetime(entry: NotificationEntry, reason: Int): Boolean {
+        val shouldExtend = queryShouldExtendLifetime(entry)
+        if (debug) {
+            Log.d(tag, "$name.shouldExtendLifetime(key=${entry.key}, reason=$reason)" +
+                    " isExtending=${isExtending(entry.key)}" +
+                    " shouldExtend=$shouldExtend")
+        }
+        warnIfEnding()
+        if (shouldExtend && mEntriesExtended.put(entry.key, entry) == null) {
+            onStartedLifetimeExtension(entry)
+        }
+        return shouldExtend
+    }
+
+    final override fun cancelLifetimeExtension(entry: NotificationEntry) {
+        if (debug) {
+            Log.d(tag, "$name.cancelLifetimeExtension(key=${entry.key})" +
+                    " isExtending=${isExtending(entry.key)}")
+        }
+        warnIfEnding()
+        mEntriesExtended.remove(entry.key)
+        onCanceledLifetimeExtension(entry)
+    }
+
+    abstract fun queryShouldExtendLifetime(entry: NotificationEntry): Boolean
+    open fun onStartedLifetimeExtension(entry: NotificationEntry) {}
+    open fun onCanceledLifetimeExtension(entry: NotificationEntry) {}
+
+    final override fun setCallback(callback: NotifLifetimeExtender.OnEndLifetimeExtensionCallback) {
+        mCallback = callback
+    }
+
+    final override fun dump(fd: FileDescriptor, pw: PrintWriter, args: Array<out String>) {
+        pw.println("LifetimeExtender: $name:")
+        pw.println("  mEntriesExtended: ${mEntriesExtended.size}")
+        mEntriesExtended.forEach { pw.println("  * ${it.key}") }
+    }
+}
\ No newline at end of file
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/NotificationRemoteInputManagerTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/NotificationRemoteInputManagerTest.java
index add5c6a7c1d6..4ed722470334 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/NotificationRemoteInputManagerTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/NotificationRemoteInputManagerTest.java
@@ -99,7 +99,10 @@ public class NotificationRemoteInputManagerTest extends SysuiTestCase {
 
         mRemoteInputManager = new TestableNotificationRemoteInputManager(mContext,
                 mock(FeatureFlags.class),
-                mLockscreenUserManager, mSmartReplyController, mEntryManager,
+                mLockscreenUserManager,
+                mSmartReplyController,
+                mEntryManager,
+                mock(RemoteInputNotificationRebuilder.class),
                 () -> Optional.of(mock(StatusBar.class)),
                 mStateController,
                 Handler.createAsync(Looper.myLooper()),
@@ -137,6 +140,7 @@ public class NotificationRemoteInputManagerTest extends SysuiTestCase {
     public void testShouldExtendLifetime_remoteInputActive() {
         when(mController.isRemoteInputActive(mEntry)).thenReturn(true);
 
+        assertTrue(mRemoteInputManager.isRemoteInputActive(mEntry));
         assertTrue(mRemoteInputActiveExtender.shouldExtendLifetime(mEntry));
     }
 
@@ -145,6 +149,7 @@ public class NotificationRemoteInputManagerTest extends SysuiTestCase {
         NotificationRemoteInputManager.FORCE_REMOTE_INPUT_HISTORY = true;
         when(mController.isSpinning(mEntry.getKey())).thenReturn(true);
 
+        assertTrue(mRemoteInputManager.shouldKeepForRemoteInputHistory(mEntry));
         assertTrue(mRemoteInputHistoryExtender.shouldExtendLifetime(mEntry));
     }
 
@@ -153,6 +158,7 @@ public class NotificationRemoteInputManagerTest extends SysuiTestCase {
         NotificationRemoteInputManager.FORCE_REMOTE_INPUT_HISTORY = true;
         mEntry.lastRemoteInputSent = SystemClock.elapsedRealtime();
 
+        assertTrue(mRemoteInputManager.shouldKeepForRemoteInputHistory(mEntry));
         assertTrue(mRemoteInputHistoryExtender.shouldExtendLifetime(mEntry));
     }
 
@@ -161,6 +167,7 @@ public class NotificationRemoteInputManagerTest extends SysuiTestCase {
         NotificationRemoteInputManager.FORCE_REMOTE_INPUT_HISTORY = true;
         when(mSmartReplyController.isSendingSmartReply(mEntry.getKey())).thenReturn(true);
 
+        assertTrue(mRemoteInputManager.shouldKeepForSmartReplyHistory(mEntry));
         assertTrue(mSmartReplyHistoryExtender.shouldExtendLifetime(mEntry));
     }
 
@@ -185,6 +192,7 @@ public class NotificationRemoteInputManagerTest extends SysuiTestCase {
                 NotificationLockscreenUserManager lockscreenUserManager,
                 SmartReplyController smartReplyController,
                 NotificationEntryManager notificationEntryManager,
+                RemoteInputNotificationRebuilder rebuilder,
                 Lazy<Optional<StatusBar>> statusBarOptionalLazy,
                 StatusBarStateController statusBarStateController,
                 Handler mainHandler,
@@ -198,6 +206,7 @@ public class NotificationRemoteInputManagerTest extends SysuiTestCase {
                     lockscreenUserManager,
                     smartReplyController,
                     notificationEntryManager,
+                    rebuilder,
                     statusBarOptionalLazy,
                     statusBarStateController,
                     mainHandler,
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/SmartReplyControllerTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/SmartReplyControllerTest.java
index 0a61cdbdb4a6..99c965a9e57f 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/SmartReplyControllerTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/SmartReplyControllerTest.java
@@ -42,7 +42,6 @@ import com.android.systemui.dump.DumpManager;
 import com.android.systemui.flags.FeatureFlags;
 import com.android.systemui.plugins.statusbar.StatusBarStateController;
 import com.android.systemui.statusbar.notification.NotificationEntryManager;
-import com.android.systemui.statusbar.notification.collection.NotifPipeline;
 import com.android.systemui.statusbar.notification.collection.NotificationEntry;
 import com.android.systemui.statusbar.notification.collection.NotificationEntryBuilder;
 import com.android.systemui.statusbar.phone.StatusBar;
@@ -100,7 +99,7 @@ public class SmartReplyControllerTest extends SysuiTestCase {
                 mock(FeatureFlags.class),
                 mock(NotificationLockscreenUserManager.class), mSmartReplyController,
                 mNotificationEntryManager,
-                mock(NotifPipeline.class),
+                new RemoteInputNotificationRebuilder(mContext),
                 () -> Optional.of(mock(StatusBar.class)),
                 mStatusBarStateController,
                 Handler.createAsync(Looper.myLooper()),
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
new file mode 100644
index 000000000000..0ce6ada51f23
--- /dev/null
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/coordinator/RemoteInputCoordinatorTest.kt
@@ -0,0 +1,145 @@
+/*
+ * Copyright (C) 2021 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.statusbar.notification.collection.coordinator
+
+import android.os.Handler
+import android.service.notification.StatusBarNotification
+import android.testing.AndroidTestingRunner
+import android.testing.TestableLooper.RunWithLooper
+import androidx.test.filters.SmallTest
+import com.android.systemui.SysuiTestCase
+import com.android.systemui.dump.DumpManager
+import com.android.systemui.statusbar.NotificationRemoteInputManager
+import com.android.systemui.statusbar.NotificationRemoteInputManager.RemoteInputListener
+import com.android.systemui.statusbar.RemoteInputNotificationRebuilder
+import com.android.systemui.statusbar.SmartReplyController
+import com.android.systemui.statusbar.notification.collection.NotifPipeline
+import com.android.systemui.statusbar.notification.collection.NotificationEntry
+import com.android.systemui.statusbar.notification.collection.NotificationEntryBuilder
+import com.android.systemui.statusbar.notification.collection.notifcollection.InternalNotifUpdater
+import com.android.systemui.statusbar.notification.collection.notifcollection.NotifCollectionListener
+import com.android.systemui.statusbar.notification.collection.notifcollection.NotifLifetimeExtender
+import com.android.systemui.statusbar.notification.collection.notifcollection.NotifLifetimeExtender.OnEndLifetimeExtensionCallback
+import com.android.systemui.util.mockito.any
+import com.android.systemui.util.mockito.withArgCaptor
+import com.google.common.truth.Truth.assertThat
+import org.junit.Before
+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.verify
+import org.mockito.MockitoAnnotations.initMocks
+
+@SmallTest
+@RunWith(AndroidTestingRunner::class)
+@RunWithLooper
+class RemoteInputCoordinatorTest : SysuiTestCase() {
+    private lateinit var coordinator: RemoteInputCoordinator
+    private lateinit var listener: RemoteInputListener
+    private lateinit var collectionListener: NotifCollectionListener
+
+    private lateinit var entry1: NotificationEntry
+    private lateinit var entry2: NotificationEntry
+
+    @Mock private lateinit var lifetimeExtensionCallback: OnEndLifetimeExtensionCallback
+    @Mock private lateinit var rebuilder: RemoteInputNotificationRebuilder
+    @Mock private lateinit var remoteInputManager: NotificationRemoteInputManager
+    @Mock private lateinit var mainHandler: Handler
+    @Mock private lateinit var smartReplyController: SmartReplyController
+    @Mock private lateinit var pipeline: NotifPipeline
+    @Mock private lateinit var notifUpdater: InternalNotifUpdater
+    @Mock private lateinit var dumpManager: DumpManager
+    @Mock private lateinit var sbn: StatusBarNotification
+
+    @Before
+    fun setUp() {
+        initMocks(this)
+        coordinator = RemoteInputCoordinator(
+                dumpManager,
+                rebuilder,
+                remoteInputManager,
+                mainHandler,
+                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())
+        }
+        collectionListener = withArgCaptor {
+            verify(pipeline).addCollectionListener(capture())
+        }
+        entry1 = NotificationEntryBuilder().setId(1).build()
+        entry2 = NotificationEntryBuilder().setId(2).build()
+        `when`(rebuilder.rebuildForCanceledSmartReplies(any())).thenReturn(sbn)
+        `when`(rebuilder.rebuildForRemoteInputReply(any())).thenReturn(sbn)
+        `when`(rebuilder.rebuildForSendingSmartReply(any(), any())).thenReturn(sbn)
+    }
+
+    val remoteInputActiveExtender get() = coordinator.mRemoteInputActiveExtender
+    val remoteInputHistoryExtender get() = coordinator.mRemoteInputHistoryExtender
+    val smartReplyHistoryExtender get() = coordinator.mSmartReplyHistoryExtender
+
+    @Test
+    fun testRemoteInputActive() {
+        `when`(remoteInputManager.isRemoteInputActive(entry1)).thenReturn(true)
+        assertThat(remoteInputActiveExtender.shouldExtendLifetime(entry1, 0)).isTrue()
+        assertThat(remoteInputHistoryExtender.shouldExtendLifetime(entry1, 0)).isFalse()
+        assertThat(smartReplyHistoryExtender.shouldExtendLifetime(entry1, 0)).isFalse()
+        assertThat(listener.isNotificationKeptForRemoteInputHistory(entry1.key)).isFalse()
+    }
+
+    @Test
+    fun testRemoteInputHistory() {
+        `when`(remoteInputManager.shouldKeepForRemoteInputHistory(entry1)).thenReturn(true)
+        assertThat(remoteInputActiveExtender.shouldExtendLifetime(entry1, 0)).isFalse()
+        assertThat(remoteInputHistoryExtender.shouldExtendLifetime(entry1, 0)).isTrue()
+        assertThat(smartReplyHistoryExtender.shouldExtendLifetime(entry1, 0)).isFalse()
+        assertThat(listener.isNotificationKeptForRemoteInputHistory(entry1.key)).isTrue()
+    }
+
+    @Test
+    fun testSmartReplyHistory() {
+        `when`(remoteInputManager.shouldKeepForSmartReplyHistory(entry1)).thenReturn(true)
+        assertThat(remoteInputActiveExtender.shouldExtendLifetime(entry1, 0)).isFalse()
+        assertThat(remoteInputHistoryExtender.shouldExtendLifetime(entry1, 0)).isFalse()
+        assertThat(smartReplyHistoryExtender.shouldExtendLifetime(entry1, 0)).isTrue()
+        assertThat(listener.isNotificationKeptForRemoteInputHistory(entry1.key)).isTrue()
+    }
+
+    @Test
+    fun testNotificationWithRemoteInputActiveIsRemovedOnCollapse() {
+        `when`(remoteInputManager.isRemoteInputActive(entry1)).thenReturn(true)
+        assertThat(remoteInputActiveExtender.isExtending(entry1.key)).isFalse()
+
+        // Nothing should happen on panel collapse before we start extending the lifetime
+        listener.onPanelCollapsed()
+        assertThat(remoteInputActiveExtender.isExtending(entry1.key)).isFalse()
+        verify(lifetimeExtensionCallback, never()).onEndLifetimeExtension(any(), any())
+
+        // Start extending lifetime & validate that the extension is ended
+        assertThat(remoteInputActiveExtender.shouldExtendLifetime(entry1, 0)).isTrue()
+        assertThat(remoteInputActiveExtender.isExtending(entry1.key)).isTrue()
+        listener.onPanelCollapsed()
+        verify(lifetimeExtensionCallback).onEndLifetimeExtension(remoteInputActiveExtender, entry1)
+        assertThat(remoteInputActiveExtender.isExtending(entry1.key)).isFalse()
+    }
+}
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/notifcollection/SelfTrackingLifetimeExtenderTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/notifcollection/SelfTrackingLifetimeExtenderTest.kt
new file mode 100644
index 000000000000..37ad8357aa95
--- /dev/null
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/notifcollection/SelfTrackingLifetimeExtenderTest.kt
@@ -0,0 +1,230 @@
+/*
+ * Copyright (C) 2021 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.statusbar.notification.collection.notifcollection
+
+import android.os.Handler
+import android.testing.AndroidTestingRunner
+import android.testing.TestableLooper.RunWithLooper
+import androidx.test.filters.SmallTest
+import com.android.systemui.SysuiTestCase
+import com.android.systemui.statusbar.notification.collection.NotificationEntry
+import com.android.systemui.statusbar.notification.collection.NotificationEntryBuilder
+import com.android.systemui.statusbar.notification.collection.notifcollection.NotifLifetimeExtender.OnEndLifetimeExtensionCallback
+import com.android.systemui.util.mockito.withArgCaptor
+import com.google.common.truth.Truth.assertThat
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mockito.Mock
+import org.mockito.Mockito.`when`
+import org.mockito.Mockito.any
+import org.mockito.Mockito.anyLong
+import org.mockito.Mockito.eq
+import org.mockito.Mockito.never
+import org.mockito.Mockito.times
+import org.mockito.Mockito.verify
+import org.mockito.MockitoAnnotations.initMocks
+import java.util.function.Consumer
+import java.util.function.Predicate
+
+@SmallTest
+@RunWith(AndroidTestingRunner::class)
+@RunWithLooper
+class SelfTrackingLifetimeExtenderTest : SysuiTestCase() {
+    private lateinit var extender: TestableSelfTrackingLifetimeExtender
+
+    private lateinit var entry1: NotificationEntry
+    private lateinit var entry2: NotificationEntry
+
+    @Mock
+    private lateinit var callback: OnEndLifetimeExtensionCallback
+    @Mock
+    private lateinit var mainHandler: Handler
+    @Mock
+    private lateinit var shouldExtend: Predicate<NotificationEntry>
+    @Mock
+    private lateinit var onStarted: Consumer<NotificationEntry>
+    @Mock
+    private lateinit var onCanceled: Consumer<NotificationEntry>
+
+    @Before
+    fun setUp() {
+        initMocks(this)
+        extender = TestableSelfTrackingLifetimeExtender()
+        extender.setCallback(callback)
+        entry1 = NotificationEntryBuilder().setId(1).build()
+        entry2 = NotificationEntryBuilder().setId(2).build()
+    }
+
+    @Test
+    fun testName() {
+        assertThat(extender.name).isEqualTo("Testable")
+    }
+
+    @Test
+    fun testNoExtend() {
+        `when`(shouldExtend.test(entry1)).thenReturn(false)
+        assertThat(extender.shouldExtendLifetime(entry1, 0)).isFalse()
+        assertThat(extender.isExtending(entry1.key)).isFalse()
+        verify(onStarted, never()).accept(entry1)
+        verify(onCanceled, never()).accept(entry1)
+    }
+
+    @Test
+    fun testExtendThenCancelForRepost() {
+        `when`(shouldExtend.test(entry1)).thenReturn(true)
+        assertThat(extender.shouldExtendLifetime(entry1, 0)).isTrue()
+        verify(onStarted).accept(entry1)
+        verify(onCanceled, never()).accept(entry1)
+        assertThat(extender.isExtending(entry1.key)).isTrue()
+        extender.cancelLifetimeExtension(entry1)
+        verify(onCanceled).accept(entry1)
+    }
+
+    @Test
+    fun testExtendThenCancel_thenEndDoesNothing() {
+        testExtendThenCancelForRepost()
+        assertThat(extender.isExtending(entry1.key)).isFalse()
+
+        extender.endLifetimeExtension(entry1.key)
+        extender.endLifetimeExtensionAfterDelay(entry1.key, 1000)
+        verify(callback, never()).onEndLifetimeExtension(any(), any())
+        verify(mainHandler, never()).postDelayed(any(), anyLong())
+    }
+
+    @Test
+    fun testExtendThenEnd() {
+        `when`(shouldExtend.test(entry1)).thenReturn(true)
+        assertThat(extender.shouldExtendLifetime(entry1, 0)).isTrue()
+        verify(onStarted).accept(entry1)
+        assertThat(extender.isExtending(entry1.key)).isTrue()
+        extender.endLifetimeExtension(entry1.key)
+        verify(callback).onEndLifetimeExtension(extender, entry1)
+        verify(onCanceled, never()).accept(entry1)
+    }
+
+    @Test
+    fun testExtendThenEndAfterDelay() {
+        `when`(shouldExtend.test(entry1)).thenReturn(true)
+        assertThat(extender.shouldExtendLifetime(entry1, 0)).isTrue()
+        verify(onStarted).accept(entry1)
+        assertThat(extender.isExtending(entry1.key)).isTrue()
+
+        // Call the method and capture the posted runnable
+        extender.endLifetimeExtensionAfterDelay(entry1.key, 1234)
+        val runnable = withArgCaptor<Runnable> {
+            verify(mainHandler).postDelayed(capture(), eq(1234.toLong()))
+        }
+        assertThat(extender.isExtending(entry1.key)).isTrue()
+        verify(callback, never()).onEndLifetimeExtension(any(), any())
+
+        // now run the posted runnable and ensure it works as expected
+        runnable.run()
+        verify(callback).onEndLifetimeExtension(extender, entry1)
+        assertThat(extender.isExtending(entry1.key)).isFalse()
+        verify(onCanceled, never()).accept(entry1)
+    }
+
+    @Test
+    fun testExtendThenEndAll() {
+        `when`(shouldExtend.test(entry1)).thenReturn(true)
+        `when`(shouldExtend.test(entry2)).thenReturn(true)
+        assertThat(extender.shouldExtendLifetime(entry1, 0)).isTrue()
+        verify(onStarted).accept(entry1)
+        assertThat(extender.isExtending(entry1.key)).isTrue()
+        assertThat(extender.isExtending(entry2.key)).isFalse()
+        assertThat(extender.shouldExtendLifetime(entry2, 0)).isTrue()
+        verify(onStarted).accept(entry2)
+        assertThat(extender.isExtending(entry1.key)).isTrue()
+        assertThat(extender.isExtending(entry2.key)).isTrue()
+        extender.endAllLifetimeExtensions()
+        verify(callback).onEndLifetimeExtension(extender, entry1)
+        verify(callback).onEndLifetimeExtension(extender, entry2)
+        verify(onCanceled, never()).accept(entry1)
+        verify(onCanceled, never()).accept(entry2)
+    }
+
+    @Test
+    fun testExtendWithinEndCanReExtend() {
+        `when`(shouldExtend.test(entry1)).thenReturn(true)
+        assertThat(extender.shouldExtendLifetime(entry1, 0)).isTrue()
+        verify(onStarted, times(1)).accept(entry1)
+
+        `when`(callback.onEndLifetimeExtension(extender, entry1)).thenAnswer {
+            assertThat(extender.shouldExtendLifetime(entry1, 0)).isTrue()
+        }
+        extender.endLifetimeExtension(entry1.key)
+        verify(onStarted, times(2)).accept(entry1)
+        assertThat(extender.isExtending(entry1.key)).isTrue()
+    }
+
+    @Test
+    fun testExtendWithinEndCanNotReExtend() {
+        `when`(shouldExtend.test(entry1)).thenReturn(true, false)
+        assertThat(extender.shouldExtendLifetime(entry1, 0)).isTrue()
+        verify(onStarted, times(1)).accept(entry1)
+
+        `when`(callback.onEndLifetimeExtension(extender, entry1)).thenAnswer {
+            assertThat(extender.shouldExtendLifetime(entry1, 0)).isFalse()
+        }
+        extender.endLifetimeExtension(entry1.key)
+        verify(onStarted, times(1)).accept(entry1)
+        assertThat(extender.isExtending(entry1.key)).isFalse()
+    }
+
+    @Test
+    fun testExtendWithinEndAllCanReExtend() {
+        `when`(shouldExtend.test(entry1)).thenReturn(true)
+        assertThat(extender.shouldExtendLifetime(entry1, 0)).isTrue()
+        verify(onStarted, times(1)).accept(entry1)
+
+        `when`(callback.onEndLifetimeExtension(extender, entry1)).thenAnswer {
+            assertThat(extender.shouldExtendLifetime(entry1, 0)).isTrue()
+        }
+        extender.endAllLifetimeExtensions()
+        verify(onStarted, times(2)).accept(entry1)
+        assertThat(extender.isExtending(entry1.key)).isTrue()
+    }
+
+    @Test
+    fun testExtendWithinEndAllCanNotReExtend() {
+        `when`(shouldExtend.test(entry1)).thenReturn(true, false)
+        assertThat(extender.shouldExtendLifetime(entry1, 0)).isTrue()
+        verify(onStarted, times(1)).accept(entry1)
+
+        `when`(callback.onEndLifetimeExtension(extender, entry1)).thenAnswer {
+            assertThat(extender.shouldExtendLifetime(entry1, 0)).isFalse()
+        }
+        extender.endAllLifetimeExtensions()
+        verify(onStarted, times(1)).accept(entry1)
+        assertThat(extender.isExtending(entry1.key)).isFalse()
+    }
+
+    inner class TestableSelfTrackingLifetimeExtender(debug: Boolean = false) :
+            SelfTrackingLifetimeExtender("Test", "Testable", debug, mainHandler) {
+
+        override fun queryShouldExtendLifetime(entry: NotificationEntry) =
+                shouldExtend.test(entry)
+
+        override fun onStartedLifetimeExtension(entry: NotificationEntry) {
+            onStarted.accept(entry)
+        }
+
+        override fun onCanceledLifetimeExtension(entry: NotificationEntry) {
+            onCanceled.accept(entry)
+        }
+    }
+}