Restore AOD after SysUI Restart.

If the device was sleeping when SystemUI restarts, it would lose its
connection to loze its connection to the DreamService.

With this cl, we store when the device was dreaming and attempt to
reconnect when sysui restarts. This happens particularly when Flags
change.

Fixes: 268060399
Test: manual
Change-Id: Ie8583a6bbf9522d3a0fe2e5d3cfd94c43cc48f8d
diff --git a/packages/SystemUI/src/com/android/systemui/flags/FeatureFlagsDebugStartable.kt b/packages/SystemUI/src/com/android/systemui/flags/FeatureFlagsDebugStartable.kt
index dc7fc28..06ca0ad 100644
--- a/packages/SystemUI/src/com/android/systemui/flags/FeatureFlagsDebugStartable.kt
+++ b/packages/SystemUI/src/com/android/systemui/flags/FeatureFlagsDebugStartable.kt
@@ -22,6 +22,7 @@
 import com.android.systemui.dump.DumpManager
 import com.android.systemui.statusbar.commandline.CommandRegistry
 import com.android.systemui.util.InitializationChecker
+import com.android.systemui.util.concurrency.DelayableExecutor
 import dagger.Binds
 import dagger.Module
 import dagger.multibindings.ClassKey
@@ -36,7 +37,9 @@
     private val flagCommand: FlagCommand,
     private val featureFlags: FeatureFlagsDebug,
     private val broadcastSender: BroadcastSender,
-    private val initializationChecker: InitializationChecker
+    private val initializationChecker: InitializationChecker,
+    private val restartDozeListener: RestartDozeListener,
+    private val delayableExecutor: DelayableExecutor
 ) : CoreStartable {
 
     init {
@@ -52,6 +55,9 @@
             // protected broadcast should only be sent for the main process
             val intent = Intent(FlagManager.ACTION_SYSUI_STARTED)
             broadcastSender.sendBroadcast(intent)
+
+            restartDozeListener.init()
+            delayableExecutor.executeDelayed({ restartDozeListener.maybeRestartSleep() }, 1000)
         }
     }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/flags/FeatureFlagsReleaseStartable.kt b/packages/SystemUI/src/com/android/systemui/flags/FeatureFlagsReleaseStartable.kt
index d088d74..133e67f 100644
--- a/packages/SystemUI/src/com/android/systemui/flags/FeatureFlagsReleaseStartable.kt
+++ b/packages/SystemUI/src/com/android/systemui/flags/FeatureFlagsReleaseStartable.kt
@@ -18,6 +18,8 @@
 
 import com.android.systemui.CoreStartable
 import com.android.systemui.dump.DumpManager
+import com.android.systemui.util.InitializationChecker
+import com.android.systemui.util.concurrency.DelayableExecutor
 import dagger.Binds
 import dagger.Module
 import dagger.multibindings.ClassKey
@@ -26,7 +28,13 @@
 
 class FeatureFlagsReleaseStartable
 @Inject
