Merge "Create bugfix flag to enable system dialog transitions in Desktop Mode." into main
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellModule.java b/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellModule.java
index 44fce81..601cf70 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellModule.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellModule.java
@@ -67,6 +67,7 @@
 import com.android.wm.shell.desktopmode.CloseDesktopTaskTransitionHandler;
 import com.android.wm.shell.desktopmode.DefaultDragToDesktopTransitionHandler;
 import com.android.wm.shell.desktopmode.DesktopActivityOrientationChangeHandler;
+import com.android.wm.shell.desktopmode.DesktopBackNavigationTransitionHandler;
 import com.android.wm.shell.desktopmode.DesktopDisplayEventHandler;
 import com.android.wm.shell.desktopmode.DesktopImmersiveController;
 import com.android.wm.shell.desktopmode.DesktopMixedTransitionHandler;
@@ -915,6 +916,16 @@
 
     @WMSingleton
     @Provides
+    static DesktopBackNavigationTransitionHandler provideDesktopBackNavigationTransitionHandler(
+            @ShellMainThread ShellExecutor mainExecutor,
+            @ShellAnimationThread ShellExecutor animExecutor,
+            DisplayController displayController) {
+        return new DesktopBackNavigationTransitionHandler(mainExecutor, animExecutor,
+                displayController);
+    }
+
+    @WMSingleton
+    @Provides
     static DesktopModeDragAndDropTransitionHandler provideDesktopModeDragAndDropTransitionHandler(
             Transitions transitions) {
         return new DesktopModeDragAndDropTransitionHandler(transitions);
@@ -964,6 +975,7 @@
             Optional<DesktopRepository> desktopRepository,
             Transitions transitions,
             ShellTaskOrganizer shellTaskOrganizer,
+            Optional<DesktopMixedTransitionHandler> desktopMixedTransitionHandler,
             ShellInit shellInit) {
         return desktopRepository.flatMap(
                 repository ->
@@ -973,6 +985,7 @@
                                         repository,
                                         transitions,
                                         shellTaskOrganizer,
+                                        desktopMixedTransitionHandler.get(),
                                         shellInit)));
     }
 
@@ -985,6 +998,7 @@
             FreeformTaskTransitionHandler freeformTaskTransitionHandler,
             CloseDesktopTaskTransitionHandler closeDesktopTaskTransitionHandler,
             Optional<DesktopImmersiveController> desktopImmersiveController,
+            DesktopBackNavigationTransitionHandler desktopBackNavigationTransitionHandler,
             InteractionJankMonitor interactionJankMonitor,
             @ShellMainThread Handler handler,
             ShellInit shellInit,
@@ -1001,6 +1015,7 @@
                         freeformTaskTransitionHandler,
                         closeDesktopTaskTransitionHandler,
                         desktopImmersiveController.get(),
+                        desktopBackNavigationTransitionHandler,
                         interactionJankMonitor,
                         handler,
                         shellInit,
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopBackNavigationTransitionHandler.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopBackNavigationTransitionHandler.kt
new file mode 100644
index 0000000..83b0f84
--- /dev/null
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopBackNavigationTransitionHandler.kt
@@ -0,0 +1,97 @@
+/*
+ * Copyright (C) 2024 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.wm.shell.desktopmode
+
+import android.animation.Animator
+import android.app.WindowConfiguration.WINDOWING_MODE_FREEFORM
+import android.os.IBinder
+import android.util.DisplayMetrics
+import android.view.SurfaceControl.Transaction
+import android.window.TransitionInfo
+import android.window.TransitionRequestInfo
+import android.window.WindowContainerTransaction
+import com.android.wm.shell.common.DisplayController
+import com.android.wm.shell.common.ShellExecutor
+import com.android.wm.shell.shared.TransitionUtil
+import com.android.wm.shell.shared.animation.MinimizeAnimator.create
+import com.android.wm.shell.transition.Transitions
+
+/**
+ * The [Transitions.TransitionHandler] that handles transitions for tasks that are closing or going
+ * to back as part of back navigation. This handler is used only for animating transitions.
+ */
+class DesktopBackNavigationTransitionHandler(
+    private val mainExecutor: ShellExecutor,
+    private val animExecutor: ShellExecutor,
+    private val displayController: DisplayController,
+) : Transitions.TransitionHandler {
+
+    /** Shouldn't handle anything */
+    override fun handleRequest(
+        transition: IBinder,
+        request: TransitionRequestInfo,
+    ): WindowContainerTransaction? = null
+
+    /** Animates a transition with minimizing tasks */
+    override fun startAnimation(
+        transition: IBinder,
+        info: TransitionInfo,
+        startTransaction: Transaction,
+        finishTransaction: Transaction,
+        finishCallback: Transitions.TransitionFinishCallback,
+    ): Boolean {
+        if (!TransitionUtil.isClosingType(info.type)) return false
+
+        val animations = mutableListOf<Animator>()
+        val onAnimFinish: (Animator) -> Unit = { animator ->
+            mainExecutor.execute {
+                // Animation completed
+                animations.remove(animator)
+                if (animations.isEmpty()) {
+                    // All animations completed, finish the transition
+                    finishCallback.onTransitionFinished(/* wct= */ null)
+                }
+            }
+        }
+
+        animations +=
+            info.changes
+                .filter {
+                    it.mode == info.type &&
+                            it.taskInfo?.windowingMode == WINDOWING_MODE_FREEFORM
+                }
+                .mapNotNull { createMinimizeAnimation(it, finishTransaction, onAnimFinish) }
+        if (animations.isEmpty()) return false
+        animExecutor.execute { animations.forEach(Animator::start) }
+        return true
+    }
+
+    private fun createMinimizeAnimation(
+        change: TransitionInfo.Change,
+        finishTransaction: Transaction,
+        onAnimFinish: (Animator) -> Unit
+    ): Animator? {
+        val t = Transaction()
+        val sc = change.leash
+        finishTransaction.hide(sc)
+        val displayMetrics: DisplayMetrics? =
+            change.taskInfo?.let {
+                displayController.getDisplayContext(it.displayId)?.getResources()?.displayMetrics
+            }
+        return displayMetrics?.let { create(it, change, t, onAnimFinish) }
+    }
+}
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopMixedTransitionHandler.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopMixedTransitionHandler.kt
index 01c680d..2001f97 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopMixedTransitionHandler.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopMixedTransitionHandler.kt
@@ -52,6 +52,7 @@
     private val freeformTaskTransitionHandler: FreeformTaskTransitionHandler,
     private val closeDesktopTaskTransitionHandler: CloseDesktopTaskTransitionHandler,
     private val desktopImmersiveController: DesktopImmersiveController,
