Merge "[flexiglass] Resets the scene when the device goes to sleep." into udc-qpr-dev
diff --git a/packages/SystemUI/src/com/android/systemui/authentication/data/repository/AuthenticationRepository.kt b/packages/SystemUI/src/com/android/systemui/authentication/data/repository/AuthenticationRepository.kt
index 43fb363..444491f 100644
--- a/packages/SystemUI/src/com/android/systemui/authentication/data/repository/AuthenticationRepository.kt
+++ b/packages/SystemUI/src/com/android/systemui/authentication/data/repository/AuthenticationRepository.kt
@@ -144,12 +144,7 @@
     private val lockPatternUtils: LockPatternUtils,
 ) : AuthenticationRepository {
 
-    override val isUnlocked: StateFlow<Boolean> =
-        keyguardRepository.isKeyguardUnlocked.stateIn(
-            scope = applicationScope,
-            started = SharingStarted.WhileSubscribed(),
-            initialValue = false,
-        )
+    override val isUnlocked: StateFlow<Boolean> = keyguardRepository.isKeyguardUnlocked
 
     override suspend fun isLockscreenEnabled(): Boolean {
         return withContext(backgroundDispatcher) {
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/data/repository/KeyguardRepository.kt b/packages/SystemUI/src/com/android/systemui/keyguard/data/repository/KeyguardRepository.kt
index e7704d6..d119920 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/data/repository/KeyguardRepository.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/data/repository/KeyguardRepository.kt
@@ -87,7 +87,7 @@
     val isKeyguardShowing: Flow<Boolean>
 
     /** Is the keyguard in a unlocked state? */
-    val isKeyguardUnlocked: Flow<Boolean>
+    val isKeyguardUnlocked: StateFlow<Boolean>
 
     /** Is an activity showing over the keyguard? */
     val isKeyguardOccluded: Flow<Boolean>
@@ -299,7 +299,7 @@
             }
             .distinctUntilChanged()
 
-    override val isKeyguardUnlocked: Flow<Boolean> =
+    override val isKeyguardUnlocked: StateFlow<Boolean> =
         conflatedCallbackFlow {
                 val callback =
                     object : KeyguardStateController.Callback {
@@ -330,7 +330,11 @@
 
                 awaitClose { keyguardStateController.removeCallback(callback) }
             }
-            .distinctUntilChanged()
+            .stateIn(
+                scope = scope,
+                started = SharingStarted.WhileSubscribed(),
+                initialValue = keyguardStateController.isUnlocked,
+            )
 
     override val isKeyguardGoingAway: Flow<Boolean> = conflatedCallbackFlow {
         val callback =
diff --git a/packages/SystemUI/src/com/android/systemui/scene/domain/startable/SystemUiDefaultSceneContainerStartable.kt b/packages/SystemUI/src/com/android/systemui/scene/domain/startable/SystemUiDefaultSceneContainerStartable.kt
index 59f82f0..285ff74 100644
--- a/packages/SystemUI/src/com/android/systemui/scene/domain/startable/SystemUiDefaultSceneContainerStartable.kt
+++ b/packages/SystemUI/src/com/android/systemui/scene/domain/startable/SystemUiDefaultSceneContainerStartable.kt
@@ -22,6 +22,8 @@
 import com.android.systemui.dagger.qualifiers.Application
 import com.android.systemui.flags.FeatureFlags
 import com.android.systemui.flags.Flags
+import com.android.systemui.keyguard.domain.interactor.KeyguardInteractor
+import com.android.systemui.keyguard.shared.model.WakefulnessState
 import com.android.systemui.scene.domain.interactor.SceneInteractor
 import com.android.systemui.scene.shared.model.SceneContainerNames
 import com.android.systemui.scene.shared.model.SceneKey
@@ -45,6 +47,7 @@
     @Application private val applicationScope: CoroutineScope,
     private val sceneInteractor: SceneInteractor,
     private val authenticationInteractor: AuthenticationInteractor,
+    private val keyguardInteractor: KeyguardInteractor,
     private val featureFlags: FeatureFlags,
 ) : CoreStartable {
 
@@ -78,7 +81,7 @@
                     when {
                         isUnlocked ->
                             when (currentSceneKey) {
-                                // When the device becomes unlocked in Bouncer, go to the Gone.
+                                // When the device becomes unlocked in Bouncer, go to Gone.
                                 is SceneKey.Bouncer -> SceneKey.Gone
                                 // When the device becomes unlocked in Lockscreen, go to Gone if
                                 // bypass is enabled.
@@ -101,15 +104,30 @@
                     }
                 }
                 .filterNotNull()
-                .collect { targetSceneKey ->
-                    sceneInteractor.setCurrentScene(
-                        containerName = CONTAINER_NAME,
-                        scene = SceneModel(targetSceneKey),
-                    )
+                .collect { targetSceneKey -> switchToScene(targetSceneKey) }
+        }
+
+        applicationScope.launch {
+            keyguardInteractor.wakefulnessModel
+                .map { it.state == WakefulnessState.ASLEEP }
+                .distinctUntilChanged()
+                .collect { isAsleep ->
+                    if (isAsleep) {
+                        // When the device goes to sleep, reset the current scene.
+                        val isUnlocked = authenticationInteractor.isUnlocked.value
+                        switchToScene(if (isUnlocked) SceneKey.Gone else SceneKey.Lockscreen)
+                    }
                 }
         }
     }
 
+    private fun switchToScene(targetSceneKey: SceneKey) {
+        sceneInteractor.setCurrentScene(
+            containerName = CONTAINER_NAME,
+            scene = SceneModel(targetSceneKey),
+        )
+    }
+
     companion object {
         private const val CONTAINER_NAME = SceneContainerNames.SYSTEM_UI_DEFAULT
     }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/scene/domain/startable/SystemUiDefaultSceneContainerStartableTest.kt b/packages/SystemUI/tests/src/com/android/systemui/scene/domain/startable/SystemUiDefaultSceneContainerStartableTest.kt
index df5e7bc..3e9ddcb 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/scene/domain/startable/SystemUiDefaultSceneContainerStartableTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/scene/domain/startable/SystemUiDefaultSceneContainerStartableTest.kt
@@ -20,6 +20,9 @@
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.coroutines.collectLastValue
 import com.android.systemui.flags.Flags
+import com.android.systemui.keyguard.shared.model.WakeSleepReason
+import com.android.systemui.keyguard.shared.model.WakefulnessModel
+import com.android.systemui.keyguard.shared.model.WakefulnessState
 import com.android.systemui.scene.SceneTestUtils
 import com.android.systemui.scene.shared.model.SceneContainerNames
 import com.android.systemui.scene.shared.model.SceneKey
@@ -47,12 +50,18 @@
         utils.authenticationInteractor(
             repository = authenticationRepository,
         )
+    private val keyguardRepository = utils.keyguardRepository()
+    private val keyguardInteractor =
+        utils.keyguardInteractor(
+            repository = keyguardRepository,
+        )
 
     private val underTest =
         SystemUiDefaultSceneContainerStartable(
             applicationScope = testScope.backgroundScope,
             sceneInteractor = sceneInteractor,
             authenticationInteractor = authenticationInteractor,
+            keyguardInteractor = keyguardInteractor,
             featureFlags = featureFlags,
         )
 
@@ -280,6 +289,94 @@
             assertThat(currentSceneKey).isEqualTo(SceneKey.Lockscreen)
         }
 
+    @Test
+    fun switchToGoneWhenDeviceSleepsUnlocked_featureEnabled() =
+        testScope.runTest {
+            val currentSceneKey by
+                collectLastValue(
+                    sceneInteractor.currentScene(SceneContainerNames.SYSTEM_UI_DEFAULT).map {
+                        it.key
+                    }
+                )
+            prepareState(
+                isFeatureEnabled = true,
+                isDeviceUnlocked = true,
+                initialSceneKey = SceneKey.Shade,
+            )
+            assertThat(currentSceneKey).isEqualTo(SceneKey.Shade)
+            underTest.start()
+
+            keyguardRepository.setWakefulnessModel(ASLEEP)
+
+            assertThat(currentSceneKey).isEqualTo(SceneKey.Gone)
+        }
+
+    @Test
+    fun switchToGoneWhenDeviceSleepsUnlocked_featureDisabled() =
+        testScope.runTest {
+            val currentSceneKey by
+                collectLastValue(
+                    sceneInteractor.currentScene(SceneContainerNames.SYSTEM_UI_DEFAULT).map {
+                        it.key
+                    }
+                )
+            prepareState(
+                isFeatureEnabled = false,
+                isDeviceUnlocked = true,
+                initialSceneKey = SceneKey.Shade,
+            )
+            assertThat(currentSceneKey).isEqualTo(SceneKey.Shade)
+            underTest.start()
+
+            keyguardRepository.setWakefulnessModel(ASLEEP)
+
+            assertThat(currentSceneKey).isEqualTo(SceneKey.Shade)
+        }
+
+    @Test
+    fun switchToLockscreenWhenDeviceSleepsLocked_featureEnabled() =
+        testScope.runTest {
+            val currentSceneKey by
+                collectLastValue(
+                    sceneInteractor.currentScene(SceneContainerNames.SYSTEM_UI_DEFAULT).map {
+                        it.key
+                    }
+                )
+            prepareState(
+                isFeatureEnabled = true,
+                isDeviceUnlocked = false,
+                initialSceneKey = SceneKey.Shade,
+            )
+            assertThat(currentSceneKey).isEqualTo(SceneKey.Shade)
+            underTest.start()
+
+            keyguardRepository.setWakefulnessModel(ASLEEP)
+
+            assertThat(currentSceneKey).isEqualTo(SceneKey.Lockscreen)
+        }
+
+    @Test
+    fun switchToLockscreenWhenDeviceSleepsLocked_featureDisabled() =
+        testScope.runTest {
+            val currentSceneKey by
+                collectLastValue(
+                    sceneInteractor.currentScene(SceneContainerNames.SYSTEM_UI_DEFAULT).map {
+                        it.key
+                    }
+                )
+            prepareState(
+                isFeatureEnabled = false,
+                isDeviceUnlocked = false,
+                initialSceneKey = SceneKey.Shade,
+            )
+            assertThat(currentSceneKey).isEqualTo(SceneKey.Shade)
+            underTest.start()
+
+            keyguardRepository.setWakefulnessModel(ASLEEP)
+
+            assertThat(currentSceneKey).isEqualTo(SceneKey.Shade)
+        }
+
     private fun prepareState(
         isFeatureEnabled: Boolean = true,
         isDeviceUnlocked: Boolean = false,
@@ -293,4 +390,13 @@
             sceneInteractor.setCurrentScene(SceneContainerNames.SYSTEM_UI_DEFAULT, SceneModel(it))
         }
     }
+
+    companion object {
+        private val ASLEEP =
+            WakefulnessModel(
+                state = WakefulnessState.ASLEEP,
+                lastWakeReason = WakeSleepReason.POWER_BUTTON,
+                lastSleepReason = WakeSleepReason.POWER_BUTTON
+            )
+    }
 }
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 60a4951..6309740 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
@@ -48,7 +48,7 @@
     override val isKeyguardShowing: Flow<Boolean> = _isKeyguardShowing
 
     private val _isKeyguardUnlocked = MutableStateFlow(false)
