Merge "Fix UMO missing on lock screen when starting from hub" into main
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 f9d5073..0250c9d 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
@@ -19,17 +19,23 @@
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import com.android.systemui.SysuiTestCase
+import com.android.systemui.communal.data.repository.fakeCommunalSceneRepository
+import com.android.systemui.communal.shared.model.CommunalScenes
 import com.android.systemui.coroutines.collectLastValue
 import com.android.systemui.keyguard.data.repository.fakeKeyguardTransitionRepository
 import com.android.systemui.keyguard.shared.model.KeyguardState
 import com.android.systemui.kosmos.testScope
 import com.android.systemui.testKosmos
 import com.google.common.truth.Truth.assertThat
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+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
 
+@OptIn(ExperimentalCoroutinesApi::class)
 @SmallTest
 @RunWith(AndroidJUnit4::class)
 class CommunalTransitionViewModelTest : SysuiTestCase() {
@@ -49,13 +55,9 @@
     fun testIsUmoOnCommunalDuringTransitionBetweenLockscreenAndGlanceableHub() =
         testScope.runTest {
             val isUmoOnCommunal by collectLastValue(underTest.isUmoOnCommunal)
-            assertThat(isUmoOnCommunal).isNull()
+            runCurrent()
 
-            keyguardTransitionRepository.sendTransitionSteps(
-                from = KeyguardState.LOCKSCREEN,
-                to = KeyguardState.GLANCEABLE_HUB,
-                testScope
-            )
+            enterCommunal(from = KeyguardState.LOCKSCREEN)
             assertThat(isUmoOnCommunal).isTrue()
 
             keyguardTransitionRepository.sendTransitionSteps(
@@ -70,13 +72,9 @@
     fun testIsUmoOnCommunalDuringTransitionBetweenDreamingAndGlanceableHub() =
         testScope.runTest {
             val isUmoOnCommunal by collectLastValue(underTest.isUmoOnCommunal)
-            assertThat(isUmoOnCommunal).isNull()
+            runCurrent()
 
-            keyguardTransitionRepository.sendTransitionSteps(
-                from = KeyguardState.DREAMING,
-                to = KeyguardState.GLANCEABLE_HUB,
-                testScope
-            )
+            enterCommunal(from = KeyguardState.DREAMING)
             assertThat(isUmoOnCommunal).isTrue()
 
             keyguardTransitionRepository.sendTransitionSteps(
@@ -91,13 +89,9 @@
     fun testIsUmoOnCommunalDuringTransitionBetweenOccludedAndGlanceableHub() =
         testScope.runTest {
             val isUmoOnCommunal by collectLastValue(underTest.isUmoOnCommunal)
-            assertThat(isUmoOnCommunal).isNull()
+            runCurrent()
 
-            keyguardTransitionRepository.sendTransitionSteps(
-                from = KeyguardState.OCCLUDED,
-                to = KeyguardState.GLANCEABLE_HUB,
-                testScope
-            )
+            enterCommunal(from = KeyguardState.OCCLUDED)
             assertThat(isUmoOnCommunal).isTrue()
 
             keyguardTransitionRepository.sendTransitionSteps(
@@ -105,7 +99,33 @@
                 to = KeyguardState.OCCLUDED,
                 testScope
             )
-
             assertThat(isUmoOnCommunal).isFalse()
         }
+
+    @Test
+    fun isUmoOnCommunal_noLongerVisible_returnsFalse() =
+        testScope.runTest {
+            val isUmoOnCommunal by collectLastValue(underTest.isUmoOnCommunal)
+            runCurrent()
+
+            enterCommunal(from = KeyguardState.LOCKSCREEN)
+            assertThat(isUmoOnCommunal).isTrue()
+
+            // Communal is no longer visible.
+            kosmos.fakeCommunalSceneRepository.changeScene(CommunalScenes.Blank)
+            runCurrent()
+
+            // isUmoOnCommunal returns false, even without any keyguard transition.
+            assertThat(isUmoOnCommunal).isFalse()
+        }
+
+    private suspend fun TestScope.enterCommunal(from: KeyguardState) {
+        keyguardTransitionRepository.sendTransitionSteps(
+            from = from,
+            to = KeyguardState.GLANCEABLE_HUB,
+            testScope
+        )
+        kosmos.fakeCommunalSceneRepository.changeScene(CommunalScenes.Communal)
+        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 7a20ebc..fce18a26 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
@@ -18,6 +18,7 @@
 
 import android.graphics.Color
 import com.android.systemui.communal.domain.interactor.CommunalInteractor
+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.keyguard.domain.interactor.KeyguardTransitionInteractor
@@ -29,11 +30,11 @@
 import com.android.systemui.keyguard.ui.viewmodel.GlanceableHubToLockscreenTransitionViewModel
 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 javax.inject.Inject
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.combine
-import kotlinx.coroutines.flow.distinctUntilChanged
 import kotlinx.coroutines.flow.filter
 import kotlinx.coroutines.flow.map
 import kotlinx.coroutines.flow.merge
@@ -50,16 +51,18 @@
     dreamToGlanceableHubTransitionViewModel: DreamingToGlanceableHubTransitionViewModel,
     glanceableHubToDreamTransitionViewModel: GlanceableHubToDreamingTransitionViewModel,
     communalInteractor: CommunalInteractor,
-    keyguardTransitionInteractor: KeyguardTransitionInteractor,
+    communalSceneInteractor: CommunalSceneInteractor,
+    keyguardTransitionInteractor: KeyguardTransitionInteractor
 ) {
     // Show UMO on glanceable hub immediately on transition into glanceable hub
     private val showUmoFromOccludedToGlanceableHub: Flow<Boolean> =
         keyguardTransitionInteractor
-            .transition(Edge.create(from = KeyguardState.OCCLUDED))
+            .transition(
+                Edge.create(from = KeyguardState.OCCLUDED, to = KeyguardState.GLANCEABLE_HUB)
+            )
             .filter {
-                it.to == KeyguardState.GLANCEABLE_HUB &&
-                    (it.transitionState == TransitionState.STARTED ||
-                        it.transitionState == TransitionState.CANCELED)
+                (it.transitionState == TransitionState.STARTED ||
+                    it.transitionState == TransitionState.CANCELED)
             }
             .map { it.transitionState == TransitionState.STARTED }
 
@@ -82,7 +85,13 @@
      * of UMO should be updated.
      */
     val isUmoOnCommunal: Flow<Boolean> =
-        merge(
+        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,
@@ -90,7 +99,7 @@
                 showUmoFromOccludedToGlanceableHub,
                 showUmoFromGlanceableHubToOccluded,
             )
-            .distinctUntilChanged()
+        )
 
     /** Whether to show communal when exiting the occluded state. */
     val showCommunalFromOccluded: Flow<Boolean> = communalInteractor.showCommunalFromOccluded
diff --git a/packages/SystemUI/tests/src/com/android/systemui/media/controls/ui/controller/MediaHierarchyManagerTest.kt b/packages/SystemUI/tests/src/com/android/systemui/media/controls/ui/controller/MediaHierarchyManagerTest.kt
index bba01bd..6c350cb 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/media/controls/ui/controller/MediaHierarchyManagerTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/media/controls/ui/controller/MediaHierarchyManagerTest.kt
@@ -25,6 +25,8 @@
 import androidx.test.filters.SmallTest
 import com.android.keyguard.KeyguardViewController
 import com.android.systemui.SysuiTestCase
+import com.android.systemui.communal.data.repository.fakeCommunalSceneRepository
+import com.android.systemui.communal.shared.model.CommunalScenes
 import com.android.systemui.communal.ui.viewmodel.communalTransitionViewModel
 import com.android.systemui.controls.controller.ControlsControllerImplTest.Companion.eq
 import com.android.systemui.dreams.DreamOverlayStateController
@@ -498,6 +500,8 @@
                 to = KeyguardState.GLANCEABLE_HUB,
                 testScope = testScope,
             )
+            kosmos.fakeCommunalSceneRepository.changeScene(CommunalScenes.Communal)
+            runCurrent()
             mediaHierarchyManager.qsExpansion = 0f
             mediaHierarchyManager.setTransitionToFullShadeAmount(123f)
 
@@ -542,6 +546,8 @@
                 to = KeyguardState.GLANCEABLE_HUB,
                 testScope = testScope,
             )
+            kosmos.fakeCommunalSceneRepository.changeScene(CommunalScenes.Communal)
+            runCurrent()
             verify(mediaCarouselController)
                 .onDesiredLocationChanged(
                     eq(MediaHierarchyManager.LOCATION_COMMUNAL_HUB),
@@ -557,6 +563,8 @@
                 to = KeyguardState.LOCKSCREEN,
                 testScope = testScope,
             )
+            kosmos.fakeCommunalSceneRepository.changeScene(CommunalScenes.Blank)
+            runCurrent()
             verify(mediaCarouselController)
                 .onDesiredLocationChanged(
                     eq(MediaHierarchyManager.LOCATION_QQS),
@@ -579,6 +587,8 @@
                 to = KeyguardState.GLANCEABLE_HUB,
                 testScope = testScope,
             )
+            kosmos.fakeCommunalSceneRepository.changeScene(CommunalScenes.Communal)
+            runCurrent()
             verify(mediaCarouselController)
                 .onDesiredLocationChanged(
                     eq(MediaHierarchyManager.LOCATION_COMMUNAL_HUB),
@@ -600,6 +610,8 @@
                 to = KeyguardState.GLANCEABLE_HUB,
                 testScope = testScope,
             )
+            kosmos.fakeCommunalSceneRepository.changeScene(CommunalScenes.Communal)
+            runCurrent()
             verify(mediaCarouselController)
                 .onDesiredLocationChanged(
                     eq(MediaHierarchyManager.LOCATION_COMMUNAL_HUB),
@@ -635,6 +647,8 @@
                 to = KeyguardState.GLANCEABLE_HUB,
                 testScope = testScope,
             )
+            kosmos.fakeCommunalSceneRepository.changeScene(CommunalScenes.Communal)
+            runCurrent()
             // Mock the behavior for dreaming that pulling down shade will immediately set QS as
             // expanded
             expandQS()
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 e3c218d..3f38408 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
@@ -17,6 +17,7 @@
 package com.android.systemui.communal.ui.viewmodel
 
 import com.android.systemui.communal.domain.interactor.communalInteractor
+import com.android.systemui.communal.domain.interactor.communalSceneInteractor
 import com.android.systemui.communal.util.communalColors
 import com.android.systemui.keyguard.domain.interactor.keyguardTransitionInteractor
 import com.android.systemui.keyguard.ui.viewmodel.dreamingToGlanceableHubTransitionViewModel
@@ -37,6 +38,7 @@
             dreamToGlanceableHubTransitionViewModel = dreamingToGlanceableHubTransitionViewModel,
             glanceableHubToDreamTransitionViewModel = glanceableHubToDreamingTransitionViewModel,
             communalInteractor = communalInteractor,
+            communalSceneInteractor = communalSceneInteractor,
             keyguardTransitionInteractor = keyguardTransitionInteractor,
             communalColors = communalColors,
         )