+    private val desktopBackNavigationTransitionHandler: DesktopBackNavigationTransitionHandler,
     private val interactionJankMonitor: InteractionJankMonitor,
     @ShellMainThread private val handler: Handler,
     shellInit: ShellInit,
@@ -161,6 +162,14 @@
                 finishTransaction,
                 finishCallback
             )
+            is PendingMixedTransition.Minimize -> animateMinimizeTransition(
+                pending,
+                transition,
+                info,
+                startTransaction,
+                finishTransaction,
+                finishCallback
+            )
         }
     }
 
@@ -272,6 +281,42 @@
         )
     }
 
+    private fun animateMinimizeTransition(
+        pending: PendingMixedTransition.Minimize,
+        transition: IBinder,
+        info: TransitionInfo,
+        startTransaction: SurfaceControl.Transaction,
+        finishTransaction: SurfaceControl.Transaction,
+        finishCallback: TransitionFinishCallback,
+    ): Boolean {
+        if (!DesktopModeFlags.ENABLE_DESKTOP_WINDOWING_BACK_NAVIGATION.isTrue) return false
+
+        val minimizeChange = findDesktopTaskChange(info, pending.minimizingTask)
+        if (minimizeChange == null) {
+            logW("Should have minimizing desktop task")
+            return false
+        }
+        if (pending.isLastTask) {
+            // Dispatch close desktop task animation to the default transition handlers.
+            return dispatchToLeftoverHandler(
+                transition,
+                info,
+                startTransaction,
+                finishTransaction,
+                finishCallback
+            )
+        }
+
+        // Animate minimizing desktop task transition with [DesktopBackNavigationTransitionHandler].
+        return desktopBackNavigationTransitionHandler.startAnimation(
+            transition,
+            info,
+            startTransaction,
+            finishTransaction,
+            finishCallback,
+        )
+    }
+
     override fun onTransitionConsumed(
         transition: IBinder,
         aborted: Boolean,
@@ -400,6 +445,14 @@
             val minimizingTask: Int?,
             val exitingImmersiveTask: Int?,
         ) : PendingMixedTransition()
+
+        /** A task is minimizing. This should be used for task going to back and some closing cases
+         * with back navigation. */
+        data class Minimize(
+            override val transition: IBinder,
+            val minimizingTask: Int,
+            val isLastTask: Boolean,
+        ) : PendingMixedTransition()
     }
 
     private fun logV(msg: String, vararg arguments: Any?) {
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksController.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksController.kt
index 3a7cc49..223038f 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksController.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksController.kt
@@ -1290,7 +1290,11 @@
                     // Check if freeform task launch during recents should be handled
                     shouldHandleMidRecentsFreeformLaunch -> handleMidRecentsFreeformTaskLaunch(task)
                     // Check if the closing task needs to be handled
-                    TransitionUtil.isClosingType(request.type) -> handleTaskClosing(task)
+                    TransitionUtil.isClosingType(request.type) -> handleTaskClosing(
+                        task,
+                        transition,
+                        request.type
+                    )
                     // Check if the top task shouldn't be allowed to enter desktop mode
                     isIncompatibleTask(task) -> handleIncompatibleTaskLaunch(task)
                     // Check if fullscreen task should be updated
@@ -1620,7 +1624,7 @@
     }
 
     /** Handle task closing by removing wallpaper activity if it's the last active task */