-    override val isKeyguardUnlocked: Flow<Boolean> = _isKeyguardUnlocked
+    override val isKeyguardUnlocked: StateFlow<Boolean> = _isKeyguardUnlocked.asStateFlow()
 
     private val _isKeyguardOccluded = MutableStateFlow(false)
     override val isKeyguardOccluded: Flow<Boolean> = _isKeyguardOccluded
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/scene/SceneTestUtils.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/scene/SceneTestUtils.kt
index d797962..47e1daf4 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/scene/SceneTestUtils.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/scene/SceneTestUtils.kt
@@ -22,11 +22,20 @@
 import com.android.systemui.authentication.data.repository.FakeAuthenticationRepository
 import com.android.systemui.authentication.domain.interactor.AuthenticationInteractor
 import com.android.systemui.bouncer.data.repository.BouncerRepository
+import com.android.systemui.bouncer.data.repository.FakeKeyguardBouncerRepository
 import com.android.systemui.bouncer.domain.interactor.BouncerInteractor
 import com.android.systemui.bouncer.ui.viewmodel.BouncerViewModel
+import com.android.systemui.common.ui.data.repository.FakeConfigurationRepository
 import com.android.systemui.flags.FakeFeatureFlags
 import com.android.systemui.flags.Flags
+import com.android.systemui.keyguard.data.repository.FakeCommandQueue
+import com.android.systemui.keyguard.data.repository.FakeKeyguardRepository
+import com.android.systemui.keyguard.data.repository.KeyguardRepository
+import com.android.systemui.keyguard.domain.interactor.KeyguardInteractor
 import com.android.systemui.keyguard.domain.interactor.LockscreenSceneInteractor
