2 files changed, 26 insertions, 18 deletions

diff --git a/packages/SystemUI/src/com/android/systemui/flags/RestartDozeListener.kt b/packages/SystemUI/src/com/android/systemui/flags/RestartDozeListener.kt
index b49d60dbdde1..dc0de2cef349 100644
--- a/packages/SystemUI/src/com/android/systemui/flags/RestartDozeListener.kt
+++ b/packages/SystemUI/src/com/android/systemui/flags/RestartDozeListener.kt
@@ -39,15 +39,15 @@ constructor(
 ) {
 
     companion object {
-        @VisibleForTesting val RESTART_NAP_KEY = "restart_nap_after_start"
+        @VisibleForTesting val RESTART_SLEEP_KEY = "restart_nap_after_start"
     }
 
     private var inited = false
 
     val listener =
         object : StatusBarStateController.StateListener {
-            override fun onDreamingChanged(isDreaming: Boolean) {
-                storeSleepState(isDreaming)
+            override fun onDozingChanged(isDozing: Boolean) {
+                storeSleepState(isDozing)
             }
         }
 
@@ -67,9 +67,13 @@ constructor(
     fun maybeRestartSleep() {
         bgExecutor.executeDelayed(
             {
-                if (settings.getBool(RESTART_NAP_KEY, false)) {
+                if (settings.getBool(RESTART_SLEEP_KEY, false)) {
                     Log.d("RestartDozeListener", "Restarting sleep state")
-                    powerManager.wakeUp(systemClock.uptimeMillis())
+                    powerManager.wakeUp(
+                        systemClock.uptimeMillis(),
+                        PowerManager.WAKE_REASON_APPLICATION,
+                        "RestartDozeListener"
+                    )
                     powerManager.goToSleep(systemClock.uptimeMillis())
                 }
             },
@@ -78,6 +82,6 @@ constructor(
     }
 
     private fun storeSleepState(sleeping: Boolean) {
-        bgExecutor.execute { settings.putBool(RESTART_NAP_KEY, sleeping) }
+        bgExecutor.execute { settings.putBool(RESTART_SLEEP_KEY, sleeping) }
     }
 }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/flags/RestartDozeListenerTest.kt b/packages/SystemUI/tests/src/com/android/systemui/flags/RestartDozeListenerTest.kt
index 2db4596c80a7..e287f19b2455 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/flags/RestartDozeListenerTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/flags/RestartDozeListenerTest.kt
@@ -21,12 +21,15 @@ import android.test.suitebuilder.annotation.SmallTest
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.plugins.statusbar.StatusBarStateController
 import com.android.systemui.util.concurrency.FakeExecutor
+import com.android.systemui.util.mockito.eq
 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.ArgumentMatchers.anyInt
+import org.mockito.ArgumentMatchers.anyString
 import org.mockito.Mock
 import org.mockito.Mockito.anyLong
 import org.mockito.Mockito.never
@@ -58,36 +61,37 @@ class RestartDozeListenerTest : SysuiTestCase() {
     }
 
     @Test
-    fun testStoreDreamState_onDreamingStarted() {
-        listener.onDreamingChanged(true)
+    fun testStoreDreamState_onDozingStarted() {
+        listener.onDozingChanged(true)
         executor.runAllReady()
-        assertThat(settings.getBool(RestartDozeListener.RESTART_NAP_KEY)).isTrue()
+        assertThat(settings.getBool(RestartDozeListener.RESTART_SLEEP_KEY)).isTrue()
     }
 
     @Test
-    fun testStoreDreamState_onDreamingStopped() {
-        listener.onDreamingChanged(false)
+    fun testStoreDozeState_onDozingStopped() {
+        listener.onDozingChanged(false)
         executor.runAllReady()
-        assertThat(settings.getBool(RestartDozeListener.RESTART_NAP_KEY)).isFalse()
+        assertThat(settings.getBool(RestartDozeListener.RESTART_SLEEP_KEY)).isFalse()
     }
 
     @Test
-    fun testRestoreDreamState_dreamingShouldStart() {
-        settings.putBool(RestartDozeListener.RESTART_NAP_KEY, true)
+    fun testRestoreDozeState_dozingShouldStart() {
+        settings.putBool(RestartDozeListener.RESTART_SLEEP_KEY, true)
         restartDozeListener.maybeRestartSleep()
         executor.advanceClockToLast()
         executor.runAllReady()
-        verify(powerManager).wakeUp(clock.uptimeMillis())
+        verify(powerManager)
+            .wakeUp(eq(clock.uptimeMillis()), eq(PowerManager.WAKE_REASON_APPLICATION), anyString())
         verify(powerManager).goToSleep(clock.uptimeMillis())
     }
 
     @Test
-    fun testRestoreDreamState_dreamingShouldNot() {
-        settings.putBool(RestartDozeListener.RESTART_NAP_KEY, false)
+    fun testRestoreDozeState_dozingShouldNotStart() {
+        settings.putBool(RestartDozeListener.RESTART_SLEEP_KEY, false)
         restartDozeListener.maybeRestartSleep()
         executor.advanceClockToLast()
         executor.runAllReady()
-        verify(powerManager, never()).wakeUp(anyLong())
+        verify(powerManager, never()).wakeUp(anyLong(), anyInt(), anyString())
         verify(powerManager, never()).goToSleep(anyLong())
     }
 }