Merge "Transitions - Listen for bouncer show & hide" into tm-qpr-dev am: be5367c98a

Original change: https://googleplex-android-review.googlesource.com/c/platform/frameworks/base/+/20639790 Change-Id: If23b9be9b8a4c0f256272e12eb391d307c2077bc Signed-off-by: Automerger Merge Worker <android-build-automerger-merge-worker@system.gserviceaccount.com>
2022-12-16 13:30:38 +00:00
parent 97d020c60a be5367c98a
commit c54aaa1677
8 changed files with 358 additions and 144 deletions
--- a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/LockscreenBouncerTransitionInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/LockscreenBouncerTransitionInteractor.kt
@@ -22,7 +22,7 @@ import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.dagger.qualifiers.Application
 import com.android.systemui.keyguard.data.repository.KeyguardTransitionRepository
 import com.android.systemui.keyguard.shared.model.KeyguardState
-import com.android.systemui.keyguard.shared.model.StatusBarState.SHADE_LOCKED
+import com.android.systemui.keyguard.shared.model.StatusBarState.KEYGUARD
 import com.android.systemui.keyguard.shared.model.TransitionInfo
 import com.android.systemui.keyguard.shared.model.TransitionState
 import com.android.systemui.keyguard.shared.model.WakefulnessState
@@ -50,27 +50,22 @@ constructor(

    override fun start() {
        listenForDraggingUpToBouncer()
-        listenForBouncerHiding()
+        listenForBouncer()
    }

-    private fun listenForBouncerHiding() {
+    private fun listenForBouncer() {
        scope.launch {
            keyguardInteractor.isBouncerShowing
                .sample(
                    combine(
                        keyguardInteractor.wakefulnessModel,
                        keyguardTransitionInteractor.startedKeyguardTransitionStep,
-                    ) { wakefulnessModel, transitionStep ->
-                        Pair(wakefulnessModel, transitionStep)
-                    }
-                ) { bouncerShowing, wakefulnessAndTransition ->
-                    Triple(
-                        bouncerShowing,
-                        wakefulnessAndTransition.first,
-                        wakefulnessAndTransition.second
-                    )
-                }
-                .collect { (isBouncerShowing, wakefulnessState, lastStartedTransitionStep) ->
+                        ::Pair
+                    ),
+                    ::toTriple
+                )
+                .collect { triple ->
+                    val (isBouncerShowing, wakefulnessState, lastStartedTransitionStep) = triple
                    if (
                        !isBouncerShowing && lastStartedTransitionStep.to == KeyguardState.BOUNCER
                    ) {
@@ -91,7 +86,19 @@ constructor(
                                animator = getAnimator(),
                            )
                        )
+                    } else if (
+                        isBouncerShowing && lastStartedTransitionStep.to == KeyguardState.LOCKSCREEN
+                    ) {
+                        keyguardTransitionRepository.startTransition(
+                            TransitionInfo(
+                                ownerName = name,
+                                from = KeyguardState.LOCKSCREEN,
+                                to = KeyguardState.BOUNCER,
+                                animator = getAnimator(),
+                            )
+                        )
                    }
+                    Unit
                }
        }
    }
