Merge "Shared notification container - Update isLockscreen" into main
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/ui/viewmodel/SharedNotificationContainerViewModel.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/ui/viewmodel/SharedNotificationContainerViewModel.kt
index f750fed..1229cb9 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/ui/viewmodel/SharedNotificationContainerViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/ui/viewmodel/SharedNotificationContainerViewModel.kt
@@ -31,6 +31,7 @@
 import kotlinx.coroutines.flow.distinctUntilChanged
 import kotlinx.coroutines.flow.flatMapLatest
 import kotlinx.coroutines.flow.map
+import kotlinx.coroutines.flow.onStart
 
 /** View-model for the shared notification container, used by both the shade and keyguard spaces */
 class SharedNotificationContainerViewModel
@@ -45,8 +46,8 @@
 ) {
     private val statesForConstrainedNotifications =
         setOf(
-            KeyguardState.LOCKSCREEN,
             KeyguardState.AOD,
+            KeyguardState.LOCKSCREEN,
             KeyguardState.DOZING,
             KeyguardState.ALTERNATE_BOUNCER,
             KeyguardState.PRIMARY_BOUNCER
@@ -68,8 +69,17 @@
 
     /** If the user is visually on one of the unoccluded lockscreen states. */
     val isOnLockscreen: Flow<Boolean> =
-        keyguardTransitionInteractor.finishedKeyguardState
-            .map { statesForConstrainedNotifications.contains(it) }
+        combine(
+                keyguardTransitionInteractor.finishedKeyguardState.map {
+                    statesForConstrainedNotifications.contains(it)
+                },
+                keyguardTransitionInteractor
+                    .transitionValue(KeyguardState.LOCKSCREEN)
+                    .onStart { emit(0f) }
+                    .map { it > 0 }
+            ) { constrainedNotificationState, transitioningToOrFromLockscreen ->
+                constrainedNotificationState || transitioningToOrFromLockscreen
+            }
             .distinctUntilChanged()
 
     /** Are we purely on the keyguard without the shade/qs? */
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/ui/viewmodel/SharedNotificationContainerViewModelTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/ui/viewmodel/SharedNotificationContainerViewModelTest.kt
index ac11ff2..0a7dc4e 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/ui/viewmodel/SharedNotificationContainerViewModelTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/ui/viewmodel/SharedNotificationContainerViewModelTest.kt
@@ -192,10 +192,26 @@
             val isOnLockscreen by collectLastValue(underTest.isOnLockscreen)
 
             keyguardTransitionRepository.sendTransitionStep(
-                TransitionStep(to = KeyguardState.GONE, transitionState = TransitionState.FINISHED)
+                TransitionStep(
+                    from = KeyguardState.LOCKSCREEN,
+                    to = KeyguardState.GONE,
+                    value = 1f,
+                    transitionState = TransitionState.FINISHED
+                )
             )
             assertThat(isOnLockscreen).isFalse()
 
+            // While progressing from lockscreen, should still be true
+            keyguardTransitionRepository.sendTransitionStep(
+                TransitionStep(
+                    from = KeyguardState.LOCKSCREEN,
+                    to = KeyguardState.GONE,
+                    value = 0.8f,
+                    transitionState = TransitionState.RUNNING
+                )
+            )
+            assertThat(isOnLockscreen).isTrue()
+
             keyguardTransitionRepository.sendTransitionStep(
                 TransitionStep(
                     to = KeyguardState.LOCKSCREEN,