-    private fun handleTaskClosing(task: RunningTaskInfo): WindowContainerTransaction? {
+    private fun handleTaskClosing(task: RunningTaskInfo, transition: IBinder, requestType: Int): WindowContainerTransaction? {
         logV("handleTaskClosing")
         if (!isDesktopModeShowing(task.displayId))
             return null
@@ -1636,8 +1640,15 @@
         if (!DesktopModeFlags.ENABLE_DESKTOP_WINDOWING_BACK_NAVIGATION.isTrue()) {
             taskRepository.addClosingTask(task.displayId, task.taskId)
             desktopTilingDecorViewModel.removeTaskIfTiled(task.displayId, task.taskId)
+        } else if (requestType == TRANSIT_CLOSE) {
+            // Handle closing tasks, tasks that are going to back are handled in
+            // [DesktopTasksTransitionObserver].
+            desktopMixedTransitionHandler.addPendingMixedTransition(
+                DesktopMixedTransitionHandler.PendingMixedTransition.Minimize(
+                    transition, task.taskId, taskRepository.getVisibleTaskCount(task.displayId) == 1
+                )
+            )
         }
-
         taskbarDesktopTaskListener?.onTaskbarCornerRoundingUpdate(
             doesAnyTaskRequireTaskbarRounding(
                 task.displayId,
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksTransitionObserver.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksTransitionObserver.kt
index d1534da..c39c715 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksTransitionObserver.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksTransitionObserver.kt
@@ -46,6 +46,7 @@
     private val desktopRepository: DesktopRepository,
     private val transitions: Transitions,
     private val shellTaskOrganizer: ShellTaskOrganizer,
+    private val desktopMixedTransitionHandler: DesktopMixedTransitionHandler,
     shellInit: ShellInit
 ) : Transitions.TransitionObserver {
 
@@ -71,7 +72,7 @@
         // TODO: b/332682201 Update repository state
         updateWallpaperToken(info)
         if (DesktopModeFlags.ENABLE_DESKTOP_WINDOWING_BACK_NAVIGATION.isTrue()) {
-            handleBackNavigation(info)
+            handleBackNavigation(transition, info)
             removeTaskIfNeeded(info)
         }
         removeWallpaperOnLastTaskClosingIfNeeded(transition, info)
@@ -95,7 +96,7 @@
         }
     }
 
-    private fun handleBackNavigation(info: TransitionInfo) {
+    private fun handleBackNavigation(transition: IBinder, info: TransitionInfo) {
         // When default back navigation happens, transition type is TO_BACK and the change is
         // TO_BACK. Mark the task going to back as minimized.
         if (info.type == TRANSIT_TO_BACK) {
@@ -105,10 +106,14 @@
                     continue
                 }
 
-                if (desktopRepository.getVisibleTaskCount(taskInfo.displayId) > 0 &&
+                val visibleTaskCount = desktopRepository.getVisibleTaskCount(taskInfo.displayId)
+                if (visibleTaskCount > 0 &&
                     change.mode == TRANSIT_TO_BACK &&
                     taskInfo.windowingMode == WINDOWING_MODE_FREEFORM) {
                     desktopRepository.minimizeTask(taskInfo.displayId, taskInfo.taskId)
+                    desktopMixedTransitionHandler.addPendingMixedTransition(
+                        DesktopMixedTransitionHandler.PendingMixedTransition.Minimize(
+                            transition, taskInfo.taskId, visibleTaskCount == 1))
                 }
             }
         }
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopBackNavigationTransitionHandlerTest.kt b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopBackNavigationTransitionHandlerTest.kt
new file mode 100644
index 0000000..6df8d6f
--- /dev/null
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopBackNavigationTransitionHandlerTest.kt
@@ -0,0 +1,173 @@
+/*
+ * Copyright (C) 2024 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.wm.shell.desktopmode
+
+import android.app.ActivityManager.RunningTaskInfo
+import android.app.WindowConfiguration.ACTIVITY_TYPE_STANDARD
+import android.app.WindowConfiguration.WINDOWING_MODE_FREEFORM
+import android.app.WindowConfiguration.WINDOWING_MODE_FULLSCREEN
+import android.app.WindowConfiguration.WindowingMode
+import android.testing.AndroidTestingRunner
+import android.testing.TestableLooper.RunWithLooper
+import android.view.SurfaceControl
+import android.view.WindowManager
+import android.view.WindowManager.TRANSIT_CLOSE
+import android.window.TransitionInfo
+import androidx.test.filters.SmallTest
+import com.android.wm.shell.ShellTestCase
+import com.android.wm.shell.TestRunningTaskInfoBuilder
+import com.android.wm.shell.common.DisplayController
+import com.android.wm.shell.common.ShellExecutor
+import org.junit.Assert.assertFalse
+import org.junit.Assert.assertNull
+import org.junit.Assert.assertTrue
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mockito.kotlin.any
+import org.mockito.kotlin.mock
+import org.mockito.kotlin.whenever
+
+@SmallTest
+@RunWithLooper
+@RunWith(AndroidTestingRunner::class)
+class DesktopBackNavigationTransitionHandlerTest : ShellTestCase() {
+
+    private val testExecutor = mock<ShellExecutor>()
+    private val closingTaskLeash = mock<SurfaceControl>()
+    private val displayController = mock<DisplayController>()
+
+    private lateinit var handler: DesktopBackNavigationTransitionHandler
+
+    @Before
+    fun setUp() {
+        handler =
+            DesktopBackNavigationTransitionHandler(
+                testExecutor,
+                testExecutor,
+                displayController
+            )
+        whenever(displayController.getDisplayContext(any())).thenReturn(mContext)
+    }
+
+    @Test
+    fun handleRequest_returnsNull() {
+        assertNull(handler.handleRequest(mock(), mock()))
+    }
+
+    @Test
+    fun startAnimation_openTransition_returnsFalse() {
+        val animates =
+            handler.startAnimation(
+                transition = mock(),
+                info =
+                createTransitionInfo(
+                    type = WindowManager.TRANSIT_OPEN,
+                    task = createTask(WINDOWING_MODE_FREEFORM)
+                ),
+                startTransaction = mock(),
+                finishTransaction = mock(),
+                finishCallback = {}
+            )
+
+        assertFalse("Should not animate open transition", animates)
+    }
+
+    @Test
+    fun startAnimation_toBackTransitionFullscreenTask_returnsFalse() {
+        val animates =
+            handler.startAnimation(
+                transition = mock(),
+                info = createTransitionInfo(task = createTask(WINDOWING_MODE_FULLSCREEN)),
+                startTransaction = mock(),
+                finishTransaction = mock(),
+                finishCallback = {}
+            )
+
+        assertFalse("Should not animate fullscreen task to back transition", animates)
+    }
+
+    @Test
+    fun startAnimation_toBackTransitionOpeningFreeformTask_returnsFalse() {
+        val animates =
+            handler.startAnimation(
+                transition = mock(),
+                info =
+                createTransitionInfo(
+                    changeMode = WindowManager.TRANSIT_OPEN,
+                    task = createTask(WINDOWING_MODE_FREEFORM)
+                ),
+                startTransaction = mock(),
+                finishTransaction = mock(),
+                finishCallback = {}
+            )
+
+        assertFalse("Should not animate opening freeform task to back transition", animates)
+    }
+
+    @Test
+    fun startAnimation_toBackTransitionToBackFreeformTask_returnsTrue() {
+        val animates =
+            handler.startAnimation(
+                transition = mock(),
+                info = createTransitionInfo(task = createTask(WINDOWING_MODE_FREEFORM)),
+                startTransaction = mock(),
+                finishTransaction = mock(),
+                finishCallback = {}
+            )
+
+        assertTrue("Should animate going to back freeform task close transition", animates)
+    }
+
+    @Test
+    fun startAnimation_closeTransitionClosingFreeformTask_returnsTrue() {
+        val animates =
+            handler.startAnimation(
+                transition = mock(),
+                info = createTransitionInfo(
+                    type = TRANSIT_CLOSE,
+                    changeMode = TRANSIT_CLOSE,
+                    task = createTask(WINDOWING_MODE_FREEFORM)
+                ),
+                startTransaction = mock(),
+                finishTransaction = mock(),
+                finishCallback = {}
+            )
+
+        assertTrue("Should animate going to back freeform task close transition", animates)
+    }
+    private fun createTransitionInfo(
+        type: Int = WindowManager.TRANSIT_TO_BACK,
+        changeMode: Int = WindowManager.TRANSIT_TO_BACK,
+        task: RunningTaskInfo
+    ): TransitionInfo =
+        TransitionInfo(type, 0 /* flags */).apply {
+            addChange(
+                TransitionInfo.Change(mock(), closingTaskLeash).apply {
+                    mode = changeMode
+                    parent = null
+                    taskInfo = task
+                }
+            )
+        }
+
+    private fun createTask(@WindowingMode windowingMode: Int): RunningTaskInfo =
+        TestRunningTaskInfoBuilder()
+            .setActivityType(ACTIVITY_TYPE_STANDARD)
+            .setWindowingMode(windowingMode)
+            .build()
+}
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopMixedTransitionHandlerTest.kt b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopMixedTransitionHandlerTest.kt
index b06c2da..f21f264 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopMixedTransitionHandlerTest.kt
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopMixedTransitionHandlerTest.kt
@@ -32,6 +32,7 @@
 import android.view.SurfaceControl
 import android.view.WindowManager
 import android.view.WindowManager.TRANSIT_OPEN
+import android.view.WindowManager.TRANSIT_TO_BACK
 import android.view.WindowManager.TransitionType
 import android.window.TransitionInfo
 import android.window.WindowContainerTransaction
@@ -77,16 +78,28 @@
 
     @JvmField @Rule val setFlagsRule = SetFlagsRule()
 
-    @Mock lateinit var transitions: Transitions
-    @Mock lateinit var desktopRepository: DesktopRepository
-    @Mock lateinit var freeformTaskTransitionHandler: FreeformTaskTransitionHandler
-    @Mock lateinit var closeDesktopTaskTransitionHandler: CloseDesktopTaskTransitionHandler
-    @Mock lateinit var desktopImmersiveController: DesktopImmersiveController
-    @Mock lateinit var interactionJankMonitor: InteractionJankMonitor
-    @Mock lateinit var mockHandler: Handler
-    @Mock lateinit var closingTaskLeash: SurfaceControl
-    @Mock lateinit var shellInit: ShellInit
-    @Mock lateinit var rootTaskDisplayAreaOrganizer: RootTaskDisplayAreaOrganizer
+    @Mock
+    lateinit var transitions: Transitions
+    @Mock
+    lateinit var desktopRepository: DesktopRepository
+    @Mock
+    lateinit var freeformTaskTransitionHandler: FreeformTaskTransitionHandler
+    @Mock
+    lateinit var closeDesktopTaskTransitionHandler: CloseDesktopTaskTransitionHandler
+    @Mock
+    lateinit var desktopBackNavigationTransitionHandler: DesktopBackNavigationTransitionHandler
+    @Mock
+    lateinit var desktopImmersiveController: DesktopImmersiveController
+    @Mock
+    lateinit var interactionJankMonitor: InteractionJankMonitor
+    @Mock
+    lateinit var mockHandler: Handler
+    @Mock
+    lateinit var closingTaskLeash: SurfaceControl
+    @Mock
+    lateinit var shellInit: ShellInit
+    @Mock
+    lateinit var rootTaskDisplayAreaOrganizer: RootTaskDisplayAreaOrganizer
 
     private lateinit var mixedHandler: DesktopMixedTransitionHandler
 
@@ -100,6 +113,7 @@
                 freeformTaskTransitionHandler,
                 closeDesktopTaskTransitionHandler,
                 desktopImmersiveController,
+                desktopBackNavigationTransitionHandler,
                 interactionJankMonitor,
                 mockHandler,
                 shellInit,
@@ -595,6 +609,87 @@
         assertThat(mixedHandler.pendingMixedTransitions).isEmpty()
     }
 
