Rework isAbleToDream signal

Still seeing bugs where the device comes out of AOD, goes
to LOCKSCREEN, but then back to DREAMING because SysUI
is not told fast enough that the dream ended. Add a very
large delay when leaving dozing before processing any
dreaming information. This should be OK as AOD->DREAMING
isn't really supported anyways

Test: atest KeyguardInteractorTest
Fixes: 356356505
Flag: com.android.systemui.migrate_clocks_to_blueprint
Change-Id: I6733a2dad73389d096bb3032e42b45155a0a15a2

9 files changed, 157 insertions, 61 deletions

diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/domain/interactor/CommunalSceneTransitionInteractorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/domain/interactor/CommunalSceneTransitionInteractorTest.kt
index ad7385344fac..d6712f09cd4e 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/domain/interactor/CommunalSceneTransitionInteractorTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/domain/interactor/CommunalSceneTransitionInteractorTest.kt
@@ -225,7 +225,7 @@ class CommunalSceneTransitionInteractorTest : SysuiTestCase() {
             kosmos.fakeKeyguardRepository.setKeyguardOccluded(true)
             kosmos.fakeKeyguardRepository.setDreaming(true)
             kosmos.fakeKeyguardRepository.setDreamingWithOverlay(true)
-            advanceTimeBy(100L)
+            advanceTimeBy(600L)
 
             sceneTransitions.value = hubToBlank
 
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/view/viewmodel/CommunalViewModelTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/view/viewmodel/CommunalViewModelTest.kt
index 12552489496d..cc945d63e15f 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/view/viewmodel/CommunalViewModelTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/view/viewmodel/CommunalViewModelTest.kt
@@ -533,6 +533,8 @@ class CommunalViewModelTest(flags: FlagsParameterization) : SysuiTestCase() {
             keyguardRepository.setDozeTransitionModel(
                 DozeTransitionModel(from = DozeStateModel.DOZE, to = DozeStateModel.FINISH)
             )
+            advanceTimeBy(600L)
+
             keyguardRepository.setDreaming(true)
             keyguardRepository.setDreamingWithOverlay(true)
             advanceTimeBy(60L)
@@ -641,6 +643,7 @@ class CommunalViewModelTest(flags: FlagsParameterization) : SysuiTestCase() {
             keyguardRepository.setDozeTransitionModel(
                 DozeTransitionModel(from = DozeStateModel.DOZE, to = DozeStateModel.FINISH)
             )
+            advanceTimeBy(600L)
             keyguardRepository.setDreaming(true)
             keyguardRepository.setDreamingWithOverlay(true)
             advanceTimeBy(60L)
@@ -699,6 +702,7 @@ class CommunalViewModelTest(flags: FlagsParameterization) : SysuiTestCase() {
             keyguardRepository.setDozeTransitionModel(
                 DozeTransitionModel(from = DozeStateModel.DOZE, to = DozeStateModel.FINISH)
             )
+            advanceTimeBy(600L)
             keyguardRepository.setDreaming(true)
             keyguardRepository.setDreamingWithOverlay(true)
             advanceTimeBy(60L)
diff --git a/packages/SystemUI/tests/src/com/android/systemui/keyguard/domain/interactor/FromDreamingTransitionInteractorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/domain/interactor/FromDreamingTransitionInteractorTest.kt
index 032794c43f08..dc225a399250 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/keyguard/domain/interactor/FromDreamingTransitionInteractorTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/domain/interactor/FromDreamingTransitionInteractorTest.kt
@@ -14,24 +14,10 @@
  * limitations under the License.
  */
 
-/*
- * Copyright (C) 2024 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 com.android.systemui.keyguard.shared.model.DozeStateModel.Companion.isDozeOff
+import com.android.systemui.coroutines.collectLastValue
 import android.platform.test.annotations.EnableFlags
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
@@ -42,24 +28,28 @@ import com.android.systemui.keyguard.data.repository.FakeKeyguardTransitionRepos
 import com.android.systemui.keyguard.data.repository.fakeKeyguardRepository
 import com.android.systemui.keyguard.data.repository.fakeKeyguardTransitionRepository
 import com.android.systemui.keyguard.data.repository.keyguardOcclusionRepository
+import com.android.systemui.keyguard.domain.interactor.keyguardInteractor
 import com.android.systemui.keyguard.shared.model.BiometricUnlockMode
 import com.android.systemui.keyguard.shared.model.KeyguardState
 import com.android.systemui.keyguard.util.KeyguardTransitionRepositorySpySubject.Companion.assertThat
 import com.android.systemui.kosmos.testScope
 import com.android.systemui.power.domain.interactor.PowerInteractor.Companion.setAsleepForTest
+import com.android.systemui.power.domain.interactor.PowerInteractor.Companion.setAwakeForTest
 import com.android.systemui.power.domain.interactor.powerInteractor
 import com.android.systemui.statusbar.domain.interactor.keyguardOcclusionInteractor
 import com.android.systemui.testKosmos
-import kotlin.test.Test
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.runBlocking
+import kotlinx.coroutines.test.advanceTimeBy
 import kotlinx.coroutines.test.runCurrent
 import kotlinx.coroutines.test.runTest
 import org.junit.Before
 import org.junit.Ignore
+import org.junit.Test
 import org.junit.runner.RunWith
 import org.mockito.Mockito.reset
 import org.mockito.Mockito.spy
+import com.google.common.truth.Truth.assertThat
 
 @OptIn(ExperimentalCoroutinesApi::class)
 @SmallTest
@@ -79,21 +69,6 @@ class FromDreamingTransitionInteractorTest : SysuiTestCase() {
     @Before
     fun setup() {
         underTest.start()
-
-        kosmos.fakeKeyguardRepository.setDreaming(true)
-        kosmos.keyguardOcclusionInteractor.setWmNotifiedShowWhenLockedActivityOnTop(true)
-
-        // Transition to DOZING and set the power interactor asleep.
-        powerInteractor.setAsleepForTest()
-        runBlocking {
-            transitionRepository.sendTransitionSteps(
-                from = KeyguardState.LOCKSCREEN,
-                to = KeyguardState.DREAMING,
-                testScope
-            )
-            kosmos.fakeKeyguardRepository.setBiometricUnlockState(BiometricUnlockMode.NONE)
-            reset(transitionRepository)
-        }
     }
 
     @Test
@@ -146,20 +121,27 @@ class FromDreamingTransitionInteractorTest : SysuiTestCase() {
     @EnableFlags(Flags.FLAG_KEYGUARD_WM_STATE_REFACTOR)
     fun testTransitionsToLockscreen_whenOccludingActivityEnds() =
         testScope.runTest {
+            runCurrent()
+
             kosmos.fakeKeyguardRepository.setDreaming(true)
-            kosmos.keyguardOcclusionRepository.setShowWhenLockedActivityInfo(onTop = true)
+            kosmos.keyguardOcclusionInteractor.setWmNotifiedShowWhenLockedActivityOnTop(true)
+            // Transition to DREAMING and set the power interactor awake
+            powerInteractor.setAwakeForTest()
+
             transitionRepository.sendTransitionSteps(
                 from = KeyguardState.LOCKSCREEN,
                 to = KeyguardState.DREAMING,
-                testScope,
+                testScope
             )
-            runCurrent()
+            kosmos.fakeKeyguardRepository.setBiometricUnlockState(BiometricUnlockMode.NONE)
 
+            // Get past initial setup
+            advanceTimeBy(600L)
             reset(transitionRepository)
 
             kosmos.keyguardOcclusionRepository.setShowWhenLockedActivityInfo(onTop = false)
             kosmos.fakeKeyguardRepository.setDreaming(false)
-            runCurrent()
+            advanceTimeBy(60L)
 
             assertThat(transitionRepository)
                 .startedTransition(
@@ -171,6 +153,13 @@ class FromDreamingTransitionInteractorTest : SysuiTestCase() {
     @Test
     fun testTransitionToAlternateBouncer() =
         testScope.runTest {
+            transitionRepository.sendTransitionSteps(
+                from = KeyguardState.LOCKSCREEN,
+                to = KeyguardState.DREAMING,
+                testScope,
+            )
+            reset(transitionRepository)
+
             kosmos.fakeKeyguardBouncerRepository.setAlternateVisible(true)
             runCurrent()
 
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/domain/interactor/KeyguardInteractorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/domain/interactor/KeyguardInteractorTest.kt
index fc827a1478c7..ebefb4d51943 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/domain/interactor/KeyguardInteractorTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/domain/interactor/KeyguardInteractorTest.kt
@@ -33,11 +33,15 @@ import com.android.systemui.keyguard.data.repository.fakeCommandQueue
 import com.android.systemui.keyguard.data.repository.fakeKeyguardRepository
 import com.android.systemui.keyguard.data.repository.fakeKeyguardTransitionRepository
 import com.android.systemui.keyguard.shared.model.CameraLaunchType
+import com.android.systemui.keyguard.shared.model.DozeStateModel
+import com.android.systemui.keyguard.shared.model.DozeTransitionModel
 import com.android.systemui.keyguard.shared.model.KeyguardState
 import com.android.systemui.keyguard.shared.model.StatusBarState
 import com.android.systemui.keyguard.shared.model.TransitionState
 import com.android.systemui.keyguard.shared.model.TransitionStep
 import com.android.systemui.kosmos.testScope
+import com.android.systemui.power.domain.interactor.PowerInteractor.Companion.setAwakeForTest
+import com.android.systemui.power.domain.interactor.powerInteractor
 import com.android.systemui.res.R
 import com.android.systemui.scene.domain.interactor.sceneInteractor
 import com.android.systemui.scene.shared.model.Scenes
@@ -47,6 +51,7 @@ import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.flowOf
+import kotlinx.coroutines.test.advanceTimeBy
 import kotlinx.coroutines.test.runCurrent
 import kotlinx.coroutines.test.runTest
 import org.junit.Before
@@ -67,6 +72,7 @@ class KeyguardInteractorTest : SysuiTestCase() {
     private val configRepository by lazy { kosmos.fakeConfigurationRepository }
     private val bouncerRepository by lazy { kosmos.keyguardBouncerRepository }
     private val shadeRepository by lazy { kosmos.shadeRepository }
+    private val powerInteractor by lazy { kosmos.powerInteractor }
     private val keyguardTransitionRepository by lazy { kosmos.fakeKeyguardTransitionRepository }
 
     private val transitionState: MutableStateFlow<ObservableTransitionState> =
@@ -350,6 +356,59 @@ class KeyguardInteractorTest : SysuiTestCase() {
         }
 
     @Test
+    fun isAbleToDream_falseWhenDozing() =
+        testScope.runTest {
+            val isAbleToDream by collectLastValue(underTest.isAbleToDream)
+
+            repository.setDozeTransitionModel(
+                DozeTransitionModel(from = DozeStateModel.INITIALIZED, to = DozeStateModel.DOZE_AOD)
+            )
+
+            assertThat(isAbleToDream).isEqualTo(false)
+        }
+
+    @Test
+    fun isAbleToDream_falseWhenNotDozingAndNotDreaming() =
+        testScope.runTest {
+            val isAbleToDream by collectLastValue(underTest.isAbleToDream)
+
+            repository.setDozeTransitionModel(
+                DozeTransitionModel(from = DozeStateModel.DOZE, to = DozeStateModel.FINISH)
+            )
+            powerInteractor.setAwakeForTest()
+            advanceTimeBy(1000L)
+
+            assertThat(isAbleToDream).isEqualTo(false)
+        }
+
+    @Test
+    fun isAbleToDream_trueWhenNotDozingAndIsDreaming_afterDelay() =
+        testScope.runTest {
+            val isAbleToDream by collectLastValue(underTest.isAbleToDream)
+            runCurrent()
+
+            repository.setDreaming(true)
+            repository.setDozeTransitionModel(
+                DozeTransitionModel(from = DozeStateModel.DOZE, to = DozeStateModel.FINISH)
+            )
+            powerInteractor.setAwakeForTest()
+            runCurrent()
+
+            // After some delay, still false
+            advanceTimeBy(300L)
+            assertThat(isAbleToDream).isEqualTo(false)
+
+            // After more delay, is true
+            advanceTimeBy(300L)
+            assertThat(isAbleToDream).isEqualTo(true)
+
+            // Also changes back after the minimal debounce
+            repository.setDreaming(false)
+            advanceTimeBy(55L)
+            assertThat(isAbleToDream).isEqualTo(false)
+        }
+
+    @Test
     @EnableSceneContainer
     fun animationDozingTransitions() =
         testScope.runTest {
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/domain/interactor/KeyguardTransitionScenariosTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/domain/interactor/KeyguardTransitionScenariosTest.kt
index 9762fd8e2158..90e13a57cefe 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/domain/interactor/KeyguardTransitionScenariosTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/domain/interactor/KeyguardTransitionScenariosTest.kt
@@ -258,7 +258,7 @@ class KeyguardTransitionScenariosTest(flags: FlagsParameterization?) : SysuiTest
             keyguardRepository.setDozeTransitionModel(
                 DozeTransitionModel(from = DozeStateModel.DOZE, to = DozeStateModel.FINISH)
             )
-            runCurrent()
+            advanceTimeBy(600L)
 
             // GIVEN a prior transition has run to LOCKSCREEN
             runTransitionAndSetWakefulness(KeyguardState.GONE, KeyguardState.LOCKSCREEN)
@@ -287,7 +287,7 @@ class KeyguardTransitionScenariosTest(flags: FlagsParameterization?) : SysuiTest
             keyguardRepository.setDozeTransitionModel(
                 DozeTransitionModel(from = DozeStateModel.DOZE, to = DozeStateModel.FINISH)
             )
-            runCurrent()
+            advanceTimeBy(600L)
 
             // GIVEN a prior transition has run to LOCKSCREEN
             runTransitionAndSetWakefulness(KeyguardState.GONE, KeyguardState.LOCKSCREEN)
@@ -625,6 +625,9 @@ class KeyguardTransitionScenariosTest(flags: FlagsParameterization?) : SysuiTest
     @DisableSceneContainer
     fun dreamingToGoneWithKeyguardNotShowing() =
         testScope.runTest {
+            // Setup - Move past initial delay with [KeyguardInteractor#isAbleToDream]
+            advanceTimeBy(600L)
+
             // GIVEN a prior transition has run to DREAMING
             keyguardRepository.setDreamingWithOverlay(true)
             runTransitionAndSetWakefulness(KeyguardState.LOCKSCREEN, KeyguardState.DREAMING)
@@ -757,12 +760,8 @@ class KeyguardTransitionScenariosTest(flags: FlagsParameterization?) : SysuiTest
     @DisableSceneContainer
     fun goneToDreaming() =
         testScope.runTest {
-            // GIVEN a device that is not dreaming or dozing
-            keyguardRepository.setDreamingWithOverlay(false)
-            keyguardRepository.setDozeTransitionModel(
-                DozeTransitionModel(from = DozeStateModel.DOZE, to = DozeStateModel.FINISH)
-            )
-            runCurrent()
+            // Setup - Move past initial delay with [KeyguardInteractor#isAbleToDream]
+            advanceTimeBy(600L)
 
             // GIVEN a prior transition has run to GONE
             runTransitionAndSetWakefulness(KeyguardState.LOCKSCREEN, KeyguardState.GONE)
@@ -1130,6 +1129,9 @@ class KeyguardTransitionScenariosTest(flags: FlagsParameterization?) : SysuiTest
     @DisableSceneContainer
     fun primaryBouncerToGlanceableHubWhileDreaming() =
         testScope.runTest {
+            // Setup - Move past initial delay with [KeyguardInteractor#isAbleToDream]
+            advanceTimeBy(600L)
+
             // GIVEN the device is idle on the glanceable hub
             communalSceneInteractor.changeScene(CommunalScenes.Communal)
             runCurrent()
@@ -1144,6 +1146,7 @@ class KeyguardTransitionScenariosTest(flags: FlagsParameterization?) : SysuiTest
             // GIVEN that we are dreaming and occluded
             keyguardRepository.setDreaming(true)
             keyguardRepository.setKeyguardOccluded(true)
+            advanceTimeBy(60L)
 
             // WHEN the primaryBouncer stops showing
             bouncerRepository.setPrimaryShow(false)
@@ -2181,12 +2184,14 @@ class KeyguardTransitionScenariosTest(flags: FlagsParameterization?) : SysuiTest
     @DisableFlags(FLAG_COMMUNAL_SCENE_KTF_REFACTOR)
     fun glanceableHubToDreaming() =
         testScope.runTest {
+            runCurrent()
+
             // GIVEN that we are dreaming and not dozing
             keyguardRepository.setDreaming(true)
             keyguardRepository.setDozeTransitionModel(
                 DozeTransitionModel(from = DozeStateModel.DOZE, to = DozeStateModel.FINISH)
             )
-            runCurrent()
+            advanceTimeBy(600L)
 
             // GIVEN a prior transition has run to GLANCEABLE_HUB
             runTransitionAndSetWakefulness(KeyguardState.DREAMING, KeyguardState.GLANCEABLE_HUB)
@@ -2233,7 +2238,7 @@ class KeyguardTransitionScenariosTest(flags: FlagsParameterization?) : SysuiTest
             keyguardRepository.setDozeTransitionModel(
                 DozeTransitionModel(from = DozeStateModel.DOZE, to = DozeStateModel.FINISH)
             )
-            advanceTimeBy(100L)
+            advanceTimeBy(600L)
 
             // GIVEN a prior transition has run to GLANCEABLE_HUB
             communalSceneInteractor.changeScene(CommunalScenes.Communal)
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/FromDreamingTransitionInteractor.kt b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/FromDreamingTransitionInteractor.kt
index 3775d191949e..17c1e823a1ca 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/FromDreamingTransitionInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/FromDreamingTransitionInteractor.kt
@@ -133,7 +133,12 @@ constructor(
                 transitionInteractor.startedKeyguardState.replayCache.last() ==
                     KeyguardState.DREAMING
             ) {
-                startTransitionTo(KeyguardState.LOCKSCREEN)
+                if (powerInteractor.detailedWakefulness.value.isAwake()) {
+                    startTransitionTo(
+                        KeyguardState.LOCKSCREEN,
+                        ownerReason = "Dream has ended and device is awake"
+                    )
+                }
             }
         }
     }
@@ -144,7 +149,7 @@ constructor(
             scope.launch {
                 combine(
                         keyguardInteractor.isKeyguardOccluded,
-                        keyguardInteractor.isDreaming
+                        keyguardInteractor.isAbleToDream
                             // Debounce the dreaming signal since there is a race condition between
                             // the occluded and dreaming signals. We therefore add a small delay
                             // to give enough time for occluded to flip to false when the dream
@@ -172,7 +177,7 @@ constructor(
         }
 
         scope.launch {
-            keyguardInteractor.isDreaming
+            keyguardInteractor.isAbleToDream
                 .filter { !it }
                 .sample(deviceEntryInteractor.isUnlocked, ::Pair)
                 .collect { (_, dismissable) ->
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardInteractor.kt b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardInteractor.kt
index 42490c4176c6..69e10d9a6009 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardInteractor.kt
@@ -14,7 +14,6 @@
  *  limitations under the License.
  *
  */
-
 @file:OptIn(ExperimentalCoroutinesApi::class)
 
 package com.android.systemui.keyguard.domain.interactor
@@ -156,14 +155,23 @@ constructor(
 
     val isPulsing: Flow<Boolean> = dozeTransitionModel.map { it.to == DozeStateModel.DOZE_PULSING }
 
+    /** Whether the system is dreaming with an overlay active */
+    val isDreamingWithOverlay: Flow<Boolean> = repository.isDreamingWithOverlay
+
     /**
      * Whether the system is dreaming. [isDreaming] will be always be true when [isDozing] is true,
-     * but not vice-versa.
+     * but not vice-versa. Also accounts for [isDreamingWithOverlay]
      */
-    val isDreaming: StateFlow<Boolean> = repository.isDreaming
-
-    /** Whether the system is dreaming with an overlay active */
-    val isDreamingWithOverlay: Flow<Boolean> = repository.isDreamingWithOverlay
+    val isDreaming: StateFlow<Boolean> =
+        merge(
+                repository.isDreaming,
+                repository.isDreamingWithOverlay,
+            )
+            .stateIn(
+                scope = applicationScope,
+                started = SharingStarted.Eagerly,
+                initialValue = false,
+            )
 
     /** Whether the system is dreaming and the active dream is hosted in lockscreen */
     val isActiveDreamLockscreenHosted: StateFlow<Boolean> = repository.isActiveDreamLockscreenHosted
@@ -179,12 +187,25 @@ constructor(
      * Allow a brief moment to prevent rapidly oscillating between true/false signals.
      */
     val isAbleToDream: Flow<Boolean> =
-        merge(isDreaming, isDreamingWithOverlay)
-            .combine(dozeTransitionModel) { isDreaming, dozeTransitionModel ->
-                isDreaming && isDozeOff(dozeTransitionModel.to)
+        dozeTransitionModel
+            .flatMapLatest { dozeTransitionModel ->
+                if (isDozeOff(dozeTransitionModel.to)) {
+                    // When dozing stops, it is a very early signal that the device is exiting the
+                    // dream state. DreamManagerService eventually notifies window manager, which
+                    // invokes SystemUI through KeyguardService. Because of this substantial delay,
+                    // do not immediately process any dreaming information when exiting AOD. It
+                    // should actually be quite strange to leave AOD and then go straight to
+                    // DREAMING so this should be fine.
+                    delay(500L)
+                    isDreaming
+                        .sample(powerInteractor.isAwake) { isDreaming, isAwake ->
+                            isDreaming && isAwake
+                        }
+                        .debounce(50L)
+                } else {
+                    flowOf(false)
+                }
             }
-            .sample(powerInteractor.isAwake) { isAbleToDream, isAwake -> isAbleToDream && isAwake }
-            .debounce(50L)
             .stateIn(
                 scope = applicationScope,
                 started = SharingStarted.WhileSubscribed(),
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardTransitionAuditLogger.kt b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardTransitionAuditLogger.kt
index e132eb7ec32a..b89eb2723fab 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardTransitionAuditLogger.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardTransitionAuditLogger.kt
@@ -98,6 +98,16 @@ constructor(
         }
 
         scope.launch {
+            keyguardInteractor.isDreaming.collect { logger.log(TAG, VERBOSE, "isDreaming", it) }
+        }
+
+        scope.launch {
+            keyguardInteractor.isDreamingWithOverlay.collect {
+                logger.log(TAG, VERBOSE, "isDreamingWithOverlay", it)
+            }
+        }
+
+        scope.launch {
             keyguardInteractor.isAbleToDream.collect {
                 logger.log(TAG, VERBOSE, "isAbleToDream", it)
             }
diff --git 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
index 727de9e95872..a1c2f79536c4 100644
--- 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
@@ -216,6 +216,9 @@ class FakeKeyguardRepository @Inject constructor() : KeyguardRepository {
 
     override fun setDreaming(isDreaming: Boolean) {
         _isDreaming.value = isDreaming
+        // Intentionally set both for testing, to avoid races with merge() in the interactor that
+        // would make testing difficult
+        _isDreamingWithOverlay.value = isDreaming
     }
 
     fun setDreamingWithOverlay(isDreaming: Boolean) {