Keyguard Transitions - Update check for canceling

Instead of relying on an emitted transition, check the
animator directly to see if a cancel will be issued, which
is the actual source of truth for automated transitions.

Also, fix the transition runner. The handler needs to be
updated once per test run in setup, not for each call to
startTransition, which was causing race conditions.

Also, add a bit more logging for the test where this issue
was detected.

Test: atest SystemUITests:KeyguardTransitionRepositoryTest
Bug: 383318506
Flag: EXEMPT bugfix
Change-Id: I265cb41b1c4494827c91cb98287ab2b7e2102b4a

4 files changed, 74 insertions, 49 deletions

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 5798e0776c4f..338b06824e85 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
@@ -24,8 +24,9 @@ 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
 import kotlinx.coroutines.delay
+import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.MutableSharedFlow
 import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.collect
 import kotlinx.coroutines.launch
@@ -36,12 +37,10 @@ import org.junit.Assert.fail
  * 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 {
-
-    private var frameCount = 1L
-    private var frames = MutableStateFlow(Pair<Long, FrameCallback?>(0L, null))
-    private var job: Job? = null
+class KeyguardTransitionRunner(
+    val frames: Flow<Long>,
+    val repository: KeyguardTransitionRepository,
+) {
     @Volatile private var isTerminated = false
 
     /**
@@ -54,21 +53,12 @@ class KeyguardTransitionRunner(val repository: KeyguardTransitionRepository) :
         maxFrames: Int = 100,
         frameCallback: Consumer<Long>? = null,
     ) {
-        // AnimationHandler uses ThreadLocal storage, and ValueAnimators MUST start from main
-        // thread
-        withContext(Dispatchers.Main) {
-            info.animator!!.getAnimationHandler().setProvider(this@KeyguardTransitionRunner)
-        }
-
-        job =
+        val job =
             scope.launch {
-                frames.collect {
-                    val (frameNumber, callback) = it
-
+                frames.collect { frameNumber ->
                     isTerminated = frameNumber >= maxFrames
                     if (!isTerminated) {
                         try {
-                            withContext(Dispatchers.Main) { callback?.doFrame(frameNumber) }
                             frameCallback?.accept(frameNumber)
                         } catch (e: IllegalStateException) {
                             e.printStackTrace()
@@ -78,27 +68,46 @@ class KeyguardTransitionRunner(val repository: KeyguardTransitionRepository) :
             }
         withContext(Dispatchers.Main) { repository.startTransition(info) }
 
-        waitUntilComplete(info.animator!!)
+        waitUntilComplete(info, info.animator!!)
+        job.cancel()
     }
 
-    private suspend fun waitUntilComplete(animator: ValueAnimator) {
+    private suspend fun waitUntilComplete(info: TransitionInfo, animator: ValueAnimator) {
         withContext(Dispatchers.Main) {
             val startTime = System.currentTimeMillis()
             while (!isTerminated && animator.isRunning()) {
                 delay(1)
                 if (System.currentTimeMillis() - startTime > MAX_TEST_DURATION) {
-                    fail("Failed test due to excessive runtime of: $MAX_TEST_DURATION")
+                    fail("Failed due to excessive runtime of: $MAX_TEST_DURATION, info: $info")
                 }
             }
-
-            animator.getAnimationHandler().setProvider(null)
         }
+    }
 
-        job?.cancel()
+    companion object {
+        private const val MAX_TEST_DURATION = 300L
+    }
+}
+
+class FrameCallbackProvider(val scope: CoroutineScope) : AnimationFrameCallbackProvider {
+    private val callback = MutableSharedFlow<FrameCallback?>(replay = 2)
+    private var frameCount = 0L
+    val frames = MutableStateFlow(frameCount)
+
+    init {
+        scope.launch {
+            callback.collect {
+                withContext(Dispatchers.Main) {
+                    delay(1)
+                    it?.doFrame(frameCount)
+                }
+            }
+        }
     }
 
     override fun postFrameCallback(cb: FrameCallback) {
-        frames.value = Pair(frameCount++, cb)
+        frames.value = ++frameCount
+        callback.tryEmit(cb)
     }
 
     override fun postCommitCallback(runnable: Runnable) {}
@@ -108,8 +117,4 @@ class KeyguardTransitionRunner(val repository: KeyguardTransitionRepository) :
     override fun getFrameDelay() = 1L
 
     override fun setFrameDelay(delay: Long) {}
-
-    companion object {
-        private const val MAX_TEST_DURATION = 200L
-    }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/KeyguardViewMediator.java b/packages/SystemUI/src/com/android/systemui/keyguard/KeyguardViewMediator.java
index 63ac5094c400..4bc6742b5040 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/KeyguardViewMediator.java
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/KeyguardViewMediator.java
@@ -813,7 +813,7 @@ public class KeyguardViewMediator implements CoreStartable, Dumpable,
             if (targetUserId != mSelectedUserInteractor.getSelectedUserId()) {
                 return;
             }
-            if (DEBUG) Log.d(TAG, "keyguardDone");
+            Log.d(TAG, "keyguardDone");
             tryKeyguardDone();
         }
 
@@ -832,7 +832,7 @@ public class KeyguardViewMediator implements CoreStartable, Dumpable,
         @Override
         public void keyguardDonePending(int targetUserId) {
             Trace.beginSection("KeyguardViewMediator.mViewMediatorCallback#keyguardDonePending");
-            if (DEBUG) Log.d(TAG, "keyguardDonePending");
+            Log.d(TAG, "keyguardDonePending");
             if (targetUserId != mSelectedUserInteractor.getSelectedUserId()) {
                 Trace.endSection();
                 return;
@@ -2735,10 +2735,8 @@ public class KeyguardViewMediator implements CoreStartable, Dumpable,
     }
 
     private void tryKeyguardDone() {
-        if (DEBUG) {
-            Log.d(TAG, "tryKeyguardDone: pending - " + mKeyguardDonePending + ", animRan - "
-                    + mHideAnimationRun + " animRunning - " + mHideAnimationRunning);
-        }
+        Log.d(TAG, "tryKeyguardDone: pending - " + mKeyguardDonePending + ", animRan - "
+                + mHideAnimationRun + " animRunning - " + mHideAnimationRunning);
         if (!mKeyguardDonePending && mHideAnimationRun && !mHideAnimationRunning) {
             handleKeyguardDone();
         } else if (mSurfaceBehindRemoteAnimationRunning) {
@@ -3040,7 +3038,7 @@ public class KeyguardViewMediator implements CoreStartable, Dumpable,
     }
 
     private final Runnable mHideAnimationFinishedRunnable = () -> {
-        Log.e(TAG, "mHideAnimationFinishedRunnable#run");
+        Log.d(TAG, "mHideAnimationFinishedRunnable#run");
         mHideAnimationRunning = false;
         tryKeyguardDone();
     };
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 354fc3d82342..24f2493c626d 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
@@ -212,8 +212,11 @@ constructor(@Main val mainDispatcher: CoroutineDispatcher) : KeyguardTransitionR
                 Log.i(TAG, "Duplicate call to start the transition, rejecting: $info")
                 return@withContext null
             }
+            val isAnimatorRunning = lastAnimator?.isRunning() ?: false
+            val isManualTransitionRunning =
+                updateTransitionId != null && lastStep.transitionState != TransitionState.FINISHED
             val startingValue =
-                if (lastStep.transitionState != TransitionState.FINISHED) {
+                if (isAnimatorRunning || isManualTransitionRunning) {
                     Log.i(TAG, "Transition still active: $lastStep, canceling")
                     when (info.modeOnCanceled) {
                         TransitionModeOnCanceled.LAST_VALUE -> lastStep.value
diff --git a/packages/SystemUI/tests/src/com/android/systemui/keyguard/data/repository/KeyguardTransitionRepositoryTest.kt b/packages/SystemUI/tests/src/com/android/systemui/keyguard/data/repository/KeyguardTransitionRepositoryTest.kt
index 7a3089f33276..77c40a1e8eef 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/keyguard/data/repository/KeyguardTransitionRepositoryTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/keyguard/data/repository/KeyguardTransitionRepositoryTest.kt
@@ -16,6 +16,7 @@
 
 package com.android.systemui.keyguard.data.repository
 
+import android.animation.AnimationHandler
 import android.animation.Animator
 import android.animation.ValueAnimator
 import android.platform.test.annotations.EnableFlags
@@ -36,6 +37,7 @@ import com.android.systemui.keyguard.shared.model.TransitionInfo
 import com.android.systemui.keyguard.shared.model.TransitionModeOnCanceled
 import com.android.systemui.keyguard.shared.model.TransitionState
 import com.android.systemui.keyguard.shared.model.TransitionStep
+import com.android.systemui.keyguard.util.FrameCallbackProvider
 import com.android.systemui.keyguard.util.KeyguardTransitionRunner
 import com.android.systemui.kosmos.testDispatcher
 import com.android.systemui.kosmos.testScope
@@ -52,9 +54,12 @@ import kotlinx.coroutines.flow.dropWhile
 import kotlinx.coroutines.flow.launchIn
 import kotlinx.coroutines.flow.onEach
 import kotlinx.coroutines.launch
+import kotlinx.coroutines.runBlocking
 import kotlinx.coroutines.test.advanceTimeBy
 import kotlinx.coroutines.test.runCurrent
 import kotlinx.coroutines.test.runTest
+import kotlinx.coroutines.withContext
+import org.junit.After
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
@@ -70,13 +75,29 @@ class KeyguardTransitionRepositoryTest : SysuiTestCase() {
 
     private lateinit var underTest: KeyguardTransitionRepository
     private lateinit var runner: KeyguardTransitionRunner
+    private lateinit var callbackProvider: FrameCallbackProvider
 
     private val animatorListener = mock<Animator.AnimatorListener>()
 
     @Before
     fun setUp() {
         underTest = KeyguardTransitionRepositoryImpl(Dispatchers.Main)
-        runner = KeyguardTransitionRunner(underTest)
+        runBlocking {
+            callbackProvider = FrameCallbackProvider(testScope.backgroundScope)
+            withContext(Dispatchers.Main) {
+                // AnimationHandler uses ThreadLocal storage, and ValueAnimators MUST start from
+                // main thread
+                AnimationHandler.getInstance().setProvider(callbackProvider)
+            }
+            runner = KeyguardTransitionRunner(callbackProvider.frames, underTest)
+        }
+    }
+
+    @After
+    fun tearDown() {
+        runBlocking {
+            withContext(Dispatchers.Main) { AnimationHandler.getInstance().setProvider(null) }
+        }
     }
 
     @Test
@@ -84,13 +105,11 @@ class KeyguardTransitionRepositoryTest : SysuiTestCase() {
         testScope.runTest {
             val steps = mutableListOf<TransitionStep>()
             val job = underTest.transition(AOD, LOCKSCREEN).onEach { steps.add(it) }.launchIn(this)
-
             runner.startTransition(
                 this,
                 TransitionInfo(OWNER_NAME, AOD, LOCKSCREEN, getAnimator()),
                 maxFrames = 100,
             )
-
             assertSteps(steps, listWithStep(BigDecimal(.1)), AOD, LOCKSCREEN)
             job.cancel()
         }
@@ -119,12 +138,12 @@ class KeyguardTransitionRepositoryTest : SysuiTestCase() {
                 ),
             )
 
-            val firstTransitionSteps = listWithStep(step = BigDecimal(.1), stop = BigDecimal(.1))
-            assertSteps(steps.subList(0, 4), firstTransitionSteps, AOD, LOCKSCREEN)
+            val firstTransitionSteps = listWithStep(step = BigDecimal(.1), stop = BigDecimal(.2))
+            assertSteps(steps.subList(0, 5), firstTransitionSteps, AOD, LOCKSCREEN)
 
-            // Second transition starts from .1 (LAST_VALUE)
-            val secondTransitionSteps = listWithStep(step = BigDecimal(.1), start = BigDecimal(.1))
-            assertSteps(steps.subList(4, steps.size), secondTransitionSteps, LOCKSCREEN, AOD)
+            // Second transition starts from .2 (LAST_VALUE)
+            val secondTransitionSteps = listWithStep(step = BigDecimal(.1), start = BigDecimal(.2))
+            assertSteps(steps.subList(5, steps.size), secondTransitionSteps, LOCKSCREEN, AOD)
 
             job.cancel()
             job2.cancel()
@@ -154,12 +173,12 @@ class KeyguardTransitionRepositoryTest : SysuiTestCase() {
                 ),
             )
 
-            val firstTransitionSteps = listWithStep(step = BigDecimal(.1), stop = BigDecimal(.1))
-            assertSteps(steps.subList(0, 4), firstTransitionSteps, AOD, LOCKSCREEN)
+            val firstTransitionSteps = listWithStep(step = BigDecimal(.1), stop = BigDecimal(.2))
+            assertSteps(steps.subList(0, 5), firstTransitionSteps, AOD, LOCKSCREEN)
 
             // Second transition starts from 0 (RESET)
             val secondTransitionSteps = listWithStep(start = BigDecimal(0), step = BigDecimal(.1))
-            assertSteps(steps.subList(4, steps.size), secondTransitionSteps, LOCKSCREEN, AOD)
+            assertSteps(steps.subList(5, steps.size), secondTransitionSteps, LOCKSCREEN, AOD)
 
             job.cancel()
             job2.cancel()
@@ -173,7 +192,7 @@ class KeyguardTransitionRepositoryTest : SysuiTestCase() {
             runner.startTransition(
                 this,
                 TransitionInfo(OWNER_NAME, AOD, LOCKSCREEN, getAnimator()),
-                maxFrames = 3,
+                maxFrames = 2,
             )
 
             // Now start 2nd transition, which will interrupt the first