+import com.android.systemui.keyguard.shared.model.WakeSleepReason
+import com.android.systemui.keyguard.shared.model.WakefulnessModel
+import com.android.systemui.keyguard.shared.model.WakefulnessState
 import com.android.systemui.scene.data.repository.SceneContainerRepository
 import com.android.systemui.scene.domain.interactor.SceneInteractor
 import com.android.systemui.scene.shared.model.SceneContainerConfig
@@ -53,7 +62,11 @@
 ) {
     val testDispatcher = StandardTestDispatcher()
     val testScope = TestScope(testDispatcher)
-    val featureFlags = FakeFeatureFlags().apply { set(Flags.SCENE_CONTAINER, true) }
+    val featureFlags =
+        FakeFeatureFlags().apply {
+            set(Flags.SCENE_CONTAINER, true)
+            set(Flags.FACE_AUTH_REFACTOR, false)
+        }
     private val userRepository: UserRepository by lazy {
         FakeUserRepository().apply {
             val users = listOf(UserInfo(/* id=  */ 0, "name", /* flags= */ 0))
@@ -67,6 +80,17 @@
             currentTime = { testScope.currentTime },
         )
     }
+    val keyguardRepository: FakeKeyguardRepository by lazy {
+        FakeKeyguardRepository().apply {
+            setWakefulnessModel(
+                WakefulnessModel(
+                    WakefulnessState.AWAKE,
+                    WakeSleepReason.OTHER,
+                    WakeSleepReason.OTHER,
+                )
+            )
+        }
+    }
     private val context = test.context
 
     fun fakeSceneContainerRepository(
@@ -122,6 +146,20 @@
         )
     }
 
+    fun keyguardRepository(): FakeKeyguardRepository {
+        return keyguardRepository
+    }
+
+    fun keyguardInteractor(repository: KeyguardRepository): KeyguardInteractor {
+        return KeyguardInteractor(
+            repository = repository,
+            commandQueue = FakeCommandQueue(),
+            featureFlags = featureFlags,
+            bouncerRepository = FakeKeyguardBouncerRepository(),
+            configurationRepository = FakeConfigurationRepository()
+        )
+    }
+
     fun bouncerInteractor(
         authenticationInteractor: AuthenticationInteractor,
         sceneInteractor: SceneInteractor,