+    @Test
+    @EnableFlags(Flags.FLAG_ENABLE_DESKTOP_WINDOWING_BACK_NAVIGATION)
+    fun startAnimation_withMinimizingDesktopTask_callsBackNavigationHandler() {
+        val minimizingTask = createTask(WINDOWING_MODE_FREEFORM)
+        val transition = Binder()
+        whenever(desktopRepository.getExpandedTaskCount(any())).thenReturn(2)
+        whenever(
+            desktopBackNavigationTransitionHandler.startAnimation(any(), any(), any(), any(), any())
+        )
+            .thenReturn(true)
+        mixedHandler.addPendingMixedTransition(
+            PendingMixedTransition.Minimize(
+                transition = transition,
+                minimizingTask = minimizingTask.taskId,
+                isLastTask = false,
+            )
+        )
+
+        val minimizingTaskChange = createChange(minimizingTask)
+        val started = mixedHandler.startAnimation(
+            transition = transition,
+            info =
+                createTransitionInfo(
+                TRANSIT_TO_BACK,
+                listOf(minimizingTaskChange)
+            ),
+            startTransaction = mock(),
+            finishTransaction = mock(),
+            finishCallback = {}
+        )
+
+        assertTrue("Should delegate animation to back navigation transition handler", started)
+        verify(desktopBackNavigationTransitionHandler)
+            .startAnimation(
+                eq(transition),
+                argThat { info -> info.changes.contains(minimizingTaskChange) },
+                any(), any(), any())
+    }
+
+    @Test
+    @EnableFlags(Flags.FLAG_ENABLE_DESKTOP_WINDOWING_BACK_NAVIGATION)
+    fun startAnimation_withMinimizingLastDesktopTask_dispatchesTransition() {
+        val minimizingTask = createTask(WINDOWING_MODE_FREEFORM)
+        val transition = Binder()
+        whenever(desktopRepository.getExpandedTaskCount(any())).thenReturn(2)
+        whenever(
+            desktopBackNavigationTransitionHandler.startAnimation(any(), any(), any(), any(), any())
+        )
+            .thenReturn(true)
+        mixedHandler.addPendingMixedTransition(
+            PendingMixedTransition.Minimize(
+                transition = transition,
+                minimizingTask = minimizingTask.taskId,
+                isLastTask = true,
+            )
+        )
+
+        val minimizingTaskChange = createChange(minimizingTask)
+        mixedHandler.startAnimation(
+            transition = transition,
+            info =
+            createTransitionInfo(
+                TRANSIT_TO_BACK,
+                listOf(minimizingTaskChange)
+            ),
+            startTransaction = mock(),
+            finishTransaction = mock(),
+            finishCallback = {}
+        )
+
+        verify(transitions)
+            .dispatchTransition(
+                eq(transition),
+                argThat { info -> info.changes.contains(minimizingTaskChange) },
+                any(),
+                any(),
+                any(),
+                eq(mixedHandler)
+            )
+    }
+
     private fun createTransitionInfo(
         type: Int = WindowManager.TRANSIT_CLOSE,
         changeMode: Int = WindowManager.TRANSIT_CLOSE,
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopTasksTransitionObserverTest.kt b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopTasksTransitionObserverTest.kt
index 737439c..7f1c1db 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopTasksTransitionObserverTest.kt
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopTasksTransitionObserverTest.kt
@@ -76,6 +76,7 @@
     private val context = mock<Context>()
     private val shellTaskOrganizer = mock<ShellTaskOrganizer>()
     private val taskRepository = mock<DesktopRepository>()
+    private val mixedHandler = mock<DesktopMixedTransitionHandler>()
 
     private lateinit var transitionObserver: DesktopTasksTransitionObserver
     private lateinit var shellInit: ShellInit
@@ -87,7 +88,7 @@
 
         transitionObserver =
             DesktopTasksTransitionObserver(
-                context, taskRepository, transitions, shellTaskOrganizer, shellInit
+                context, taskRepository, transitions, shellTaskOrganizer, mixedHandler, shellInit
             )
     }
 