-constructor(dumpManager: DumpManager, featureFlags: FeatureFlags) : CoreStartable {
+constructor(
+    dumpManager: DumpManager,
+    featureFlags: FeatureFlags,
+    private val initializationChecker: InitializationChecker,
+    private val restartDozeListener: RestartDozeListener,
+    private val delayableExecutor: DelayableExecutor
+) : CoreStartable {
 
     init {
         dumpManager.registerCriticalDumpable(FeatureFlagsRelease.TAG) { pw, args ->
@@ -35,7 +43,10 @@
     }
 
     override fun start() {
-        // no-op
+        if (initializationChecker.initializeComponents()) {
+            restartDozeListener.init()
+            delayableExecutor.executeDelayed({ restartDozeListener.maybeRestartSleep() }, 1000)
+        }
     }
 }
 
diff --git a/packages/SystemUI/src/com/android/systemui/flags/RestartDozeListener.kt b/packages/SystemUI/src/com/android/systemui/flags/RestartDozeListener.kt
new file mode 100644
index 0000000..bd74f4e
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/flags/RestartDozeListener.kt
@@ -0,0 +1,72 @@
+/*
+ * 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.os.PowerManager
+import android.util.Log
+import com.android.internal.annotations.VisibleForTesting
+import com.android.systemui.dagger.SysUISingleton
+import com.android.systemui.plugins.statusbar.StatusBarStateController
+import com.android.systemui.util.settings.SecureSettings
+import com.android.systemui.util.time.SystemClock
+import javax.inject.Inject
+
+@SysUISingleton
+class RestartDozeListener
+@Inject
+constructor(
+    private val settings: SecureSettings,
+    private val statusBarStateController: StatusBarStateController,
+    private val powerManager: PowerManager,
+    private val systemClock: SystemClock,
+) {
+
+    companion object {
+        @VisibleForTesting val RESTART_NAP_KEY = "restart_nap_after_start"
+    }
+
+    private var inited = false
+
+    val listener =
+        object : StatusBarStateController.StateListener {
+            override fun onDreamingChanged(isDreaming: Boolean) {
+                settings.putBool(RESTART_NAP_KEY, isDreaming)
+            }
+        }
+
+    fun init() {
+        if (inited) {
+            return
+        }
+        inited = true
+
+        statusBarStateController.addCallback(listener)
+    }
+
+    fun destroy() {
+        statusBarStateController.removeCallback(listener)
+    }
+
+    fun maybeRestartSleep() {
+        if (settings.getBool(RESTART_NAP_KEY, false)) {
+            Log.d("RestartDozeListener", "Restarting sleep state")
+            powerManager.wakeUp(systemClock.uptimeMillis())
+            powerManager.goToSleep(systemClock.uptimeMillis())
+            settings.putBool(RESTART_NAP_KEY, false)
+        }
+    }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/flags/ServerFlagReader.kt b/packages/SystemUI/src/com/android/systemui/flags/ServerFlagReader.kt
index e225b10..7acd3f3 100644
--- a/packages/SystemUI/src/com/android/systemui/flags/ServerFlagReader.kt
+++ b/packages/SystemUI/src/com/android/systemui/flags/ServerFlagReader.kt
@@ -62,11 +62,10 @@
                 return
             }
 
-
             for ((listener, flags) in listeners) {
                 propLoop@ for (propName in properties.keyset) {
                     for (flag in flags) {
-                        if (propName == getServerOverrideName(flag.id) || propName == flag.name) {
+                        if (propName == flag.name) {
                             listener.onChange(flag)
                             break@propLoop
                         }
@@ -103,10 +102,6 @@
         }
         listeners.add(Pair(listener, flags))
     }
-
-    private fun getServerOverrideName(flagId: Int): String {
-        return "flag_override_$flagId"
-    }
 }
 
 @Module
diff --git a/packages/SystemUI/tests/src/com/android/systemui/flags/RestartDozeListenerTest.kt b/packages/SystemUI/tests/src/com/android/systemui/flags/RestartDozeListenerTest.kt
new file mode 100644
index 0000000..de0e511
--- /dev/null
+++ b/packages/SystemUI/tests/src/com/android/systemui/flags/RestartDozeListenerTest.kt
@@ -0,0 +1,85 @@
+/*
+ * 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.os.PowerManager
+import android.test.suitebuilder.annotation.SmallTest
+import com.android.systemui.SysuiTestCase
+import com.android.systemui.plugins.statusbar.StatusBarStateController
+import com.android.systemui.util.settings.FakeSettings
+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.anyLong
+import org.mockito.Mockito.never
+import org.mockito.Mockito.verify
+import org.mockito.MockitoAnnotations
+
+@SmallTest
+class RestartDozeListenerTest : SysuiTestCase() {
+
+    lateinit var restartDozeListener: RestartDozeListener
+
+    val settings = FakeSettings()
+    @Mock lateinit var statusBarStateController: StatusBarStateController
+    @Mock lateinit var powerManager: PowerManager
+    val clock = FakeSystemClock()
+    lateinit var listener: StatusBarStateController.StateListener
+
+    @Before
+    fun setup() {
+        MockitoAnnotations.initMocks(this)
+        restartDozeListener =
+            RestartDozeListener(settings, statusBarStateController, powerManager, clock)
+
+        val captor = ArgumentCaptor.forClass(StatusBarStateController.StateListener::class.java)
+        restartDozeListener.init()
+        verify(statusBarStateController).addCallback(captor.capture())
+        listener = captor.value
+    }
+
+    @Test
+    fun testStoreDreamState_onDreamingStarted() {
+        listener.onDreamingChanged(true)
+        assertThat(settings.getBool(RestartDozeListener.RESTART_NAP_KEY)).isTrue()
+    }
+
+    @Test
+    fun testStoreDreamState_onDreamingStopped() {
+        listener.onDreamingChanged(false)
+        assertThat(settings.getBool(RestartDozeListener.RESTART_NAP_KEY)).isFalse()
+    }
+
+    @Test
+    fun testRestoreDreamState_dreamingShouldStart() {
+        settings.putBool(RestartDozeListener.RESTART_NAP_KEY, true)
+        restartDozeListener.maybeRestartSleep()
+        verify(powerManager).wakeUp(clock.uptimeMillis())
+        verify(powerManager).goToSleep(clock.uptimeMillis())
+    }
+
+    @Test
+    fun testRestoreDreamState_dreamingShouldNot() {
+        settings.putBool(RestartDozeListener.RESTART_NAP_KEY, false)
+        restartDozeListener.maybeRestartSleep()
+        verify(powerManager, never()).wakeUp(anyLong())
+        verify(powerManager, never()).goToSleep(anyLong())
+    }
+}
diff --git a/packages/SystemUI/tests/src/com/android/systemui/flags/ServerFlagReaderImplTest.kt b/packages/SystemUI/tests/src/com/android/systemui/flags/ServerFlagReaderImplTest.kt
index 4ebf974..a12315b 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/flags/ServerFlagReaderImplTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/flags/ServerFlagReaderImplTest.kt
@@ -50,10 +50,10 @@
 
     @Test
     fun testChange_alertsListener() {
-        val flag = ReleasedFlag(1, "1", "test")
+        val flag = ReleasedFlag(1, "flag_1", "test")
         serverFlagReader.listenForChanges(listOf(flag), changeListener)
 
-        deviceConfig.setProperty(NAMESPACE, "flag_override_1", "1", false)
+        deviceConfig.setProperty(NAMESPACE, "flag_1", "1", false)
         executor.runAllReady()
 
         verify(changeListener).onChange(flag)