Fix UMO missing on glanceable hub

Similar to ag/27798458, adds a safe-guard to ensure we show the umo on
the hub if we are idle on the hub and not transitioning.

Fixes: 339093481
Test: atest CommunalTransitionViewModelTest
Test: atest MediaHierarchyManagerTest
Flag: com.android.systemui.communal_hub
Change-Id: I38aac5758029d4bc1e75ebfd1031fd01394a9b93
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/ui/viewmodel/CommunalTransitionViewModelTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/ui/viewmodel/CommunalTransitionViewModelTest.kt
index 0250c9d..baeb2dd 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/ui/viewmodel/CommunalTransitionViewModelTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/ui/viewmodel/CommunalTransitionViewModelTest.kt
@@ -31,7 +31,6 @@
 import kotlinx.coroutines.test.TestScope
 import kotlinx.coroutines.test.runCurrent
 import kotlinx.coroutines.test.runTest
-import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
 
@@ -42,13 +41,11 @@
     private val kosmos = testKosmos()
     private val testScope = kosmos.testScope
 
-    val keyguardTransitionRepository = kosmos.fakeKeyguardTransitionRepository
+    private val keyguardTransitionRepository = kosmos.fakeKeyguardTransitionRepository
+    private val communalSceneRepository = kosmos.fakeCommunalSceneRepository
 
-    private lateinit var underTest: CommunalTransitionViewModel
-
-    @Before
-    fun setup() {
-        underTest = kosmos.communalTransitionViewModel
+    private val underTest: CommunalTransitionViewModel by lazy {
+        kosmos.communalTransitionViewModel
     }
 
     @Test
@@ -60,11 +57,7 @@
             enterCommunal(from = KeyguardState.LOCKSCREEN)
             assertThat(isUmoOnCommunal).isTrue()
 
-            keyguardTransitionRepository.sendTransitionSteps(
-                from = KeyguardState.GLANCEABLE_HUB,
-                to = KeyguardState.LOCKSCREEN,
-                testScope
-            )
+            exitCommunal(to = KeyguardState.LOCKSCREEN)
             assertThat(isUmoOnCommunal).isFalse()
         }
 
@@ -77,11 +70,7 @@
             enterCommunal(from = KeyguardState.DREAMING)
             assertThat(isUmoOnCommunal).isTrue()
 
-            keyguardTransitionRepository.sendTransitionSteps(
-                from = KeyguardState.GLANCEABLE_HUB,
-                to = KeyguardState.DREAMING,
-                testScope
-            )
+            exitCommunal(to = KeyguardState.DREAMING)
             assertThat(isUmoOnCommunal).isFalse()
         }
 
@@ -94,11 +83,7 @@
             enterCommunal(from = KeyguardState.OCCLUDED)
             assertThat(isUmoOnCommunal).isTrue()
 
-            keyguardTransitionRepository.sendTransitionSteps(
-                from = KeyguardState.GLANCEABLE_HUB,
-                to = KeyguardState.OCCLUDED,
-                testScope
-            )
+            exitCommunal(to = KeyguardState.OCCLUDED)
             assertThat(isUmoOnCommunal).isFalse()
         }
 
@@ -112,20 +97,42 @@
             assertThat(isUmoOnCommunal).isTrue()
 
             // Communal is no longer visible.
-            kosmos.fakeCommunalSceneRepository.changeScene(CommunalScenes.Blank)
-            runCurrent()
+            communalSceneRepository.changeScene(CommunalScenes.Blank)
 
             // isUmoOnCommunal returns false, even without any keyguard transition.
             assertThat(isUmoOnCommunal).isFalse()
         }
 
+    @Test
+    fun isUmoOnCommunal_idleOnCommunal_returnsTrue() =
+        testScope.runTest {
+            val isUmoOnCommunal by collectLastValue(underTest.isUmoOnCommunal)
+            assertThat(isUmoOnCommunal).isFalse()
+
+            // Communal is fully visible.
+            communalSceneRepository.changeScene(CommunalScenes.Communal)
+
+            // isUmoOnCommunal returns true, even without any keyguard transition.
+            assertThat(isUmoOnCommunal).isTrue()
+        }
+
     private suspend fun TestScope.enterCommunal(from: KeyguardState) {
         keyguardTransitionRepository.sendTransitionSteps(
             from = from,
             to = KeyguardState.GLANCEABLE_HUB,
             testScope
         )
-        kosmos.fakeCommunalSceneRepository.changeScene(CommunalScenes.Communal)
+        communalSceneRepository.changeScene(CommunalScenes.Communal)
+        runCurrent()
+    }
+
+    private suspend fun TestScope.exitCommunal(to: KeyguardState) {
+        keyguardTransitionRepository.sendTransitionSteps(
+            from = KeyguardState.GLANCEABLE_HUB,
+            to = to,
+            testScope
+        )
+        communalSceneRepository.changeScene(CommunalScenes.Blank)
         runCurrent()
     }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/communal/ui/viewmodel/CommunalTransitionViewModel.kt b/packages/SystemUI/src/com/android/systemui/communal/ui/viewmodel/CommunalTransitionViewModel.kt