@@ -106,6 +107,7 @@
         )
 
         verify(taskRepository).minimizeTask(task.displayId, task.taskId)
+        verify(mixedHandler).addPendingMixedTransition(any())
     }
 
     @Test
diff --git a/packages/SystemUI/aconfig/systemui.aconfig b/packages/SystemUI/aconfig/systemui.aconfig
index bd47d18..3df9603 100644
--- a/packages/SystemUI/aconfig/systemui.aconfig
+++ b/packages/SystemUI/aconfig/systemui.aconfig
@@ -1360,16 +1360,6 @@
 }
 
 flag {
-  name: "notification_pulsing_fix"
-  namespace: "systemui"
-  description: "Allow showing new pulsing notifications when the device is already pulsing."
-  bug: "335560575"
-  metadata {
-    purpose: PURPOSE_BUGFIX
-  }
-}
-
-flag {
     name: "media_lockscreen_launch_animation"
     namespace : "systemui"
     description : "Enable the origin launch animation for UMO when opening on top of lockscreen."
@@ -1784,3 +1774,13 @@
       purpose: PURPOSE_BUGFIX
     }
 }
+
+flag {
+    name: "keyguard_transition_force_finish_on_screen_off"
+    namespace: "systemui"
+    description: "Forces KTF transitions to finish if the screen turns all the way off."
+    bug: "331636736"
+    metadata {
+      purpose: PURPOSE_BUGFIX
+    }
+}
\ No newline at end of file
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/data/repository/KeyguardTransitionRepositoryTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/data/repository/KeyguardTransitionRepositoryTest.kt
index bfe89de..3d5498b 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/data/repository/KeyguardTransitionRepositoryTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/data/repository/KeyguardTransitionRepositoryTest.kt
@@ -16,11 +16,14 @@
 
 package com.android.systemui.keyguard.data.repository
 
+import android.animation.Animator
 import android.animation.ValueAnimator
+import android.platform.test.annotations.EnableFlags
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.FlakyTest
 import androidx.test.filters.SmallTest
 import com.android.app.animation.Interpolators
+import com.android.systemui.Flags
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.coroutines.collectValues
 import com.android.systemui.keyguard.shared.model.KeyguardState
@@ -41,6 +44,8 @@
 import java.math.BigDecimal
 import java.math.RoundingMode
 import java.util.UUID
+import kotlin.test.assertEquals
+import kotlin.test.assertTrue
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.delay
 import kotlinx.coroutines.flow.dropWhile
@@ -53,6 +58,7 @@
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
+import org.mockito.kotlin.mock
 
 @SmallTest
 @RunWith(AndroidJUnit4::class)
@@ -65,6 +71,8 @@
     private lateinit var underTest: KeyguardTransitionRepository
     private lateinit var runner: KeyguardTransitionRunner
 
