Delay restarting SystemUI on debug builds.

When a flag changes, delay restarting systemui by 1 second after the
screen goes idle. This prevents a race condition with entering
dream/doze/aod.

Also, refactor `FeatureFlagsReleaseRestarter` and
`FeatureFlagsDebugRestarter` into `ConditionalRestarter`, since they
share a great deal of the same functionality. This also has the
nice effect of making the tests smaller and more focused.

Fixes: 268060399
Test: atest SystemUITest && manual testing
Change-Id: Iae7ddc5965a95c72d35caafc2c90b267be7c3a6d

packages/SystemUI/src-debug/com/android/systemui/flags/FlagsModule.kt

@@ -23,6 +23,8 @@ import com.android.systemui.util.settings.SettingsUtilModule
 import dagger.Binds
 import dagger.Module
 import dagger.Provides
+import dagger.multibindings.IntoSet
+import javax.inject.Named
 
 @Module(includes = [
     FeatureFlagsDebugStartableModule::class,
@@ -35,7 +37,8 @@ abstract class FlagsModule {
     abstract fun bindsFeatureFlagDebug(impl: FeatureFlagsDebug): FeatureFlags
 
     @Binds
-    abstract fun bindsRestarter(debugRestarter: FeatureFlagsDebugRestarter): Restarter
+    @IntoSet
+    abstract fun bindsScreenIdleCondition(impl: ScreenIdleCondition): ConditionalRestarter.Condition
 
     @Module
     companion object {
@@ -44,5 +47,10 @@ abstract class FlagsModule {
         fun provideFlagManager(context: Context, @Main handler: Handler): FlagManager {
             return FlagManager(context, handler)
         }
+
+        @JvmStatic
+        @Provides
+        @Named(ConditionalRestarter.RESTART_DELAY)
+        fun provideRestartDelaySec(): Long = 1
     }
 }

packages/SystemUI/src-release/com/android/systemui/flags/FlagsModule.kt

@@ -18,6 +18,9 @@ package com.android.systemui.flags
 
 import dagger.Binds
 import dagger.Module
+import dagger.Provides
+import dagger.multibindings.IntoSet
+import javax.inject.Named
 
 @Module(includes = [
     FeatureFlagsReleaseStartableModule::class,
@@ -29,5 +32,18 @@ abstract class FlagsModule {
     abstract fun bindsFeatureFlagRelease(impl: FeatureFlagsRelease): FeatureFlags
 
     @Binds
-    abstract fun bindsRestarter(debugRestarter: FeatureFlagsReleaseRestarter): Restarter
+    @IntoSet
+    abstract fun bindsScreenIdleCondition(impl: ScreenIdleCondition): ConditionalRestarter.Condition
+
+    @Binds
+    @IntoSet
+    abstract fun bindsPluggedInCondition(impl: PluggedInCondition): ConditionalRestarter.Condition
+
+    @Module
+    companion object {
+        @JvmStatic
+        @Provides
+        @Named(ConditionalRestarter.RESTART_DELAY)
+        fun provideRestartDelaySec(): Long = 30
+    }
 }

packages/SystemUI/src/com/android/systemui/flags/ConditionalRestarter.kt

@@ -0,0 +1,103 @@
+/*
+ * 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.util.Log
+import com.android.systemui.dagger.qualifiers.Application
+import com.android.systemui.dagger.qualifiers.Background
+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.launch
+
+/** Restarts the process after all passed in [Condition]s are true. */
+class ConditionalRestarter
+@Inject
+constructor(
+    private val systemExitRestarter: SystemExitRestarter,
+    private val conditions: Set<@JvmSuppressWildcards Condition>,
+    @Named(RESTART_DELAY) private val restartDelaySec: Long,
+    @Application private val applicationScope: CoroutineScope,
+    @Background private val backgroundDispatcher: CoroutineDispatcher,
+) : Restarter {
+
+    private var restartJob: Job? = null
+    private var pendingReason = ""
+    private var androidRestartRequested = false
+
+    override fun restartSystemUI(reason: String) {
+        Log.d(FeatureFlagsDebug.TAG, "SystemUI Restart requested. Restarting when idle.")
+        scheduleRestart(reason)
+    }
+
+    override fun restartAndroid(reason: String) {
+        Log.d(FeatureFlagsDebug.TAG, "Android Restart requested. Restarting when idle.")
+        androidRestartRequested = true
+        scheduleRestart(reason)
+    }
+
+    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) {
+                        delay(TimeUnit.SECONDS.toMillis(restartDelaySec))
+                        restartNow()
+                    }
+            }
+        } else {
+            restartJob?.cancel()
+            restartJob = null
+        }
+    }
+
+    private fun restartNow() {
+        if (androidRestartRequested) {
+            systemExitRestarter.restartAndroid(pendingReason)
+        } else {
+            systemExitRestarter.restartSystemUI(pendingReason)
+        }
+    }
+
+    interface Condition {
+        /**
+         * Should return true if the system is ready to restart.
+         *
+         * A call to this function means that we want to restart and are waiting for this condition
+         * to return true.
+         *
+         * retryFn should be cached if it is _not_ ready to restart, and later called when it _is_
+         * ready to restart. At that point, this method will be called again to verify that the
+         * system is ready.
+         *
+         * Multiple calls to an instance of this method may happen for a single restart attempt if
+         * 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
+    }
+
+    companion object {
+        const val RESTART_DELAY = "restarter_restart_delay"
+    }
+}

packages/SystemUI/src/com/android/systemui/flags/FeatureFlagsDebugRestarter.kt

@@ -1,70 +0,0 @@
-/*
- * 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.flags
-
-import android.util.Log
-import com.android.systemui.keyguard.WakefulnessLifecycle
-import javax.inject.Inject
-
-/** Restarts SystemUI when the screen is locked. */
-class FeatureFlagsDebugRestarter
-@Inject
-constructor(
-    private val wakefulnessLifecycle: WakefulnessLifecycle,
-    private val systemExitRestarter: SystemExitRestarter,
-) : Restarter {
-
-    private var androidRestartRequested = false
-    private var pendingReason = ""
-
-    val observer =
-        object : WakefulnessLifecycle.Observer {
-            override fun onFinishedGoingToSleep() {
-                Log.d(FeatureFlagsDebug.TAG, "Restarting due to systemui flag change")
-                restartNow()
-            }
-        }
-
-    override fun restartSystemUI(reason: String) {
-        Log.d(FeatureFlagsDebug.TAG, "SystemUI Restart requested. Restarting on next screen off.")
-        Log.i(FeatureFlagsDebug.TAG, reason)
-        scheduleRestart(reason)
-    }
-
-    override fun restartAndroid(reason: String) {
-        Log.d(FeatureFlagsDebug.TAG, "Android Restart requested. Restarting on next screen off.")
-        androidRestartRequested = true
-        scheduleRestart(reason)
-    }
-
-    fun scheduleRestart(reason: String) {
-        pendingReason = reason
-        if (wakefulnessLifecycle.wakefulness == WakefulnessLifecycle.WAKEFULNESS_ASLEEP) {
-            restartNow()
-        } else {
-            wakefulnessLifecycle.addObserver(observer)
-        }
-    }
-
-    private fun restartNow() {
-        if (androidRestartRequested) {
-            systemExitRestarter.restartAndroid(pendingReason)
-        } else {
-            systemExitRestarter.restartSystemUI(pendingReason)
-        }
-    }
-}

packages/SystemUI/src/com/android/systemui/flags/FeatureFlagsReleaseRestarter.kt

@@ -1,101 +0,0 @@
-/*
- * 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.flags
-
-import android.util.Log
-import com.android.systemui.dagger.qualifiers.Background
-import com.android.systemui.keyguard.WakefulnessLifecycle
-import com.android.systemui.keyguard.WakefulnessLifecycle.WAKEFULNESS_ASLEEP
-import com.android.systemui.statusbar.policy.BatteryController
-import com.android.systemui.util.concurrency.DelayableExecutor
-import java.util.concurrent.TimeUnit
-import javax.inject.Inject
-
-/** Restarts SystemUI when the device appears idle. */
-class FeatureFlagsReleaseRestarter
-@Inject
-constructor(
-    private val wakefulnessLifecycle: WakefulnessLifecycle,
-    private val batteryController: BatteryController,
-    @Background private val bgExecutor: DelayableExecutor,
-    private val systemExitRestarter: SystemExitRestarter
-) : Restarter {
-    var listenersAdded = false
-    var pendingRestart: Runnable? = null
-    private var pendingReason = ""
-    var androidRestartRequested = false
-
-    val observer =
-        object : WakefulnessLifecycle.Observer {
-            override fun onFinishedGoingToSleep() {
-                scheduleRestart(pendingReason)
-            }
-        }
-
-    val batteryCallback =
-        object : BatteryController.BatteryStateChangeCallback {
-            override fun onBatteryLevelChanged(level: Int, pluggedIn: Boolean, charging: Boolean) {
-                scheduleRestart(pendingReason)
-            }
-        }
-
-    override fun restartSystemUI(reason: String) {
-        Log.d(
-            FeatureFlagsDebug.TAG,
-            "SystemUI Restart requested. Restarting when plugged in and idle."
-        )
-        scheduleRestart(reason)
-    }
-
-    override fun restartAndroid(reason: String) {
-        Log.d(
-            FeatureFlagsDebug.TAG,
-            "Android Restart requested. Restarting when plugged in and idle."
-        )
-        androidRestartRequested = true
-        scheduleRestart(reason)
-    }
-
-    private fun scheduleRestart(reason: String) {
-        // Don't bother adding listeners twice.
-        pendingReason = reason
-        if (!listenersAdded) {
-            listenersAdded = true
-            wakefulnessLifecycle.addObserver(observer)
-            batteryController.addCallback(batteryCallback)
-        }
-        if (
-            wakefulnessLifecycle.wakefulness == WAKEFULNESS_ASLEEP && batteryController.isPluggedIn
-        ) {
-            if (pendingRestart == null) {
-                pendingRestart = bgExecutor.executeDelayed(this::restartNow, 30L, TimeUnit.SECONDS)
-            }
-        } else if (pendingRestart != null) {
-            pendingRestart?.run()
-            pendingRestart = null
-        }
-    }
-
-    private fun restartNow() {
-        Log.d(FeatureFlagsRelease.TAG, "Restarting due to systemui flag change")
-        if (androidRestartRequested) {
-            systemExitRestarter.restartAndroid(pendingReason)
-        } else {
-            systemExitRestarter.restartSystemUI(pendingReason)
-        }
-    }
-}

packages/SystemUI/src/com/android/systemui/flags/FlagsCommonModule.kt

@@ -15,6 +15,7 @@
  */
 package com.android.systemui.flags
 
+import dagger.Binds
 import dagger.Module
 import dagger.Provides
 import javax.inject.Named
@@ -22,6 +23,8 @@ import javax.inject.Named
 /** Module containing shared code for all FeatureFlag implementations. */
 @Module
 interface FlagsCommonModule {
+    @Binds fun bindsRestarter(impl: ConditionalRestarter): Restarter
+
     companion object {
         const val ALL_FLAGS = "all_flags"
 

packages/SystemUI/src/com/android/systemui/flags/PluggedInCondition.kt

@@ -0,0 +1,49 @@
+/*
+ * 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 com.android.systemui.statusbar.policy.BatteryController
+import javax.inject.Inject
+
+/** Returns true when the device is plugged in. */
+class PluggedInCondition
+@Inject
+constructor(
+    private val batteryController: 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 fun canRestartNow(retryFn: () -> Unit): Boolean {
+        if (!listenersAdded) {
+            listenersAdded = true
+            batteryController.addCallback(batteryCallback)
+        }
+
+        this.retryFn = retryFn
+
+        return batteryController.isPluggedIn
+    }
+}

packages/SystemUI/src/com/android/systemui/flags/ScreenIdleCondition.kt

@@ -0,0 +1,49 @@
+/*
+ * 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 com.android.systemui.keyguard.WakefulnessLifecycle
+import javax.inject.Inject
+
+/** Returns true when the device is "asleep" as defined by the [WakefullnessLifecycle]. */
+class ScreenIdleCondition
+@Inject
+constructor(
+    private val wakefulnessLifecycle: WakefulnessLifecycle,
+) : 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)
+        }
+
+        this.retryFn = retryFn
+
+        return wakefulnessLifecycle.wakefulness == WakefulnessLifecycle.WAKEFULNESS_ASLEEP
+    }
+}

packages/SystemUI/tests/src/com/android/systemui/flags/ConditionalRestarterTest.kt

@@ -0,0 +1,140 @@
+/*
+ * 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 androidx.test.filters.SmallTest
+import com.android.systemui.SysuiTestCase
+import com.android.systemui.util.mockito.any
+import kotlinx.coroutines.test.StandardTestDispatcher
+import kotlinx.coroutines.test.TestScope
+import kotlinx.coroutines.test.advanceUntilIdle
+import kotlinx.coroutines.test.runTest
+import org.junit.Before
+import org.junit.Test
+import org.mockito.Mock
+import org.mockito.Mockito.never
+import org.mockito.Mockito.verify
+import org.mockito.MockitoAnnotations
+
+/**
+ * Be careful with the {FeatureFlagsReleaseRestarter} in this test. It has a call to System.exit()!
+ */
+@SmallTest
+class ConditionalRestarterTest : SysuiTestCase() {
+    private lateinit var restarter: ConditionalRestarter
+
+    @Mock private lateinit var systemExitRestarter: SystemExitRestarter
+
+    val restartDelayMs = 0L
+    val dispatcher = StandardTestDispatcher()
+    val testScope = TestScope(dispatcher)
+
+    val conditionA = FakeCondition()
+    val conditionB = FakeCondition()
+
+    @Before
+    fun setup() {
+        MockitoAnnotations.initMocks(this)
+        restarter =
+            ConditionalRestarter(
+                systemExitRestarter,
+                setOf(conditionA, conditionB),
+                restartDelayMs,
+                testScope,
+                dispatcher
+            )
+    }
+
+    @Test
+    fun restart_ImmediatelySatisfied() =
+        testScope.runTest {
+            conditionA.canRestart = true
+            conditionB.canRestart = true
+            restarter.restartSystemUI("Restart for test")
+            advanceUntilIdle()
+            verify(systemExitRestarter).restartSystemUI(any())
+        }
+
+    @Test
+    fun restart_WaitsForConditionA() =
+        testScope.runTest {
+            conditionA.canRestart = false
+            conditionB.canRestart = true
+
+            restarter.restartSystemUI("Restart for test")
+            advanceUntilIdle()
+            // No restart occurs yet.
+            verify(systemExitRestarter, never()).restartSystemUI(any())
+
+            conditionA.canRestart = true
+            conditionA.retryFn?.invoke()
+            advanceUntilIdle()
+            verify(systemExitRestarter).restartSystemUI(any())
+        }
+
+    @Test
+    fun restart_WaitsForConditionB() =
+        testScope.runTest {
+            conditionA.canRestart = true
+            conditionB.canRestart = false
+
+            restarter.restartSystemUI("Restart for test")
+            advanceUntilIdle()
+            // No restart occurs yet.
+            verify(systemExitRestarter, never()).restartSystemUI(any())
+
+            conditionB.canRestart = true
+            conditionB.retryFn?.invoke()
+            advanceUntilIdle()
+            verify(systemExitRestarter).restartSystemUI(any())
+        }
+
+    @Test
+    fun restart_WaitsForAllConditions() =
+        testScope.runTest {
+            conditionA.canRestart = true
+            conditionB.canRestart = false
+
+            restarter.restartSystemUI("Restart for test")
+            advanceUntilIdle()
+            // 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()
+            // No restart occurs yet.
+            verify(systemExitRestarter, never()).restartSystemUI(any())
+
+            conditionA.canRestart = true
+            conditionA.retryFn?.invoke()
+            advanceUntilIdle()
+            verify(systemExitRestarter).restartSystemUI(any())
+        }
+
+    class FakeCondition : ConditionalRestarter.Condition {
+        var retryFn: (() -> Unit)? = null
+        var canRestart = false
+
+        override fun canRestartNow(retryFn: () -> Unit): Boolean {
+            this.retryFn = retryFn
+
+            return canRestart
+        }
+    }
+}

packages/SystemUI/tests/src/com/android/systemui/flags/FeatureFlagsReleaseRestarterTest.kt

@@ -1,146 +0,0 @@
-/*
- * Copyright (C) 2021 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package com.android.systemui.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.statusbar.policy.BatteryController
-import com.android.systemui.util.concurrency.FakeExecutor
-import com.android.systemui.util.mockito.any
-import com.android.systemui.util.time.FakeSystemClock
-import com.google.common.truth.Truth.assertThat
-import org.junit.Before
-import org.junit.Test
-import org.mockito.ArgumentCaptor
-import org.mockito.Mock
-import org.mockito.Mockito.never
-import org.mockito.Mockito.verify
-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()!
- */
-@SmallTest
-class FeatureFlagsReleaseRestarterTest : SysuiTestCase() {
-    private lateinit var restarter: FeatureFlagsReleaseRestarter
-
-    @Mock private lateinit var wakefulnessLifecycle: WakefulnessLifecycle
-    @Mock private lateinit var batteryController: BatteryController
-    @Mock private lateinit var systemExitRestarter: SystemExitRestarter
-    private val executor = FakeExecutor(FakeSystemClock())
-
-    @Before
-    fun setup() {
-        MockitoAnnotations.initMocks(this)
-        restarter =
-            FeatureFlagsReleaseRestarter(
-                wakefulnessLifecycle,
-                batteryController,
-                executor,
-                systemExitRestarter
-            )
-    }
-
-    @Test
-    fun testRestart_ScheduledWhenReady() {
-        whenever(wakefulnessLifecycle.wakefulness).thenReturn(WAKEFULNESS_ASLEEP)
-        whenever(batteryController.isPluggedIn).thenReturn(true)
-
-        assertThat(executor.numPending()).isEqualTo(0)
-        restarter.restartSystemUI("Restart for test")
-        assertThat(executor.numPending()).isEqualTo(1)
-    }
-
-    @Test
-    fun testRestart_RestartsWhenIdle() {
-        whenever(wakefulnessLifecycle.wakefulness).thenReturn(WAKEFULNESS_ASLEEP)
-        whenever(batteryController.isPluggedIn).thenReturn(true)
-
-        restarter.restartSystemUI("Restart for test")
-        verify(systemExitRestarter, never()).restartSystemUI("Restart for test")
-        executor.advanceClockToLast()
-        executor.runAllReady()
-        verify(systemExitRestarter).restartSystemUI(any())
-    }
-
-    @Test
-    fun testRestart_NotScheduledWhenAwake() {
-        whenever(wakefulnessLifecycle.wakefulness).thenReturn(WAKEFULNESS_AWAKE)
-        whenever(batteryController.isPluggedIn).thenReturn(true)
-
-        assertThat(executor.numPending()).isEqualTo(0)
-        restarter.restartSystemUI("Restart for test")
-        assertThat(executor.numPending()).isEqualTo(0)
-    }
-
-    @Test
-    fun testRestart_NotScheduledWhenNotPluggedIn() {
-        whenever(wakefulnessLifecycle.wakefulness).thenReturn(WAKEFULNESS_ASLEEP)
-        whenever(batteryController.isPluggedIn).thenReturn(false)
-
-        assertThat(executor.numPending()).isEqualTo(0)
-        restarter.restartSystemUI("Restart for test")
-        assertThat(executor.numPending()).isEqualTo(0)
-    }
-
-    @Test
-    fun testRestart_NotDoubleSheduled() {
-        whenever(wakefulnessLifecycle.wakefulness).thenReturn(WAKEFULNESS_ASLEEP)
-        whenever(batteryController.isPluggedIn).thenReturn(true)
-
-        assertThat(executor.numPending()).isEqualTo(0)
-        restarter.restartSystemUI("Restart for test")
-        restarter.restartSystemUI("Restart for test")
-        assertThat(executor.numPending()).isEqualTo(1)
-    }
-
-    @Test
-    fun testWakefulnessLifecycle_CanRestart() {
-        whenever(wakefulnessLifecycle.wakefulness).thenReturn(WAKEFULNESS_AWAKE)
-        whenever(batteryController.isPluggedIn).thenReturn(true)
-        assertThat(executor.numPending()).isEqualTo(0)
-        restarter.restartSystemUI("Restart for test")
-
-        val captor = ArgumentCaptor.forClass(WakefulnessLifecycle.Observer::class.java)
-        verify(wakefulnessLifecycle).addObserver(captor.capture())
-
-        whenever(wakefulnessLifecycle.wakefulness).thenReturn(WAKEFULNESS_ASLEEP)
-
-        captor.value.onFinishedGoingToSleep()
-        assertThat(executor.numPending()).isEqualTo(1)
-    }
-
-    @Test
-    fun testBatteryController_CanRestart() {
-        whenever(wakefulnessLifecycle.wakefulness).thenReturn(WAKEFULNESS_ASLEEP)
-        whenever(batteryController.isPluggedIn).thenReturn(false)
-        assertThat(executor.numPending()).isEqualTo(0)
-        restarter.restartSystemUI("Restart for test")
-
-        val captor =
-            ArgumentCaptor.forClass(BatteryController.BatteryStateChangeCallback::class.java)
-        verify(batteryController).addCallback(captor.capture())
-
-        whenever(batteryController.isPluggedIn).thenReturn(true)
-
-        captor.value.onBatteryLevelChanged(0, true, true)
-        assertThat(executor.numPending()).isEqualTo(1)
-    }
-}

packages/SystemUI/tests/src/com/android/systemui/flags/PluggedInConditionTest.kt

@@ -0,0 +1,76 @@
+/*
+ * 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.statusbar.policy.BatteryController
+import com.google.common.truth.Truth.assertThat
+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
+
+/**
+ * Be careful with the {FeatureFlagsReleaseRestarter} in this test. It has a call to System.exit()!
+ */
+@SmallTest
+class PluggedInConditionTest : SysuiTestCase() {
+    private lateinit var condition: PluggedInCondition
+
+    @Mock private lateinit var batteryController: BatteryController
+
+    @Before
+    fun setup() {
+        MockitoAnnotations.initMocks(this)
+        condition = PluggedInCondition(batteryController)
+    }
+
+    @Test
+    fun testCondition_unplugged() {
+        whenever(batteryController.isPluggedIn).thenReturn(false)
+
+        assertThat(condition.canRestartNow({})).isFalse()
+    }
+
+    @Test
+    fun testCondition_pluggedIn() {
+        whenever(batteryController.isPluggedIn).thenReturn(true)
+
+        assertThat(condition.canRestartNow({})).isTrue()
+    }
+
+    @Test
+    fun testCondition_invokesRetry() {
+        whenever(batteryController.isPluggedIn).thenReturn(false)
+        var retried = false
+        val retryFn = { retried = true }
+
+        // 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())
+
+        whenever(batteryController.isPluggedIn).thenReturn(true)
+
+        captor.value.onBatteryLevelChanged(0, true, true)
+        assertThat(retried).isTrue()
+    }
+}

packages/SystemUI/tests/src/com/android/systemui/flags/ScreenIdleConditionTest.kt

@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2021 The Android Open Source Project
+ * 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.
@@ -20,12 +20,11 @@ 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.util.mockito.any
+import com.google.common.truth.Truth.assertThat
 import org.junit.Before
 import org.junit.Test
 import org.mockito.ArgumentCaptor
 import org.mockito.Mock
-import org.mockito.Mockito.never
 import org.mockito.Mockito.verify
 import org.mockito.Mockito.`when` as whenever
 import org.mockito.MockitoAnnotations
@@ -34,37 +33,45 @@ import org.mockito.MockitoAnnotations
  * Be careful with the {FeatureFlagsReleaseRestarter} in this test. It has a call to System.exit()!
  */
 @SmallTest
-class FeatureFlagsDebugRestarterTest : SysuiTestCase() {
-    private lateinit var restarter: FeatureFlagsDebugRestarter
+class ScreenIdleConditionTest : SysuiTestCase() {
+    private lateinit var condition: ScreenIdleCondition
 
     @Mock private lateinit var wakefulnessLifecycle: WakefulnessLifecycle
-    @Mock private lateinit var systemExitRestarter: SystemExitRestarter
 
     @Before
     fun setup() {
         MockitoAnnotations.initMocks(this)
-        restarter = FeatureFlagsDebugRestarter(wakefulnessLifecycle, systemExitRestarter)
+        condition = ScreenIdleCondition(wakefulnessLifecycle)
     }
 
     @Test
-    fun testRestart_ImmediateWhenAsleep() {
-        whenever(wakefulnessLifecycle.wakefulness).thenReturn(WAKEFULNESS_ASLEEP)
-        restarter.restartSystemUI("Restart for test")
-        verify(systemExitRestarter).restartSystemUI(any())
-    }
-
-    @Test
-    fun testRestart_WaitsForSceenOff() {
+    fun testCondition_awake() {
         whenever(wakefulnessLifecycle.wakefulness).thenReturn(WAKEFULNESS_AWAKE)
 
-        restarter.restartSystemUI("Restart for test")
-        verify(systemExitRestarter, never()).restartSystemUI(any())
+        assertThat(condition.canRestartNow {}).isFalse()
+    }
 
+    @Test
+    fun testCondition_asleep() {
+        whenever(wakefulnessLifecycle.wakefulness).thenReturn(WAKEFULNESS_ASLEEP)
+
+        assertThat(condition.canRestartNow {}).isTrue()
+    }
+
+    @Test
+    fun testCondition_invokesRetry() {
+        whenever(wakefulnessLifecycle.wakefulness).thenReturn(WAKEFULNESS_AWAKE)
+        var retried = false
+        val retryFn = { retried = true }
+
+        // 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())
 
-        captor.value.onFinishedGoingToSleep()
+        whenever(wakefulnessLifecycle.wakefulness).thenReturn(WAKEFULNESS_ASLEEP)
 
-        verify(systemExitRestarter).restartSystemUI(any())
+        captor.value.onFinishedGoingToSleep()
+        assertThat(retried).isTrue()
     }
 }