@@ -104,24 +111,20 @@ constructor(
                    combine(
                        keyguardTransitionInteractor.finishedKeyguardState,
                        keyguardInteractor.statusBarState,
-                    ) { finishedKeyguardState, statusBarState ->
-                        Pair(finishedKeyguardState, statusBarState)
-                    }
-                ) { shadeModel, keyguardStateAndStatusBarState ->
-                    Triple(
-                        shadeModel,
-                        keyguardStateAndStatusBarState.first,
-                        keyguardStateAndStatusBarState.second
-                    )
-                }
-                .collect { (shadeModel, keyguardState, statusBarState) ->
+                        ::Pair
+                    ),
+                    ::toTriple
+                )
+                .collect { triple ->
+                    val (shadeModel, keyguardState, statusBarState) = triple
+
                    val id = transitionId
                    if (id != null) {
                        // An existing `id` means a transition is started, and calls to
                        // `updateTransition` will control it until FINISHED
                        keyguardTransitionRepository.updateTransition(
                            id,
-                            shadeModel.expansionAmount,
+                            1f - shadeModel.expansionAmount,
                            if (
                                shadeModel.expansionAmount == 0f || shadeModel.expansionAmount == 1f
                            ) {
@@ -137,7 +140,7 @@ constructor(
                        if (
                            keyguardState == KeyguardState.LOCKSCREEN &&
                                shadeModel.isUserDragging &&
-                                statusBarState != SHADE_LOCKED
+                                statusBarState == KEYGUARD
                        ) {
                            transitionId =
                                keyguardTransitionRepository.startTransition(
--- a/packages/SystemUI/src/com/android/systemui/shade/ShadeEventsModule.java
+++ b/packages/SystemUI/src/com/android/systemui/shade/ShadeEventsModule.java
@@ -16,14 +16,17 @@

 package com.android.systemui.shade;

-import com.android.systemui.dagger.SysUISingleton;
+import com.android.systemui.shade.data.repository.ShadeRepository;
+import com.android.systemui.shade.data.repository.ShadeRepositoryImpl;

 import dagger.Binds;
 import dagger.Module;

-/** Provides a {@link ShadeStateEvents} in {@link SysUISingleton} scope. */
+/** Provides Shade-related events and information. */
@Module
 public abstract class ShadeEventsModule {
    @Binds
    abstract ShadeStateEvents bindShadeEvents(ShadeExpansionStateManager impl);
+
+    @Binds abstract ShadeRepository shadeRepository(ShadeRepositoryImpl impl);
 }
--- a/packages/SystemUI/src/com/android/systemui/shade/data/repository/ShadeRepository.kt
+++ b/packages/SystemUI/src/com/android/systemui/shade/data/repository/ShadeRepository.kt
@@ -27,11 +27,17 @@ import kotlinx.coroutines.channels.awaitClose
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.distinctUntilChanged

+interface ShadeRepository {
+    /** ShadeModel information regarding shade expansion events */
+    val shadeModel: Flow<ShadeModel>
+}
+
 /** Business logic for shade interactions */
@SysUISingleton
-class ShadeRepository @Inject constructor(shadeExpansionStateManager: ShadeExpansionStateManager) {
-
-    val shadeModel: Flow<ShadeModel> =
+class ShadeRepositoryImpl
+@Inject
+constructor(shadeExpansionStateManager: ShadeExpansionStateManager) : ShadeRepository {
+    override val shadeModel: Flow<ShadeModel> =
        conflatedCallbackFlow {
                val callback =
                    object : ShadeExpansionListener {
--- 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,13 +16,10 @@

 package com.android.systemui.keyguard.data.repository

-import android.animation.AnimationHandler.AnimationFrameCallbackProvider
 import android.animation.ValueAnimator
 import android.util.Log
 import android.util.Log.TerribleFailure
 import android.util.Log.TerribleFailureHandler
-import android.view.Choreographer.FrameCallback
-import androidx.test.filters.FlakyTest
 import androidx.test.filters.SmallTest
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.animation.Interpolators
@@ -32,22 +29,17 @@ import com.android.systemui.keyguard.shared.model.KeyguardState.LOCKSCREEN
 import com.android.systemui.keyguard.shared.model.TransitionInfo
 import com.android.systemui.keyguard.shared.model.TransitionState
 import com.android.systemui.keyguard.shared.model.TransitionStep
+import com.android.systemui.keyguard.util.KeyguardTransitionRunner
 import com.google.common.truth.Truth.assertThat
 import java.math.BigDecimal
 import java.math.RoundingMode
 import java.util.UUID
-import kotlinx.coroutines.CoroutineScope
-import kotlinx.coroutines.Dispatchers
-import kotlinx.coroutines.Job
-import kotlinx.coroutines.flow.MutableStateFlow
-import kotlinx.coroutines.flow.collect
 import kotlinx.coroutines.flow.launchIn
 import kotlinx.coroutines.flow.onEach
-import kotlinx.coroutines.launch
-import kotlinx.coroutines.runBlocking
-import kotlinx.coroutines.yield
+import kotlinx.coroutines.test.TestScope
+import kotlinx.coroutines.test.runCurrent
+import kotlinx.coroutines.test.runTest
 import org.junit.After
-import org.junit.Assert.fail
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
@@ -60,12 +52,14 @@ class KeyguardTransitionRepositoryTest : SysuiTestCase() {
    private lateinit var underTest: KeyguardTransitionRepository
    private lateinit var oldWtfHandler: TerribleFailureHandler
    private lateinit var wtfHandler: WtfHandler
+    private lateinit var runner: KeyguardTransitionRunner

    @Before
    fun setUp() {
        underTest = KeyguardTransitionRepositoryImpl()
        wtfHandler = WtfHandler()
        oldWtfHandler = Log.setWtfHandler(wtfHandler)
+        runner = KeyguardTransitionRunner(underTest)
    }

    @After
@@ -75,56 +69,37 @@ class KeyguardTransitionRepositoryTest : SysuiTestCase() {

    @Test
    fun `startTransition runs animator to completion`() =
-        runBlocking(IMMEDIATE) {
-            val (animator, provider) = setupAnimator(this)
-
+        TestScope().runTest {
            val steps = mutableListOf<TransitionStep>()
            val job = underTest.transition(AOD, LOCKSCREEN).onEach { steps.add(it) }.launchIn(this)

-            underTest.startTransition(TransitionInfo(OWNER_NAME, AOD, LOCKSCREEN, animator))
-
-            val startTime = System.currentTimeMillis()
-            while (animator.isRunning()) {
-                yield()
-                if (System.currentTimeMillis() - startTime > MAX_TEST_DURATION) {
-                    fail("Failed test due to excessive runtime of: $MAX_TEST_DURATION")
-                }
-            }
+            runner.startTransition(
+                this,
+                TransitionInfo(OWNER_NAME, AOD, LOCKSCREEN, getAnimator()),
+                maxFrames = 100
+            )

            assertSteps(steps, listWithStep(BigDecimal(.1)), AOD, LOCKSCREEN)
-
            job.cancel()
-            provider.stop()
        }

    @Test
-    @FlakyTest(bugId = 260213291)
-    fun `starting second transition will cancel the first transition`() {
-        runBlocking(IMMEDIATE) {
-            val (animator, provider) = setupAnimator(this)
-
+    fun `starting second transition will cancel the first transition`() =
+        TestScope().runTest {
            val steps = mutableListOf<TransitionStep>()
            val job = underTest.transition(AOD, LOCKSCREEN).onEach { steps.add(it) }.launchIn(this)
-
-            underTest.startTransition(TransitionInfo(OWNER_NAME, AOD, LOCKSCREEN, animator))
-            // 3 yields(), alternating with the animator, results in a value 0.1, which can be
-            // canceled and tested against
-            yield()
-            yield()
-            yield()
+            runner.startTransition(
+                this,
+                TransitionInfo(OWNER_NAME, AOD, LOCKSCREEN, getAnimator()),
+                maxFrames = 3,
+            )

            // Now start 2nd transition, which will interrupt the first
            val job2 = underTest.transition(LOCKSCREEN, AOD).onEach { steps.add(it) }.launchIn(this)
-            val (animator2, provider2) = setupAnimator(this)
-            underTest.startTransition(TransitionInfo(OWNER_NAME, LOCKSCREEN, AOD, animator2))
-
-            val startTime = System.currentTimeMillis()
-            while (animator2.isRunning()) {
-                yield()
-                if (System.currentTimeMillis() - startTime > MAX_TEST_DURATION) {
-                    fail("Failed test due to excessive runtime of: $MAX_TEST_DURATION")
-                }
-            }
+            runner.startTransition(
+                this,
+                TransitionInfo(OWNER_NAME, LOCKSCREEN, AOD, getAnimator()),
+            )

            val firstTransitionSteps = listWithStep(step = BigDecimal(.1), stop = BigDecimal(.1))
            assertSteps(steps.subList(0, 4), firstTransitionSteps, AOD, LOCKSCREEN)
@@ -134,31 +109,25 @@ class KeyguardTransitionRepositoryTest : SysuiTestCase() {

            job.cancel()
            job2.cancel()
-            provider.stop()
-            provider2.stop()
        }
-    }

    @Test
    fun `Null animator enables manual control with updateTransition`() =
-        runBlocking(IMMEDIATE) {
+        TestScope().runTest {
            val steps = mutableListOf<TransitionStep>()
            val job = underTest.transition(AOD, LOCKSCREEN).onEach { steps.add(it) }.launchIn(this)

            val uuid =
                underTest.startTransition(
-                    TransitionInfo(
-                        ownerName = OWNER_NAME,
-                        from = AOD,
-                        to = LOCKSCREEN,
-                        animator = null,
-                    )
+                    TransitionInfo(OWNER_NAME, AOD, LOCKSCREEN, animator = null)
                )
+            runCurrent()

            checkNotNull(uuid).let {
                underTest.updateTransition(it, 0.5f, TransitionState.RUNNING)
                underTest.updateTransition(it, 1f, TransitionState.FINISHED)
            }
+            runCurrent()

            assertThat(steps.size).isEqualTo(3)
            assertThat(steps[0])
@@ -256,57 +225,11 @@ class KeyguardTransitionRepositoryTest : SysuiTestCase() {
        assertThat(wtfHandler.failed).isFalse()
    }

-    private fun setupAnimator(
-        scope: CoroutineScope
-    ): Pair<ValueAnimator, TestFrameCallbackProvider> {
-        val animator =
-            ValueAnimator().apply {
-                setInterpolator(Interpolators.LINEAR)
-                setDuration(ANIMATION_DURATION)
-            }
-
-        val provider = TestFrameCallbackProvider(animator, scope)
-        provider.start()
-
-        return Pair(animator, provider)
-    }
-
-    /** Gives direct control over ValueAnimator. See [AnimationHandler] */
-    private class TestFrameCallbackProvider(
-        private val animator: ValueAnimator,
-        private val scope: CoroutineScope,
-    ) : AnimationFrameCallbackProvider {
-
-        private var frameCount = 1L
-        private var frames = MutableStateFlow(Pair<Long, FrameCallback?>(0L, null))
-        private var job: Job? = null
-
-        fun start() {
-            animator.getAnimationHandler().setProvider(this)
-
-            job =
-                scope.launch {
-                    frames.collect {
-                        // Delay is required for AnimationHandler to properly register a callback
-                        yield()
-                        val (frameNumber, callback) = it
-                        callback?.doFrame(frameNumber)
-                    }
-                }
+    private fun getAnimator(): ValueAnimator {
+        return ValueAnimator().apply {
+            setInterpolator(Interpolators.LINEAR)
+            setDuration(10)
        }
-
-        fun stop() {
-            job?.cancel()
-            animator.getAnimationHandler().setProvider(null)
-        }
-
-        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) {}
    }

    private class WtfHandler : TerribleFailureHandler {
@@ -317,9 +240,6 @@ class KeyguardTransitionRepositoryTest : SysuiTestCase() {
    }

    companion object {
-        private const val MAX_TEST_DURATION = 100L
-        private const val ANIMATION_DURATION = 10L
-        private const val OWNER_NAME = "Test"
-        private val IMMEDIATE = Dispatchers.Main.immediate
+        private const val OWNER_NAME = "KeyguardTransitionRunner"
    }
 }
--- a/packages/SystemUI/tests/src/com/android/systemui/keyguard/domain/interactor/KeyguardTransitionScenariosTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/keyguard/domain/interactor/KeyguardTransitionScenariosTest.kt
@@ -0,0 +1,140 @@
+/*
+ * Copyright (C) 2022 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.keyguard.domain.interactor
+
+import android.animation.ValueAnimator
+import androidx.test.filters.SmallTest
+import com.android.systemui.SysuiTestCase
+import com.android.systemui.animation.Interpolators
+import com.android.systemui.keyguard.data.repository.FakeKeyguardRepository
+import com.android.systemui.keyguard.data.repository.KeyguardTransitionRepository
+import com.android.systemui.keyguard.data.repository.KeyguardTransitionRepositoryImpl
+import com.android.systemui.keyguard.shared.model.KeyguardState
+import com.android.systemui.keyguard.shared.model.TransitionInfo
+import com.android.systemui.keyguard.shared.model.WakeSleepReason
+import com.android.systemui.keyguard.shared.model.WakefulnessModel
+import com.android.systemui.keyguard.shared.model.WakefulnessState
+import com.android.systemui.keyguard.util.KeyguardTransitionRunner
+import com.android.systemui.shade.data.repository.FakeShadeRepository
+import com.android.systemui.shade.data.repository.ShadeRepository
+import com.android.systemui.util.mockito.withArgCaptor
+import com.google.common.truth.Truth.assertThat
+import kotlinx.coroutines.cancelChildren
+import kotlinx.coroutines.test.TestScope
+import kotlinx.coroutines.test.runCurrent
+import kotlinx.coroutines.test.runTest
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.junit.runners.JUnit4
+import org.mockito.Mock
+import org.mockito.Mockito.verify
+import org.mockito.MockitoAnnotations
+
+/**
+ * Class for testing user journeys through the interactors. They will all be activated during setup,
+ * to ensure the expected transitions are still triggered.
+ */
+@SmallTest
+@RunWith(JUnit4::class)
+class KeyguardTransitionScenariosTest : SysuiTestCase() {
+    private lateinit var testScope: TestScope
+
+    private lateinit var keyguardRepository: FakeKeyguardRepository
+    private lateinit var shadeRepository: ShadeRepository
+
+    // Used to issue real transition steps for test input
+    private lateinit var runner: KeyguardTransitionRunner
+    private lateinit var transitionRepository: KeyguardTransitionRepository
+
+    // Used to verify transition requests for test output
+    @Mock private lateinit var mockTransitionRepository: KeyguardTransitionRepository
+
+    private lateinit var lockscreenBouncerTransitionInteractor:
+        LockscreenBouncerTransitionInteractor
+
+    @Before
+    fun setUp() {
+        MockitoAnnotations.initMocks(this)
+        testScope = TestScope()
+
+        keyguardRepository = FakeKeyguardRepository()
+        shadeRepository = FakeShadeRepository()
+
+        /* Used to issue full transition steps, to better simulate a real device */
+        transitionRepository = KeyguardTransitionRepositoryImpl()
+        runner = KeyguardTransitionRunner(transitionRepository)
+
+        lockscreenBouncerTransitionInteractor =
+            LockscreenBouncerTransitionInteractor(
+                scope = testScope,
+                keyguardInteractor = KeyguardInteractor(keyguardRepository),
+                shadeRepository = shadeRepository,
+                keyguardTransitionRepository = mockTransitionRepository,
+                keyguardTransitionInteractor = KeyguardTransitionInteractor(transitionRepository),
+            )
+        lockscreenBouncerTransitionInteractor.start()
+    }
+
+    @Test
+    fun `LOCKSCREEN to BOUNCER via bouncer showing call`() =
+        testScope.runTest {
+            // GIVEN a device that has at least woken up
+            keyguardRepository.setWakefulnessModel(startingToWake())
+            runCurrent()
+
+            // GIVEN a transition has run to LOCKSCREEN
+            runner.startTransition(
+                testScope,
+                TransitionInfo(
+                    ownerName = "",
+                    from = KeyguardState.OFF,
+                    to = KeyguardState.LOCKSCREEN,
+                    animator =
+                        ValueAnimator().apply {
+                            duration = 10
+                            interpolator = Interpolators.LINEAR
+                        },
+                )
+            )
+            runCurrent()
+
+            // WHEN the bouncer is set to show
+            keyguardRepository.setBouncerShowing(true)
+            runCurrent()
+
+            val info =
+                withArgCaptor<TransitionInfo> {
+                    verify(mockTransitionRepository).startTransition(capture())
+                }
+            // THEN a transition to BOUNCER should occur
+            assertThat(info.ownerName).isEqualTo("LockscreenBouncerTransitionInteractor")
+            assertThat(info.from).isEqualTo(KeyguardState.LOCKSCREEN)
+            assertThat(info.to).isEqualTo(KeyguardState.BOUNCER)
+            assertThat(info.animator).isNotNull()
+
+            coroutineContext.cancelChildren()
+        }
+
+    private fun startingToWake() =
+        WakefulnessModel(
+            WakefulnessState.STARTING_TO_WAKE,
+            true,
+            WakeSleepReason.OTHER,
+            WakeSleepReason.OTHER
+        )
+}
--- a/packages/SystemUI/tests/src/com/android/systemui/keyguard/util/KeyguardTransitionRunner.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/keyguard/util/KeyguardTransitionRunner.kt
@@ -0,0 +1,101 @@
+/*
+ * Copyright (C) 2022 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.keyguard.util
+
+import android.animation.AnimationHandler.AnimationFrameCallbackProvider
+import android.animation.ValueAnimator
+import android.view.Choreographer.FrameCallback
+import com.android.systemui.keyguard.data.repository.KeyguardTransitionRepository
+import com.android.systemui.keyguard.shared.model.TransitionInfo
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.Job
+import kotlinx.coroutines.delay
+import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.flow.collect
+import kotlinx.coroutines.launch
+import kotlinx.coroutines.withContext
+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
+    private var isTerminated = false
+
+    /**
+     * 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) {
+        // AnimationHandler uses ThreadLocal storage, and ValueAnimators MUST start from main
+        // thread
+        withContext(Dispatchers.Main) {
+            info.animator!!.getAnimationHandler().setProvider(this@KeyguardTransitionRunner)
+        }
+
+        job =
+            scope.launch {
+                frames.collect {
+                    val (frameNumber, callback) = it
+
+                    isTerminated = frameNumber >= maxFrames
+                    if (!isTerminated) {
+                        withContext(Dispatchers.Main) { callback?.doFrame(frameNumber) }
+                    }
+                }
+            }
+        withContext(Dispatchers.Main) { repository.startTransition(info) }
+
+        waitUntilComplete(info.animator!!)
+    }
+
+    suspend private fun waitUntilComplete(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")
+                }
+            }
+
+            animator.getAnimationHandler().setProvider(null)
+        }
+
+        job?.cancel()
+    }
+
+    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 {
+        private const val MAX_TEST_DURATION = 100L
+    }
+}
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/data/repository/FakeKeyguardRepository.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/data/repository/FakeKeyguardRepository.kt
@@ -120,6 +120,14 @@ class FakeKeyguardRepository : KeyguardRepository {
        _dozeAmount.value = dozeAmount
    }

+    fun setWakefulnessModel(model: WakefulnessModel) {
+        _wakefulnessModel.value = model
+    }
+
+    fun setBouncerShowing(isShowing: Boolean) {
+        _isBouncerShowing.value = isShowing
+    }
+
    fun setBiometricUnlockState(state: BiometricUnlockModel) {
        _biometricUnlockState.tryEmit(state)
    }
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/shade/data/repository/FakeShadeRepository.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/shade/data/repository/FakeShadeRepository.kt
@@ -0,0 +1,33 @@
+/*
+ * Copyright (C) 2022 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package com.android.systemui.shade.data.repository
+
+import com.android.systemui.shade.domain.model.ShadeModel
+import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.MutableStateFlow
+
+/** Fake implementation of [KeyguardRepository] */
+class FakeShadeRepository : ShadeRepository {
+
+    private val _shadeModel = MutableStateFlow(ShadeModel())
+    override val shadeModel: Flow<ShadeModel> = _shadeModel
+
+    fun setShadeModel(model: ShadeModel) {
+        _shadeModel.value = model
+    }
+}