+    private val animatorListener = mock<Animator.AnimatorListener>()
+
     @Before
     fun setUp() {
         underTest = KeyguardTransitionRepositoryImpl(Dispatchers.Main)
@@ -80,7 +88,7 @@
             runner.startTransition(
                 this,
                 TransitionInfo(OWNER_NAME, AOD, LOCKSCREEN, getAnimator()),
-                maxFrames = 100
+                maxFrames = 100,
             )
 
             assertSteps(steps, listWithStep(BigDecimal(.1)), AOD, LOCKSCREEN)
@@ -107,7 +115,7 @@
                     LOCKSCREEN,
                     AOD,
                     getAnimator(),
-                    TransitionModeOnCanceled.LAST_VALUE
+                    TransitionModeOnCanceled.LAST_VALUE,
                 ),
             )
 
@@ -142,7 +150,7 @@
                     LOCKSCREEN,
                     AOD,
                     getAnimator(),
-                    TransitionModeOnCanceled.RESET
+                    TransitionModeOnCanceled.RESET,
                 ),
             )
 
@@ -177,7 +185,7 @@
                     LOCKSCREEN,
                     AOD,
                     getAnimator(),
-                    TransitionModeOnCanceled.REVERSE
+                    TransitionModeOnCanceled.REVERSE,
                 ),
             )
 
@@ -476,6 +484,49 @@
             assertThat(steps.size).isEqualTo(3)
         }
 
+    @Test
+    @EnableFlags(Flags.FLAG_KEYGUARD_TRANSITION_FORCE_FINISH_ON_SCREEN_OFF)
+    fun forceFinishCurrentTransition_noFurtherStepsEmitted() =
+        testScope.runTest {
+            val steps by collectValues(underTest.transitions.dropWhile { step -> step.from == OFF })
+
+            var sentForceFinish = false
+
+            runner.startTransition(
+                this,
+                TransitionInfo(OWNER_NAME, AOD, LOCKSCREEN, getAnimator()),
+                maxFrames = 100,
+                // Force-finish on the second frame.
+                frameCallback = { frameNumber ->
+                    if (!sentForceFinish && frameNumber > 1) {
+                        testScope.launch { underTest.forceFinishCurrentTransition() }
+                        sentForceFinish = true
+                    }
+                },
+            )
+
+            val lastTwoRunningSteps =
+                steps.filter { it.transitionState == TransitionState.RUNNING }.takeLast(2)
+
+            // Make sure we stopped emitting RUNNING steps early, but then emitted a final 1f step.
+            assertTrue(lastTwoRunningSteps[0].value < 0.5f)
+            assertTrue(lastTwoRunningSteps[1].value == 1f)
+
+            assertEquals(steps.last().from, AOD)
+            assertEquals(steps.last().to, LOCKSCREEN)
+            assertEquals(steps.last().transitionState, TransitionState.FINISHED)
+        }
+
+    @Test
+    @EnableFlags(Flags.FLAG_KEYGUARD_TRANSITION_FORCE_FINISH_ON_SCREEN_OFF)
+    fun forceFinishCurrentTransition_noTransitionStarted_noStepsEmitted() =
+        testScope.runTest {
+            val steps by collectValues(underTest.transitions.dropWhile { step -> step.from == OFF })
+
+            underTest.forceFinishCurrentTransition()
+            assertEquals(0, steps.size)
+        }
+
     private fun listWithStep(
         step: BigDecimal,
         start: BigDecimal = BigDecimal.ZERO,
@@ -505,7 +556,7 @@
                     to,
                     fractions[0].toFloat(),
                     TransitionState.STARTED,
-                    OWNER_NAME
+                    OWNER_NAME,
                 )
             )
         fractions.forEachIndexed { index, fraction ->
@@ -519,7 +570,7 @@
                         to,
                         fraction.toFloat(),
                         TransitionState.RUNNING,
-                        OWNER_NAME
+                        OWNER_NAME,
                     )
                 )
         }
@@ -538,6 +589,7 @@
         return ValueAnimator().apply {
             setInterpolator(Interpolators.LINEAR)
             setDuration(10)
+            addListener(animatorListener)
         }
     }
 
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/util/KeyguardTransitionRunner.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/util/KeyguardTransitionRunner.kt
index 1abb441..5798e07 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/util/KeyguardTransitionRunner.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/util/KeyguardTransitionRunner.kt
@@ -21,6 +21,7 @@
 import android.view.Choreographer.FrameCallback
 import com.android.systemui.keyguard.data.repository.KeyguardTransitionRepository
 import com.android.systemui.keyguard.shared.model.TransitionInfo
+import java.util.function.Consumer
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.Job
@@ -35,9 +36,8 @@
  * Gives direct control over ValueAnimator, in order to make transition tests deterministic. See
  * [AnimationHandler]. Animators are required to be run on the main thread, so dispatch accordingly.
  */