index fce18a26..e1408a0 100644
--- a/packages/SystemUI/src/com/android/systemui/communal/ui/viewmodel/CommunalTransitionViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/communal/ui/viewmodel/CommunalTransitionViewModel.kt
@@ -21,6 +21,7 @@
 import com.android.systemui.communal.domain.interactor.CommunalSceneInteractor
 import com.android.systemui.communal.util.CommunalColors
 import com.android.systemui.dagger.SysUISingleton
+import com.android.systemui.dagger.qualifiers.Application
 import com.android.systemui.keyguard.domain.interactor.KeyguardTransitionInteractor
 import com.android.systemui.keyguard.shared.model.Edge
 import com.android.systemui.keyguard.shared.model.KeyguardState
@@ -31,13 +32,18 @@
 import com.android.systemui.keyguard.ui.viewmodel.LockscreenToGlanceableHubTransitionViewModel
 import com.android.systemui.scene.shared.model.Scenes
 import com.android.systemui.util.kotlin.BooleanFlowOperators.allOf
+import com.android.systemui.util.kotlin.BooleanFlowOperators.anyOf
 import javax.inject.Inject
+import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.SharingStarted
 import kotlinx.coroutines.flow.combine
 import kotlinx.coroutines.flow.filter
 import kotlinx.coroutines.flow.map
 import kotlinx.coroutines.flow.merge
+import kotlinx.coroutines.flow.onStart
+import kotlinx.coroutines.flow.stateIn
 
 /** View model for transitions related to the communal hub. */
 @OptIn(ExperimentalCoroutinesApi::class)
@@ -45,6 +51,7 @@
 class CommunalTransitionViewModel
 @Inject
 constructor(
+    @Application applicationScope: CoroutineScope,
     communalColors: CommunalColors,
     glanceableHubToLockscreenTransitionViewModel: GlanceableHubToLockscreenTransitionViewModel,
     lockscreenToGlanceableHubTransitionViewModel: LockscreenToGlanceableHubTransitionViewModel,
@@ -85,21 +92,32 @@
      * of UMO should be updated.
      */
     val isUmoOnCommunal: Flow<Boolean> =
-        allOf(
-            // Only show UMO on the hub if the hub is at least partially visible. This prevents
-            // the UMO from being missing on the lock screen when going from the hub to lock
-            // screen in some way other than through a direct transition, such as unlocking from
-            // the hub, then pressing power twice to go back to the lock screen.
-            communalSceneInteractor.isCommunalVisible,
-            merge(
-                lockscreenToGlanceableHubTransitionViewModel.showUmo,
-                glanceableHubToLockscreenTransitionViewModel.showUmo,
-                dreamToGlanceableHubTransitionViewModel.showUmo,
-                glanceableHubToDreamTransitionViewModel.showUmo,
-                showUmoFromOccludedToGlanceableHub,
-                showUmoFromGlanceableHubToOccluded,
+        anyOf(
+                communalSceneInteractor.isIdleOnCommunal,
+                allOf(
+                    // Only show UMO on the hub if the hub is at least partially visible. This
+                    // prevents
+                    // the UMO from being missing on the lock screen when going from the hub to lock
+                    // screen in some way other than through a direct transition, such as unlocking
+                    // from
+                    // the hub, then pressing power twice to go back to the lock screen.
+                    communalSceneInteractor.isCommunalVisible,
+                    merge(
+                            lockscreenToGlanceableHubTransitionViewModel.showUmo,
+                            glanceableHubToLockscreenTransitionViewModel.showUmo,
+                            dreamToGlanceableHubTransitionViewModel.showUmo,
+                            glanceableHubToDreamTransitionViewModel.showUmo,
+                            showUmoFromOccludedToGlanceableHub,
+                            showUmoFromGlanceableHubToOccluded,
+                        )
+                        .onStart { emit(false) }
+                )
             )
-        )
+            .stateIn(
+                scope = applicationScope,
+                started = SharingStarted.WhileSubscribed(),
+                initialValue = false
+            )
 
     /** Whether to show communal when exiting the occluded state. */
     val showCommunalFromOccluded: Flow<Boolean> = communalInteractor.showCommunalFromOccluded
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/communal/ui/viewmodel/CommunalTransitionViewModelKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/communal/ui/viewmodel/CommunalTransitionViewModelKosmos.kt
index 3f38408..1ae8449 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/communal/ui/viewmodel/CommunalTransitionViewModelKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/communal/ui/viewmodel/CommunalTransitionViewModelKosmos.kt
@@ -25,12 +25,14 @@
 import com.android.systemui.keyguard.ui.viewmodel.glanceableHubToLockscreenTransitionViewModel
 import com.android.systemui.keyguard.ui.viewmodel.lockscreenToGlanceableHubTransitionViewModel
 import com.android.systemui.kosmos.Kosmos
+import com.android.systemui.kosmos.applicationCoroutineScope
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 
 @OptIn(ExperimentalCoroutinesApi::class)
 val Kosmos.communalTransitionViewModel by
     Kosmos.Fixture {
         CommunalTransitionViewModel(
+            applicationScope = applicationCoroutineScope,
             glanceableHubToLockscreenTransitionViewModel =
                 glanceableHubToLockscreenTransitionViewModel,
             lockscreenToGlanceableHubTransitionViewModel =