[ROSP] Use Hydrator and AssistedInject for StatusBarPopupChipsViewModel

This change refactors `StatusBarPopupChipsViewModel` to utilize snapshot
state instead of exposing a flow.

Bug: b/394887130
Flag: com.android.systemui.status_bar_popup_chips

Test: Make sure popup chips show up in the status bar.

Change-Id: I89dc8971fcaa39992f7065526aa14a312f2d22c7

7 files changed, 86 insertions, 60 deletions

diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/pipeline/shared/ui/viewmodel/FakeHomeStatusBarViewModel.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/pipeline/shared/ui/viewmodel/FakeHomeStatusBarViewModel.kt
index 183cd8f1ae8b..eb961bd5f4ae 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/pipeline/shared/ui/viewmodel/FakeHomeStatusBarViewModel.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/pipeline/shared/ui/viewmodel/FakeHomeStatusBarViewModel.kt
@@ -57,7 +57,7 @@ class FakeHomeStatusBarViewModel(
     override val ongoingActivityChipsLegacy =
         MutableStateFlow(MultipleOngoingActivityChipsModelLegacy())
 
-    override val statusBarPopupChips = MutableStateFlow(emptyList<PopupChipModel.Shown>())
+    override val popupChips = emptyList<PopupChipModel.Shown>()
 
     override val mediaProjectionStopDialogDueToCallEndedState =
         MutableStateFlow(MediaProjectionStopDialogModel.Hidden)
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/featurepods/popups/ui/viewmodel/StatusBarPopupChipsViewModel.kt b/packages/SystemUI/src/com/android/systemui/statusbar/featurepods/popups/ui/viewmodel/StatusBarPopupChipsViewModel.kt
index caa8e6cc02c3..786794fd3b37 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/featurepods/popups/ui/viewmodel/StatusBarPopupChipsViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/featurepods/popups/ui/viewmodel/StatusBarPopupChipsViewModel.kt
@@ -16,49 +16,53 @@
 
 package com.android.systemui.statusbar.featurepods.popups.ui.viewmodel
 
-import com.android.systemui.dagger.SysUISingleton
-import com.android.systemui.dagger.qualifiers.Background
+import androidx.compose.runtime.derivedStateOf
+import androidx.compose.runtime.getValue
+import com.android.systemui.lifecycle.ExclusiveActivatable
+import com.android.systemui.lifecycle.Hydrator
 import com.android.systemui.statusbar.featurepods.media.ui.viewmodel.MediaControlChipViewModel
 import com.android.systemui.statusbar.featurepods.popups.StatusBarPopupChips
 import com.android.systemui.statusbar.featurepods.popups.shared.model.PopupChipId
 import com.android.systemui.statusbar.featurepods.popups.shared.model.PopupChipModel
-import javax.inject.Inject
-import kotlinx.coroutines.CoroutineScope
-import kotlinx.coroutines.flow.MutableStateFlow
-import kotlinx.coroutines.flow.SharingStarted
-import kotlinx.coroutines.flow.StateFlow
-import kotlinx.coroutines.flow.asStateFlow
+import dagger.assisted.AssistedFactory
+import dagger.assisted.AssistedInject
 import kotlinx.coroutines.flow.map
-import kotlinx.coroutines.flow.stateIn
 
 /**
  * View model deciding which system process chips to show in the status bar. Emits a list of
  * PopupChipModels.
  */
-@SysUISingleton
 class StatusBarPopupChipsViewModel
-@Inject
-constructor(
-    @Background scope: CoroutineScope,
-    mediaControlChipViewModel: MediaControlChipViewModel,
-) {
-    private data class PopupChipBundle(
-        val media: PopupChipModel = PopupChipModel.Hidden(chipId = PopupChipId.MediaControl)
-    )
+@AssistedInject
+constructor(mediaControlChip: MediaControlChipViewModel) : ExclusiveActivatable() {
+    private val hydrator: Hydrator = Hydrator("StatusBarPopupChipsViewModel.hydrator")
 
-    private val incomingPopupChipBundle: StateFlow<PopupChipBundle?> =
-        mediaControlChipViewModel.chip
-            .map { chip -> PopupChipBundle(media = chip) }
-            .stateIn(scope, SharingStarted.WhileSubscribed(), PopupChipBundle())
+    private val incomingPopupChipBundle: PopupChipBundle by
+        hydrator.hydratedStateOf(
+            traceName = "incomingPopupChipBundle",
+            initialValue = PopupChipBundle(),
+            source = mediaControlChip.chip.map { chip -> PopupChipBundle(media = chip) },
+        )
 
-    val shownPopupChips: StateFlow<List<PopupChipModel.Shown>> =
+    val shownPopupChips: List<PopupChipModel.Shown> by derivedStateOf {
         if (StatusBarPopupChips.isEnabled) {
-            incomingPopupChipBundle
-                .map { bundle ->
-                    listOfNotNull(bundle?.media).filterIsInstance<PopupChipModel.Shown>()
-                }
-                .stateIn(scope, SharingStarted.WhileSubscribed(), emptyList())
+            val bundle = incomingPopupChipBundle
+            listOfNotNull(bundle.media).filterIsInstance<PopupChipModel.Shown>()
         } else {
-            MutableStateFlow(emptyList<PopupChipModel.Shown>()).asStateFlow()
+            emptyList()
         }
+    }
+
+    override suspend fun onActivated(): Nothing {
+        hydrator.activate()
+    }
+
+    private data class PopupChipBundle(
+        val media: PopupChipModel = PopupChipModel.Hidden(chipId = PopupChipId.MediaControl)
+    )
+
+    @AssistedFactory
+    interface Factory {
+        fun create(): StatusBarPopupChipsViewModel
+    }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/shared/ui/composable/StatusBarRoot.kt b/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/shared/ui/composable/StatusBarRoot.kt
index a961713230c8..39a1b46292b0 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/shared/ui/composable/StatusBarRoot.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/shared/ui/composable/StatusBarRoot.kt
@@ -255,10 +255,9 @@ fun StatusBarRoot(
                                 )
 
                                 setContent {
-                                    val chips =
-                                        statusBarViewModel.statusBarPopupChips
-                                            .collectAsStateWithLifecycle()
-                                    StatusBarPopupChipsContainer(chips = chips.value)
+                                    StatusBarPopupChipsContainer(
+                                        chips = statusBarViewModel.popupChips
+                                    )
                                 }
                             }
                         endSideContent.addView(composeView, 0)
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/shared/ui/viewmodel/HomeStatusBarViewModel.kt b/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/shared/ui/viewmodel/HomeStatusBarViewModel.kt
index f396cbfc8000..9ae2cb203d91 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/shared/ui/viewmodel/HomeStatusBarViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/shared/ui/viewmodel/HomeStatusBarViewModel.kt
@@ -20,6 +20,7 @@ import android.annotation.ColorInt
 import android.graphics.Rect
 import android.view.View
 import androidx.compose.runtime.getValue
+import com.android.app.tracing.coroutines.launchTraced as launch
 import com.android.systemui.dagger.qualifiers.Background
 import com.android.systemui.keyguard.domain.interactor.KeyguardInteractor
 import com.android.systemui.keyguard.domain.interactor.KeyguardTransitionInteractor
@@ -29,6 +30,7 @@ import com.android.systemui.keyguard.shared.model.KeyguardState.GONE
 import com.android.systemui.keyguard.shared.model.KeyguardState.LOCKSCREEN
 import com.android.systemui.keyguard.shared.model.KeyguardState.OCCLUDED
 import com.android.systemui.keyguard.shared.model.TransitionState
+import com.android.systemui.lifecycle.Activatable
 import com.android.systemui.lifecycle.ExclusiveActivatable
 import com.android.systemui.lifecycle.Hydrator
 import com.android.systemui.log.table.TableLogBufferFactory
@@ -49,6 +51,7 @@ import com.android.systemui.statusbar.chips.ui.model.OngoingActivityChipModel
 import com.android.systemui.statusbar.chips.ui.viewmodel.OngoingActivityChipsViewModel
 import com.android.systemui.statusbar.events.domain.interactor.SystemStatusEventAnimationInteractor
 import com.android.systemui.statusbar.events.shared.model.SystemEventAnimationState.Idle
+import com.android.systemui.statusbar.featurepods.popups.StatusBarPopupChips
 import com.android.systemui.statusbar.featurepods.popups.shared.model.PopupChipModel
 import com.android.systemui.statusbar.featurepods.popups.ui.viewmodel.StatusBarPopupChipsViewModel
 import com.android.systemui.statusbar.headsup.shared.StatusBarNoHunBehavior
@@ -71,6 +74,8 @@ import dagger.assisted.AssistedFactory
 import dagger.assisted.AssistedInject
 import kotlinx.coroutines.CoroutineDispatcher
 import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.awaitCancellation
+import kotlinx.coroutines.coroutineScope
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.SharingStarted
 import kotlinx.coroutines.flow.StateFlow
@@ -94,7 +99,7 @@ import kotlinx.coroutines.flow.stateIn
  * [StatusBarHideIconsForBouncerManager]. We should move those pieces of logic to this class instead
  * so that it's all in one place and easily testable outside of the fragment.
  */
-interface HomeStatusBarViewModel {
+interface HomeStatusBarViewModel : Activatable {
     /** Factory to create the view model for the battery icon */
     val batteryViewModelFactory: BatteryViewModel.Factory
 
@@ -133,7 +138,7 @@ interface HomeStatusBarViewModel {
     val operatorNameViewModel: StatusBarOperatorNameViewModel
 
     /** The popup chips that should be shown on the right-hand side of the status bar. */
-    val statusBarPopupChips: StateFlow<List<PopupChipModel.Shown>>
+    val popupChips: List<PopupChipModel.Shown>
 
     /**
      * True if the current scene can show the home status bar (aka this status bar), and false if
@@ -208,7 +213,7 @@ constructor(
     shadeInteractor: ShadeInteractor,
     shareToAppChipViewModel: ShareToAppChipViewModel,
     ongoingActivityChipsViewModel: OngoingActivityChipsViewModel,
-    statusBarPopupChipsViewModel: StatusBarPopupChipsViewModel,
+    statusBarPopupChipsViewModelFactory: StatusBarPopupChipsViewModel.Factory,
     animations: SystemStatusEventAnimationInteractor,
     statusBarContentInsetsViewModelStore: StatusBarContentInsetsViewModelStore,
     @Background bgScope: CoroutineScope,
@@ -219,6 +224,8 @@ constructor(
 
     val tableLogger = tableLoggerFactory.getOrCreate(tableLogBufferName(thisDisplayId), 200)
 
+    private val statusBarPopupChips by lazy { statusBarPopupChipsViewModelFactory.create() }
+
     override val isTransitioningFromLockscreenToOccluded: StateFlow<Boolean> =
         keyguardTransitionInteractor
             .isInTransition(Edge.create(from = LOCKSCREEN, to = OCCLUDED))
@@ -246,7 +253,8 @@ constructor(
 
     override val ongoingActivityChipsLegacy = ongoingActivityChipsViewModel.chipsLegacy
 
-    override val statusBarPopupChips = statusBarPopupChipsViewModel.shownPopupChips
+    override val popupChips
+        get() = statusBarPopupChips.shownPopupChips
 
     override val isHomeStatusBarAllowedByScene: StateFlow<Boolean> =
         combine(
@@ -495,7 +503,13 @@ constructor(
     private fun Boolean.toVisibleOrInvisible(): Int = if (this) View.VISIBLE else View.INVISIBLE
 
     override suspend fun onActivated(): Nothing {
-        hydrator.activate()
+        coroutineScope {
+            launch { hydrator.activate() }
+            if (StatusBarPopupChips.isEnabled) {
+                launch { statusBarPopupChips.activate() }
+            }
+            awaitCancellation()
+        }
     }
 
     /** Inject this to create the display-dependent view model */
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/featurepods/popups/ui/viewmodel/StatusBarPopupChipsViewModelTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/featurepods/popups/ui/viewmodel/StatusBarPopupChipsViewModelTest.kt
index fcbf0fe9a37a..32b6dc0e8f3d 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/featurepods/popups/ui/viewmodel/StatusBarPopupChipsViewModelTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/featurepods/popups/ui/viewmodel/StatusBarPopupChipsViewModelTest.kt
@@ -17,19 +17,23 @@
 package com.android.systemui.statusbar.featurepods.popups.ui.viewmodel
 
 import android.platform.test.annotations.EnableFlags
+import androidx.compose.runtime.snapshots.Snapshot
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import com.android.systemui.SysuiTestCase
-import com.android.systemui.coroutines.collectLastValue
+import com.android.systemui.kosmos.runTest
 import com.android.systemui.kosmos.testScope
+import com.android.systemui.kosmos.useUnconfinedTestDispatcher
+import com.android.systemui.lifecycle.activateIn
 import com.android.systemui.media.controls.data.repository.mediaFilterRepository
 import com.android.systemui.media.controls.shared.model.MediaData
 import com.android.systemui.media.controls.shared.model.MediaDataLoadingModel
-import com.android.systemui.statusbar.featurepods.popups.shared.model.PopupChipId
 import com.android.systemui.statusbar.featurepods.popups.StatusBarPopupChips
+import com.android.systemui.statusbar.featurepods.popups.shared.model.PopupChipId
 import com.android.systemui.testKosmos
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.test.runTest
+import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
 
@@ -37,30 +41,34 @@ import org.junit.runner.RunWith
 @EnableFlags(StatusBarPopupChips.FLAG_NAME)
 @RunWith(AndroidJUnit4::class)
 class StatusBarPopupChipsViewModelTest : SysuiTestCase() {
-    private val kosmos = testKosmos()
-    private val testScope = kosmos.testScope
-    private val mediaFilterRepository = kosmos.mediaFilterRepository
-    private val underTest = kosmos.statusBarPopupChipsViewModel
+    private val kosmos = testKosmos().useUnconfinedTestDispatcher()
+    private val underTest = kosmos.statusBarPopupChipsViewModelFactory.create()
+
+    @Before
+    fun setUp() {
+        underTest.activateIn(kosmos.testScope)
+    }
 
     @Test
     fun shownPopupChips_allHidden_empty() =
-        testScope.runTest {
-            val shownPopupChips by collectLastValue(underTest.shownPopupChips)
+        kosmos.runTest {
+            val shownPopupChips = underTest.shownPopupChips
             assertThat(shownPopupChips).isEmpty()
         }
 
     @Test
     fun shownPopupChips_activeMedia_restHidden_mediaControlChipShown() =
-        testScope.runTest {
-            val shownPopupChips by collectLastValue(underTest.shownPopupChips)
-
+        kosmos.runTest {
+            val shownPopupChips = underTest.shownPopupChips
             val userMedia = MediaData(active = true, song = "test")
             val instanceId = userMedia.instanceId
 
             mediaFilterRepository.addSelectedUserMediaEntry(userMedia)
             mediaFilterRepository.addMediaDataLoadingState(MediaDataLoadingModel.Loaded(instanceId))
 
-            assertThat(shownPopupChips).hasSize(1)
-            assertThat(shownPopupChips!!.first().chipId).isEqualTo(PopupChipId.MediaControl)
+            Snapshot.takeSnapshot {
+                assertThat(shownPopupChips).hasSize(1)
+                assertThat(shownPopupChips.first().chipId).isEqualTo(PopupChipId.MediaControl)
+            }
         }
 }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/featurepods/popups/ui/viewmodel/StatusBarPopupChipsViewModelKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/featurepods/popups/ui/viewmodel/StatusBarPopupChipsViewModelKosmos.kt
index 93502f365202..b876095fefe5 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/featurepods/popups/ui/viewmodel/StatusBarPopupChipsViewModelKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/featurepods/popups/ui/viewmodel/StatusBarPopupChipsViewModelKosmos.kt
@@ -17,13 +17,14 @@
 package com.android.systemui.statusbar.featurepods.popups.ui.viewmodel
 
 import com.android.systemui.kosmos.Kosmos
-import com.android.systemui.kosmos.testScope
 import com.android.systemui.statusbar.featurepods.media.ui.viewmodel.mediaControlChipViewModel
 
-val Kosmos.statusBarPopupChipsViewModel: StatusBarPopupChipsViewModel by
+private val Kosmos.statusBarPopupChipsViewModel: StatusBarPopupChipsViewModel by
+    Kosmos.Fixture { StatusBarPopupChipsViewModel(mediaControlChip = mediaControlChipViewModel) }
+
+val Kosmos.statusBarPopupChipsViewModelFactory by
     Kosmos.Fixture {
-        StatusBarPopupChipsViewModel(
-            testScope.backgroundScope,
-            mediaControlChipViewModel = mediaControlChipViewModel,
-        )
+        object : StatusBarPopupChipsViewModel.Factory {
+            override fun create(): StatusBarPopupChipsViewModel = statusBarPopupChipsViewModel
+        }
     }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/pipeline/shared/ui/viewmodel/HomeStatusBarViewModelKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/pipeline/shared/ui/viewmodel/HomeStatusBarViewModelKosmos.kt
index fbada934c9d4..a97c651ba426 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/pipeline/shared/ui/viewmodel/HomeStatusBarViewModelKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/pipeline/shared/ui/viewmodel/HomeStatusBarViewModelKosmos.kt
@@ -29,7 +29,7 @@ import com.android.systemui.shade.domain.interactor.shadeInteractor
 import com.android.systemui.statusbar.chips.sharetoapp.ui.viewmodel.shareToAppChipViewModel
 import com.android.systemui.statusbar.chips.ui.viewmodel.ongoingActivityChipsViewModel
 import com.android.systemui.statusbar.events.domain.interactor.systemStatusEventAnimationInteractor
-import com.android.systemui.statusbar.featurepods.popups.ui.viewmodel.statusBarPopupChipsViewModel
+import com.android.systemui.statusbar.featurepods.popups.ui.viewmodel.statusBarPopupChipsViewModelFactory
 import com.android.systemui.statusbar.layout.ui.viewmodel.multiDisplayStatusBarContentInsetsViewModelStore
 import com.android.systemui.statusbar.notification.domain.interactor.activeNotificationsInteractor
 import com.android.systemui.statusbar.notification.stack.domain.interactor.headsUpNotificationInteractor
@@ -59,7 +59,7 @@ var Kosmos.homeStatusBarViewModel: HomeStatusBarViewModel by
             shadeInteractor,
             shareToAppChipViewModel,
             ongoingActivityChipsViewModel,
-            statusBarPopupChipsViewModel,
+            statusBarPopupChipsViewModelFactory,
             systemStatusEventAnimationInteractor,
             multiDisplayStatusBarContentInsetsViewModelStore,
             backgroundScope,