-class KeyguardTransitionRunner(
-    val repository: KeyguardTransitionRepository,
-) : AnimationFrameCallbackProvider {
+class KeyguardTransitionRunner(val repository: KeyguardTransitionRepository) :
+    AnimationFrameCallbackProvider {
 
     private var frameCount = 1L
     private var frames = MutableStateFlow(Pair<Long, FrameCallback?>(0L, null))
@@ -48,7 +48,12 @@
      * For transitions being directed by an animator. Will control the number of frames being
      * generated so the values are deterministic.
      */
-    suspend fun startTransition(scope: CoroutineScope, info: TransitionInfo, maxFrames: Int = 100) {
+    suspend fun startTransition(
+        scope: CoroutineScope,
+        info: TransitionInfo,
+        maxFrames: Int = 100,
+        frameCallback: Consumer<Long>? = null,
+    ) {
         // AnimationHandler uses ThreadLocal storage, and ValueAnimators MUST start from main
         // thread
         withContext(Dispatchers.Main) {
@@ -62,7 +67,12 @@
 
                     isTerminated = frameNumber >= maxFrames
                     if (!isTerminated) {
-                        withContext(Dispatchers.Main) { callback?.doFrame(frameNumber) }
+                        try {
+                            withContext(Dispatchers.Main) { callback?.doFrame(frameNumber) }
+                            frameCallback?.accept(frameNumber)
+                        } catch (e: IllegalStateException) {
+                            e.printStackTrace()
+                        }
                     }
                 }
             }
@@ -90,9 +100,13 @@
     override fun postFrameCallback(cb: FrameCallback) {
         frames.value = Pair(frameCount++, cb)
     }
+
     override fun postCommitCallback(runnable: Runnable) {}
+
     override fun getFrameTime() = frameCount
+
     override fun getFrameDelay() = 1L
+
     override fun setFrameDelay(delay: Long) {}
 
     companion object {
diff --git a/packages/SystemUI/src/com/android/systemui/doze/DozeTriggers.java b/packages/SystemUI/src/com/android/systemui/doze/DozeTriggers.java
index dd08d32..7a95a41 100644
--- a/packages/SystemUI/src/com/android/systemui/doze/DozeTriggers.java
+++ b/packages/SystemUI/src/com/android/systemui/doze/DozeTriggers.java
@@ -40,7 +40,6 @@
 import com.android.internal.logging.InstanceId;
 import com.android.internal.logging.UiEvent;
 import com.android.internal.logging.UiEventLogger;
-import com.android.systemui.Flags;
 import com.android.systemui.biometrics.AuthController;
 import com.android.systemui.broadcast.BroadcastDispatcher;
 import com.android.systemui.dock.DockManager;
@@ -566,8 +565,7 @@
         }
 
         // When already in pulsing, we can show the new Notification without requesting a new pulse.
-        if (Flags.notificationPulsingFix()
-                && dozeState == State.DOZE_PULSING && reason == DozeLog.PULSE_REASON_NOTIFICATION) {
+        if (dozeState == State.DOZE_PULSING && reason == DozeLog.PULSE_REASON_NOTIFICATION) {
             return;
         }
 
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/data/repository/KeyguardTransitionRepository.kt b/packages/SystemUI/src/com/android/systemui/keyguard/data/repository/KeyguardTransitionRepository.kt
index 3a5614f..eaf8fa9 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/data/repository/KeyguardTransitionRepository.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/data/repository/KeyguardTransitionRepository.kt
@@ -114,6 +114,18 @@
         @FloatRange(from = 0.0, to = 1.0) value: Float,
         state: TransitionState,
     )
+
+    /**
+     * Forces the current transition to emit FINISHED, foregoing any additional RUNNING steps that
+     * otherwise would have been emitted.
+     *
+     * When the screen is off, upcoming performance changes cause all Animators to cease emitting
+     * frames, which means the Animator passed to [startTransition] will never finish if it was
+     * running when the screen turned off. Also, there's simply no reason to emit RUNNING steps when
+     * the screen isn't even on. As long as we emit FINISHED, everything should end up in the
+     * correct state.
+     */
+    suspend fun forceFinishCurrentTransition()
 }
 
 @SysUISingleton
@@ -134,6 +146,7 @@
     override val transitions = _transitions.asSharedFlow().distinctUntilChanged()
     private var lastStep: TransitionStep = TransitionStep()
     private var lastAnimator: ValueAnimator? = null
+    private var animatorListener: AnimatorListenerAdapter? = null
 
     private val withContextMutex = Mutex()
     private val _currentTransitionInfo: MutableStateFlow<TransitionInfo> =
@@ -233,7 +246,7 @@
                     )
                 }
 
-                val adapter =
+                animatorListener =
                     object : AnimatorListenerAdapter() {
                         override fun onAnimationStart(animation: Animator) {
                             emitTransition(
@@ -254,9 +267,10 @@
                             animator.removeListener(this)
                             animator.removeUpdateListener(updateListener)
                             lastAnimator = null
+                            animatorListener = null
                         }
                     }
-                animator.addListener(adapter)
+                animator.addListener(animatorListener)
                 animator.addUpdateListener(updateListener)
                 animator.start()
                 return@withContext null
@@ -290,6 +304,33 @@
         }
     }
 
+    override suspend fun forceFinishCurrentTransition() {
+        withContextMutex.lock()
+
+        if (lastAnimator?.isRunning != true) {
+            return
+        }
+
+        return withContext("$TAG#forceFinishCurrentTransition", mainDispatcher) {
+            withContextMutex.unlock()
+
+            Log.d(TAG, "forceFinishCurrentTransition() - emitting FINISHED early.")
+
+            lastAnimator?.apply {
+                // Cancel the animator, but remove listeners first so we don't emit CANCELED.
+                removeAllListeners()
+                cancel()
+
+                // Emit a final 1f RUNNING step to ensure that any transitions not listening for a
+                // FINISHED step end up in the right end state.
+                emitTransition(TransitionStep(currentTransitionInfo, 1f, TransitionState.RUNNING))
+
+                // Ask the listener to emit FINISHED and clean up its state.
+                animatorListener?.onAnimationEnd(this)
+            }
+        }
+    }
+
     private suspend fun updateTransitionInternal(
         transitionId: UUID,
         @FloatRange(from = 0.0, to = 1.0) value: Float,
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardTransitionInteractor.kt b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardTransitionInteractor.kt
index b815f19..7cd2744 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardTransitionInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardTransitionInteractor.kt
@@ -19,8 +19,10 @@
 
 import android.annotation.SuppressLint
 import android.util.Log
+import com.android.app.tracing.coroutines.launchTraced as launch
 import com.android.compose.animation.scene.ObservableTransitionState
 import com.android.compose.animation.scene.SceneKey
+import com.android.systemui.Flags.keyguardTransitionForceFinishOnScreenOff
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.dagger.qualifiers.Application
 import com.android.systemui.keyguard.data.repository.KeyguardTransitionRepository
@@ -30,6 +32,8 @@
 import com.android.systemui.keyguard.shared.model.KeyguardState.UNDEFINED
 import com.android.systemui.keyguard.shared.model.TransitionState
 import com.android.systemui.keyguard.shared.model.TransitionStep
+import com.android.systemui.power.domain.interactor.PowerInteractor
+import com.android.systemui.power.shared.model.ScreenPowerState
 import com.android.systemui.scene.domain.interactor.SceneInteractor
 import com.android.systemui.scene.shared.flag.SceneContainerFlag
 import com.android.systemui.scene.shared.model.Scenes
@@ -59,7 +63,6 @@
 import kotlinx.coroutines.flow.onStart
 import kotlinx.coroutines.flow.shareIn
 import kotlinx.coroutines.flow.stateIn
-import com.android.app.tracing.coroutines.launchTraced as launch
 
 /** Encapsulates business-logic related to the keyguard transitions. */
 @OptIn(ExperimentalCoroutinesApi::class)
@@ -70,6 +73,7 @@
     @Application val scope: CoroutineScope,
     private val repository: KeyguardTransitionRepository,
     private val sceneInteractor: SceneInteractor,
+    private val powerInteractor: PowerInteractor,
 ) {
     private val transitionMap = mutableMapOf<Edge.StateToState, MutableSharedFlow<TransitionStep>>()
 
@@ -188,6 +192,18 @@
                     }
                 }
         }
