Don't restart when the lockscreen is occluded.

Converts the ConditionalRestarter.Condition class to the more compact
flow syntax.

Fixes: 295275204
Test: atest SystemUITests:com.android.systemui.flags
Change-Id: I9659e19d260ec251143ef155f4755aea1201bbad

10 files changed, 267 insertions, 134 deletions

diff --git a/packages/SystemUI/src-debug/com/android/systemui/flags/FlagsModule.kt b/packages/SystemUI/src-debug/com/android/systemui/flags/FlagsModule.kt
index fae9fec0c26d..a263361c0ab2 100644
--- a/packages/SystemUI/src-debug/com/android/systemui/flags/FlagsModule.kt
+++ b/packages/SystemUI/src-debug/com/android/systemui/flags/FlagsModule.kt
@@ -39,6 +39,12 @@ abstract class FlagsModule {
     @IntoSet
     abstract fun bindsScreenIdleCondition(impl: ScreenIdleCondition): ConditionalRestarter.Condition
 
+    @Binds
+    @IntoSet
+    abstract fun bindsNotOccludedCondition(
+        impl: NotOccludedCondition
+    ): ConditionalRestarter.Condition
+
     @Module
     companion object {
         @JvmStatic
diff --git a/packages/SystemUI/src-release/com/android/systemui/flags/FlagsModule.kt b/packages/SystemUI/src-release/com/android/systemui/flags/FlagsModule.kt
index 7aacb4efba8e..9684d5e38fa7 100644
--- a/packages/SystemUI/src-release/com/android/systemui/flags/FlagsModule.kt
+++ b/packages/SystemUI/src-release/com/android/systemui/flags/FlagsModule.kt
@@ -39,6 +39,12 @@ abstract class FlagsModule {
     @IntoSet
     abstract fun bindsPluggedInCondition(impl: PluggedInCondition): ConditionalRestarter.Condition
 
+    @Binds
+    @IntoSet
+    abstract fun bindsNotOccludedCondition(
+        impl: NotOccludedCondition
+    ): ConditionalRestarter.Condition
+
     @Module
     companion object {
         @JvmStatic
diff --git a/packages/SystemUI/src/com/android/systemui/flags/ConditionalRestarter.kt b/packages/SystemUI/src/com/android/systemui/flags/ConditionalRestarter.kt
index 83c239f169ef..93b00eef08b5 100644
--- a/packages/SystemUI/src/com/android/systemui/flags/ConditionalRestarter.kt
+++ b/packages/SystemUI/src/com/android/systemui/flags/ConditionalRestarter.kt
@@ -19,13 +19,17 @@ package com.android.systemui.flags
 import android.util.Log
 import com.android.systemui.dagger.qualifiers.Application
 import com.android.systemui.dagger.qualifiers.Background
+import com.android.systemui.flags.ConditionalRestarter.Condition
 import java.util.concurrent.TimeUnit
 import javax.inject.Inject
 import javax.inject.Named
 import kotlinx.coroutines.CoroutineDispatcher
 import kotlinx.coroutines.CoroutineScope
-import kotlinx.coroutines.Job
 import kotlinx.coroutines.delay
+import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.combine
+import kotlinx.coroutines.flow.first
+import kotlinx.coroutines.flow.transformLatest
 import kotlinx.coroutines.launch
 
 /** Restarts the process after all passed in [Condition]s are true. */
@@ -39,7 +43,6 @@ constructor(
     @Background private val backgroundDispatcher: CoroutineDispatcher,
 ) : Restarter {
 
-    private var restartJob: Job? = null
     private var pendingReason = ""
     private var androidRestartRequested = false
 
@@ -57,17 +60,19 @@ constructor(
     private fun scheduleRestart(reason: String = "") {
         pendingReason = if (reason.isEmpty()) pendingReason else reason
 
-        if (conditions.all { c -> c.canRestartNow(this::scheduleRestart) }) {
-            if (restartJob == null) {
-                restartJob =
-                    applicationScope.launch(backgroundDispatcher) {
+        applicationScope.launch(backgroundDispatcher) {
+            combine(conditions.map { condition -> condition.canRestartNow }) { it.all { it } }
+                // Once all conditions are met, delay.
+                .transformLatest { allConditionsMet ->
+                    if (allConditionsMet) {
                         delay(TimeUnit.SECONDS.toMillis(restartDelaySec))
-                        restartNow()
+                        emit(Unit)
                     }
-            }
-        } else {
-            restartJob?.cancel()
-            restartJob = null
+                }
+                // Once we have successfully delayed _once_, continue to restart.
+                .first()
+
+            restartNow()
         }
     }
 
@@ -94,7 +99,7 @@ constructor(
          * multiple [Condition]s are being checked. If any one [Condition] returns false, all the
          * [Condition]s will need to be rechecked on the next restart attempt.
          */
-        fun canRestartNow(retryFn: () -> Unit): Boolean
+        val canRestartNow: Flow<Boolean>
     }
 
     companion object {
diff --git a/packages/SystemUI/src/com/android/systemui/flags/NotOccludedCondition.kt b/packages/SystemUI/src/com/android/systemui/flags/NotOccludedCondition.kt
new file mode 100644
index 000000000000..f5b30cf4b54f
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/flags/NotOccludedCondition.kt
@@ -0,0 +1,24 @@
+package com.android.systemui.flags
+
+import com.android.systemui.keyguard.domain.interactor.KeyguardTransitionInteractor
+import com.android.systemui.keyguard.shared.model.KeyguardState
+import dagger.Lazy
+import javax.inject.Inject
+import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.map
+
+/** Returns true when the device is "asleep" as defined by the [WakefullnessLifecycle]. */
+class NotOccludedCondition
+@Inject
+constructor(
+    private val keyguardTransitionInteractorLazy: Lazy<KeyguardTransitionInteractor>,
+) : ConditionalRestarter.Condition {
+
+    override val canRestartNow: Flow<Boolean>
+        get() {
+            return keyguardTransitionInteractorLazy
+                .get()
+                .transitionValue(KeyguardState.OCCLUDED)
+                .map { it == 0f }
+        }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/flags/PluggedInCondition.kt b/packages/SystemUI/src/com/android/systemui/flags/PluggedInCondition.kt
index 3120638cb17f..dc08570447a5 100644
--- a/packages/SystemUI/src/com/android/systemui/flags/PluggedInCondition.kt
+++ b/packages/SystemUI/src/com/android/systemui/flags/PluggedInCondition.kt
@@ -16,34 +16,34 @@
 
 package com.android.systemui.flags
 
+import com.android.systemui.common.coroutine.ConflatedCallbackFlow.conflatedCallbackFlow
 import com.android.systemui.statusbar.policy.BatteryController
+import dagger.Lazy
 import javax.inject.Inject
+import kotlinx.coroutines.channels.awaitClose
 
 /** Returns true when the device is plugged in. */
 class PluggedInCondition
 @Inject
 constructor(
-    private val batteryController: BatteryController,
+    private val batteryControllerLazy: Lazy<BatteryController>,
 ) : ConditionalRestarter.Condition {
 
-    var listenersAdded = false
-    var retryFn: (() -> Unit)? = null
-
-    val batteryCallback =
-        object : BatteryController.BatteryStateChangeCallback {
-            override fun onBatteryLevelChanged(level: Int, pluggedIn: Boolean, charging: Boolean) {
-                retryFn?.invoke()
+    override val canRestartNow = conflatedCallbackFlow {
+        val batteryCallback =
+            object : BatteryController.BatteryStateChangeCallback {
+                override fun onBatteryLevelChanged(
+                    level: Int,
+                    pluggedIn: Boolean,
+                    charging: Boolean
+                ) {
+                    trySend(pluggedIn)
+                }
             }
-        }
-
-    override fun canRestartNow(retryFn: () -> Unit): Boolean {
-        if (!listenersAdded) {
-            listenersAdded = true
-            batteryController.addCallback(batteryCallback)
-        }
+        batteryControllerLazy.get().addCallback(batteryCallback)
 
-        this.retryFn = retryFn
+        trySend(batteryControllerLazy.get().isPluggedIn)
 
-        return batteryController.isPluggedIn
+        awaitClose { batteryControllerLazy.get().removeCallback(batteryCallback) }
     }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/flags/ScreenIdleCondition.kt b/packages/SystemUI/src/com/android/systemui/flags/ScreenIdleCondition.kt
index 49e61afbdcd6..3c9bc368e852 100644
--- a/packages/SystemUI/src/com/android/systemui/flags/ScreenIdleCondition.kt
+++ b/packages/SystemUI/src/com/android/systemui/flags/ScreenIdleCondition.kt
@@ -16,34 +16,19 @@
 
 package com.android.systemui.flags
 
-import com.android.systemui.keyguard.WakefulnessLifecycle
+import com.android.systemui.power.domain.interactor.PowerInteractor
+import dagger.Lazy
 import javax.inject.Inject
+import kotlinx.coroutines.flow.Flow
 
 /** Returns true when the device is "asleep" as defined by the [WakefullnessLifecycle]. */
 class ScreenIdleCondition
 @Inject
-constructor(
-    private val wakefulnessLifecycle: WakefulnessLifecycle,
-) : ConditionalRestarter.Condition {
+constructor(private val powerInteractorLazy: Lazy<PowerInteractor>) :
+    ConditionalRestarter.Condition {
 
-    var listenersAdded = false
-    var retryFn: (() -> Unit)? = null
-
-    val observer =
-        object : WakefulnessLifecycle.Observer {
-            override fun onFinishedGoingToSleep() {
-                retryFn?.invoke()
-            }
-        }
-
-    override fun canRestartNow(retryFn: () -> Unit): Boolean {
-        if (!listenersAdded) {
-            listenersAdded = true
-            wakefulnessLifecycle.addObserver(observer)
+    override val canRestartNow: Flow<Boolean>
+        get() {
+            return powerInteractorLazy.get().isAsleep
         }
-
-        this.retryFn = retryFn
-
-        return wakefulnessLifecycle.wakefulness == WakefulnessLifecycle.WAKEFULNESS_ASLEEP
-    }
 }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/flags/ConditionalRestarterTest.kt b/packages/SystemUI/tests/src/com/android/systemui/flags/ConditionalRestarterTest.kt
index 0e14591c5f53..52c6e22cfcbb 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/flags/ConditionalRestarterTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/flags/ConditionalRestarterTest.kt
@@ -18,9 +18,11 @@ package com.android.systemui.flags
 import androidx.test.filters.SmallTest
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.util.mockito.any
+import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.test.StandardTestDispatcher
 import kotlinx.coroutines.test.TestScope
-import kotlinx.coroutines.test.advanceUntilIdle
+import kotlinx.coroutines.test.runCurrent
 import kotlinx.coroutines.test.runTest
 import org.junit.Before
 import org.junit.Test
@@ -61,80 +63,70 @@ class ConditionalRestarterTest : SysuiTestCase() {
     @Test
     fun restart_ImmediatelySatisfied() =
         testScope.runTest {
-            conditionA.canRestart = true
-            conditionB.canRestart = true
+            conditionA.canRestart.emit(true)
+            conditionB.canRestart.emit(true)
             restarter.restartSystemUI("Restart for test")
-            advanceUntilIdle()
+            runCurrent()
             verify(systemExitRestarter).restartSystemUI(any())
         }
 
     @Test
     fun restart_WaitsForConditionA() =
         testScope.runTest {
-            conditionA.canRestart = false
-            conditionB.canRestart = true
+            conditionA.canRestart.emit(false)
+            conditionB.canRestart.emit(true)
 
             restarter.restartSystemUI("Restart for test")
-            advanceUntilIdle()
+            runCurrent()
             // No restart occurs yet.
             verify(systemExitRestarter, never()).restartSystemUI(any())
 
-            conditionA.canRestart = true
-            conditionA.retryFn?.invoke()
-            advanceUntilIdle()
+            conditionA.canRestart.emit(true)
+            runCurrent()
             verify(systemExitRestarter).restartSystemUI(any())
         }
 
     @Test
     fun restart_WaitsForConditionB() =
         testScope.runTest {
-            conditionA.canRestart = true
-            conditionB.canRestart = false
+            conditionA.canRestart.emit(true)
+            conditionB.canRestart.emit(false)
 
             restarter.restartSystemUI("Restart for test")
-            advanceUntilIdle()
+            runCurrent()
             // No restart occurs yet.
             verify(systemExitRestarter, never()).restartSystemUI(any())
 
-            conditionB.canRestart = true
-            conditionB.retryFn?.invoke()
-            advanceUntilIdle()
+            conditionB.canRestart.emit(true)
+            runCurrent()
             verify(systemExitRestarter).restartSystemUI(any())
         }
 
     @Test
     fun restart_WaitsForAllConditions() =
         testScope.runTest {
-            conditionA.canRestart = true
-            conditionB.canRestart = false
+            conditionA.canRestart.emit(true)
+            conditionB.canRestart.emit(false)
 
             restarter.restartSystemUI("Restart for test")
-            advanceUntilIdle()
+            runCurrent()
             // No restart occurs yet.
             verify(systemExitRestarter, never()).restartSystemUI(any())
 
             // B becomes true, but A is now false
-            conditionA.canRestart = false
-            conditionB.canRestart = true
-            conditionB.retryFn?.invoke()
-            advanceUntilIdle()
+            conditionA.canRestart.emit(false)
+            conditionB.canRestart.emit(true)
             // No restart occurs yet.
             verify(systemExitRestarter, never()).restartSystemUI(any())
 
-            conditionA.canRestart = true
-            conditionA.retryFn?.invoke()
-            advanceUntilIdle()
+            conditionA.canRestart.emit(true)
+            runCurrent()
             verify(systemExitRestarter).restartSystemUI(any())
         }
 
     class FakeCondition : ConditionalRestarter.Condition {
-        var retryFn: (() -> Unit)? = null
-        var canRestart = false
+        val canRestart = MutableStateFlow(false)
 
-        override fun canRestartNow(retryFn: () -> Unit): Boolean {
-            this.retryFn = retryFn
-
-            return canRestart
-        }
+        override val canRestartNow: Flow<Boolean> = canRestart
     }
 }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/flags/NotOccludedConditionTest.kt b/packages/SystemUI/tests/src/com/android/systemui/flags/NotOccludedConditionTest.kt
new file mode 100644
index 000000000000..db6f85f12a42
--- /dev/null
+++ b/packages/SystemUI/tests/src/com/android/systemui/flags/NotOccludedConditionTest.kt
@@ -0,0 +1,91 @@
+/*
+ * Copyright (C) 2023 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.flags
+
+import android.test.suitebuilder.annotation.SmallTest
+import com.android.systemui.SysuiTestCase
+import com.android.systemui.coroutines.collectLastValue
+import com.android.systemui.keyguard.domain.interactor.KeyguardTransitionInteractor
+import com.android.systemui.keyguard.shared.model.KeyguardState
+import com.google.common.truth.Truth.assertThat
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.test.StandardTestDispatcher
+import kotlinx.coroutines.test.TestDispatcher
+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.mockito.Mock
+import org.mockito.Mockito.`when` as whenever
+import org.mockito.MockitoAnnotations
+
+/**
+ * Be careful with the {FeatureFlagsReleaseRestarter} in this test. It has a call to System.exit()!
+ */
+@OptIn(ExperimentalCoroutinesApi::class)
+@SmallTest
+class NotOccludedConditionTest : SysuiTestCase() {
+    private lateinit var condition: NotOccludedCondition
+
+    @Mock private lateinit var keyguardTransitionInteractor: KeyguardTransitionInteractor
+    private val transitionValue = MutableStateFlow(0f)
+
+    private val testDispatcher: TestDispatcher = StandardTestDispatcher()
+    private val testScope: TestScope = TestScope(testDispatcher)
+
+    @Before
+    fun setup() {
+        MockitoAnnotations.initMocks(this)
+        whenever(keyguardTransitionInteractor.transitionValue(KeyguardState.OCCLUDED))
+            .thenReturn(transitionValue)
+        condition = NotOccludedCondition({ keyguardTransitionInteractor })
+        testScope.runCurrent()
+    }
+
+    @Test
+    fun testCondition_occluded() =
+        testScope.runTest {
+            val canRestart by collectLastValue(condition.canRestartNow)
+
+            transitionValue.emit(1f)
+            assertThat(canRestart).isFalse()
+        }
+
+    @Test
+    fun testCondition_notOccluded() =
+        testScope.runTest {
+            val canRestart by collectLastValue(condition.canRestartNow)
+
+            transitionValue.emit(0f)
+            assertThat(canRestart).isTrue()
+        }
+
+    @Test
+    fun testCondition_invokesRetry() =
+        testScope.runTest {
+            val canRestart by collectLastValue(condition.canRestartNow)
+
+            transitionValue.emit(1f)
+
+            assertThat(canRestart).isFalse()
+
+            transitionValue.emit(0f)
+
+            assertThat(canRestart).isTrue()
+        }
+}
diff --git a/packages/SystemUI/tests/src/com/android/systemui/flags/PluggedInConditionTest.kt b/packages/SystemUI/tests/src/com/android/systemui/flags/PluggedInConditionTest.kt
index 647b05a77b90..7d7abab4a0f5 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/flags/PluggedInConditionTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/flags/PluggedInConditionTest.kt
@@ -17,8 +17,13 @@ package com.android.systemui.flags
 
 import android.test.suitebuilder.annotation.SmallTest
 import com.android.systemui.SysuiTestCase
+import com.android.systemui.coroutines.collectLastValue
 import com.android.systemui.statusbar.policy.BatteryController
 import com.google.common.truth.Truth.assertThat
+import kotlinx.coroutines.test.StandardTestDispatcher
+import kotlinx.coroutines.test.TestDispatcher
+import kotlinx.coroutines.test.TestScope
+import kotlinx.coroutines.test.runTest
 import org.junit.Before
 import org.junit.Test
 import org.mockito.ArgumentCaptor
@@ -35,42 +40,51 @@ class PluggedInConditionTest : SysuiTestCase() {
     private lateinit var condition: PluggedInCondition
 
     @Mock private lateinit var batteryController: BatteryController
+    private val testDispatcher: TestDispatcher = StandardTestDispatcher()
+    private val testScope: TestScope = TestScope(testDispatcher)
+    private val callbackCaptor =
+        ArgumentCaptor.forClass(BatteryController.BatteryStateChangeCallback::class.java)
 
     @Before
     fun setup() {
         MockitoAnnotations.initMocks(this)
-        condition = PluggedInCondition(batteryController)
+
+        condition = PluggedInCondition({ batteryController })
     }
 
     @Test
-    fun testCondition_unplugged() {
-        whenever(batteryController.isPluggedIn).thenReturn(false)
+    fun testCondition_unplugged() =
+        testScope.runTest {
+            whenever(batteryController.isPluggedIn).thenReturn(false)
 
-        assertThat(condition.canRestartNow({})).isFalse()
-    }
+            val canRestart by collectLastValue(condition.canRestartNow)
+
+            assertThat(canRestart).isFalse()
+        }
 
     @Test
-    fun testCondition_pluggedIn() {
-        whenever(batteryController.isPluggedIn).thenReturn(true)
+    fun testCondition_pluggedIn() =
+        testScope.runTest {
+            whenever(batteryController.isPluggedIn).thenReturn(true)
 
-        assertThat(condition.canRestartNow({})).isTrue()
-    }
+            val canRestart by collectLastValue(condition.canRestartNow)
+
+            assertThat(canRestart).isTrue()
+        }
 
     @Test
-    fun testCondition_invokesRetry() {
-        whenever(batteryController.isPluggedIn).thenReturn(false)
-        var retried = false
-        val retryFn = { retried = true }
+    fun testCondition_invokesRetry() =
+        testScope.runTest {
+            whenever(batteryController.isPluggedIn).thenReturn(false)
 
-        // No restart yet, but we do register a listener now.
-        assertThat(condition.canRestartNow(retryFn)).isFalse()
-        val captor =
-            ArgumentCaptor.forClass(BatteryController.BatteryStateChangeCallback::class.java)
-        verify(batteryController).addCallback(captor.capture())
+            val canRestart by collectLastValue(condition.canRestartNow)
 
-        whenever(batteryController.isPluggedIn).thenReturn(true)
+            assertThat(canRestart).isFalse()
 
-        captor.value.onBatteryLevelChanged(0, true, true)
-        assertThat(retried).isTrue()
-    }
+            verify(batteryController).addCallback(callbackCaptor.capture())
+
+            callbackCaptor.value.onBatteryLevelChanged(0, true, false)
+
+            assertThat(canRestart).isTrue()
+        }
 }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/flags/ScreenIdleConditionTest.kt b/packages/SystemUI/tests/src/com/android/systemui/flags/ScreenIdleConditionTest.kt
index f7a773ea30ec..1f04828c359a 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/flags/ScreenIdleConditionTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/flags/ScreenIdleConditionTest.kt
@@ -17,15 +17,17 @@ package com.android.systemui.flags
 
 import android.test.suitebuilder.annotation.SmallTest
 import com.android.systemui.SysuiTestCase
-import com.android.systemui.keyguard.WakefulnessLifecycle
-import com.android.systemui.keyguard.WakefulnessLifecycle.WAKEFULNESS_ASLEEP
-import com.android.systemui.keyguard.WakefulnessLifecycle.WAKEFULNESS_AWAKE
+import com.android.systemui.coroutines.collectLastValue
+import com.android.systemui.power.domain.interactor.PowerInteractor
 import com.google.common.truth.Truth.assertThat
+import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.test.StandardTestDispatcher
+import kotlinx.coroutines.test.TestDispatcher
+import kotlinx.coroutines.test.TestScope
+import kotlinx.coroutines.test.runTest
 import org.junit.Before
 import org.junit.Test
-import org.mockito.ArgumentCaptor
 import org.mockito.Mock
-import org.mockito.Mockito.verify
 import org.mockito.Mockito.`when` as whenever
 import org.mockito.MockitoAnnotations
 
@@ -36,42 +38,50 @@ import org.mockito.MockitoAnnotations
 class ScreenIdleConditionTest : SysuiTestCase() {
     private lateinit var condition: ScreenIdleCondition
 
-    @Mock private lateinit var wakefulnessLifecycle: WakefulnessLifecycle
+    @Mock private lateinit var powerInteractor: PowerInteractor
+    private val isAsleep = MutableStateFlow(false)
+
+    private val testDispatcher: TestDispatcher = StandardTestDispatcher()
+    private val testScope: TestScope = TestScope(testDispatcher)
 
     @Before
     fun setup() {
         MockitoAnnotations.initMocks(this)
-        condition = ScreenIdleCondition(wakefulnessLifecycle)
+        whenever(powerInteractor.isAsleep).thenReturn(isAsleep)
+        condition = ScreenIdleCondition({ powerInteractor })
     }
 
     @Test
-    fun testCondition_awake() {
-        whenever(wakefulnessLifecycle.wakefulness).thenReturn(WAKEFULNESS_AWAKE)
+    fun testCondition_awake() =
+        testScope.runTest {
+            val canRestart by collectLastValue(condition.canRestartNow)
 
-        assertThat(condition.canRestartNow {}).isFalse()
-    }
+            isAsleep.emit(false)
+
+            assertThat(canRestart).isFalse()
+        }
 
     @Test
-    fun testCondition_asleep() {
-        whenever(wakefulnessLifecycle.wakefulness).thenReturn(WAKEFULNESS_ASLEEP)
+    fun testCondition_asleep() =
+        testScope.runTest {
+            val canRestart by collectLastValue(condition.canRestartNow)
 
-        assertThat(condition.canRestartNow {}).isTrue()
-    }
+            isAsleep.emit(true)
+
+            assertThat(canRestart).isTrue()
+        }
 
     @Test
-    fun testCondition_invokesRetry() {
-        whenever(wakefulnessLifecycle.wakefulness).thenReturn(WAKEFULNESS_AWAKE)
-        var retried = false
-        val retryFn = { retried = true }
+    fun testCondition_invokesRetry() =
+        testScope.runTest {
+            val canRestart by collectLastValue(condition.canRestartNow)
 
-        // No restart yet, but we do register a listener now.
-        assertThat(condition.canRestartNow(retryFn)).isFalse()
-        val captor = ArgumentCaptor.forClass(WakefulnessLifecycle.Observer::class.java)
-        verify(wakefulnessLifecycle).addObserver(captor.capture())
+            isAsleep.emit(false)
 
-        whenever(wakefulnessLifecycle.wakefulness).thenReturn(WAKEFULNESS_ASLEEP)
+            assertThat(canRestart).isFalse()
 
-        captor.value.onFinishedGoingToSleep()
-        assertThat(retried).isTrue()
-    }
+            isAsleep.emit(true)
+
+            assertThat(canRestart).isTrue()
+        }
 }