+
+        if (keyguardTransitionForceFinishOnScreenOff()) {
+            /**
+             * If the screen is turning off, finish the current transition immediately. Further
+             * frames won't be visible anyway.
+             */
+            scope.launch {
+                powerInteractor.screenPowerState
+                    .filter { it == ScreenPowerState.SCREEN_TURNING_OFF }
+                    .collect { repository.forceFinishCurrentTransition() }
+            }
+        }
     }
 
     fun transition(edge: Edge, edgeWithoutSceneContainer: Edge? = null): Flow<TransitionStep> {
diff --git a/packages/SystemUI/tests/src/com/android/systemui/doze/DozeTriggersTest.java b/packages/SystemUI/tests/src/com/android/systemui/doze/DozeTriggersTest.java
index 96f4a60..b4c6952 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/doze/DozeTriggersTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/doze/DozeTriggersTest.java
@@ -46,7 +46,6 @@
 
 import com.android.internal.logging.InstanceId;
 import com.android.internal.logging.UiEventLogger;
-import com.android.systemui.Flags;
 import com.android.systemui.SysuiTestCase;
 import com.android.systemui.biometrics.AuthController;
 import com.android.systemui.broadcast.BroadcastDispatcher;
@@ -222,7 +221,6 @@
     }
 
     @Test
-    @EnableFlags(Flags.FLAG_NOTIFICATION_PULSING_FIX)
     public void testOnNotification_alreadyPulsing_notificationNotSuppressed() {
         // GIVEN device is pulsing
         Runnable pulseSuppressListener = mock(Runnable.class);
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/data/repository/FakeKeyguardTransitionRepository.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/data/repository/FakeKeyguardTransitionRepository.kt
index 19e077c..8209ee1 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/data/repository/FakeKeyguardTransitionRepository.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/data/repository/FakeKeyguardTransitionRepository.kt
@@ -87,7 +87,7 @@
     ) : this(
         initInLockscreen = true,
         initiallySendTransitionStepsOnStartTransition = true,
-        testScope
+        testScope,
     )
 
     private val _currentTransitionInfo: MutableStateFlow<TransitionInfo> =
@@ -191,12 +191,12 @@
         if (lastStep != null && lastStep.transitionState != TransitionState.FINISHED) {
             sendTransitionStep(
                 step =
-                TransitionStep(
-                    transitionState = TransitionState.CANCELED,
-                    from = lastStep.from,
-                    to = lastStep.to,
-                    value = 0f,
-                )
+                    TransitionStep(
+                        transitionState = TransitionState.CANCELED,
+                        from = lastStep.from,
+                        to = lastStep.to,
+                        value = 0f,
+                    )
             )
             testScheduler.runCurrent()
         }
@@ -390,6 +390,18 @@
         @FloatRange(from = 0.0, to = 1.0) value: Float,
         state: TransitionState,
     ) = Unit
+
+    override suspend fun forceFinishCurrentTransition() {
+        _transitions.tryEmit(
+            TransitionStep(
+                _currentTransitionInfo.value.from,
+                _currentTransitionInfo.value.to,
+                1f,
+                TransitionState.FINISHED,
+                ownerName = _currentTransitionInfo.value.ownerName,
+            )
+        )
+    }
 }
 
 @Module
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/domain/interactor/KeyguardTransitionInteractorKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/domain/interactor/KeyguardTransitionInteractorKosmos.kt
index aa94c36..b9a831f 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/domain/interactor/KeyguardTransitionInteractorKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/domain/interactor/KeyguardTransitionInteractorKosmos.kt
@@ -19,6 +19,7 @@
 import com.android.systemui.keyguard.data.repository.keyguardTransitionRepository
 import com.android.systemui.kosmos.Kosmos
 import com.android.systemui.kosmos.applicationCoroutineScope
+import com.android.systemui.power.domain.interactor.powerInteractor
 import com.android.systemui.scene.domain.interactor.sceneInteractor
 
 val Kosmos.keyguardTransitionInteractor: KeyguardTransitionInteractor by
@@ -26,6 +27,7 @@
         KeyguardTransitionInteractor(
             scope = applicationCoroutineScope,
             repository = keyguardTransitionRepository,
-            sceneInteractor = sceneInteractor
+            sceneInteractor = sceneInteractor,
+            powerInteractor = powerInteractor,
         )
     }