Cache affordances and selection picker side

Bug: 302425391
Test: manually verified that quick affordances work as expected

Change-Id: I24a05119721847243afbe7e7ba326cdd328c870b
Merged-In: I24a05119721847243afbe7e7ba326cdd328c870b
diff --git a/src/com/android/customization/module/ThemePickerInjector.kt b/src/com/android/customization/module/ThemePickerInjector.kt
index a69f2ad..f31af44 100644
--- a/src/com/android/customization/module/ThemePickerInjector.kt
+++ b/src/com/android/customization/module/ThemePickerInjector.kt
@@ -267,7 +267,12 @@
         val client = getKeyguardQuickAffordancePickerProviderClient(context)
         val appContext = context.applicationContext
         return KeyguardQuickAffordancePickerInteractor(
-            KeyguardQuickAffordancePickerRepository(client, getFlags(), appContext),
+            KeyguardQuickAffordancePickerRepository(
+                client,
+                getApplicationCoroutineScope(),
+                getFlags(),
+                appContext
+            ),
             client
         ) {
             getKeyguardQuickAffordanceSnapshotRestorer(appContext)
diff --git a/src/com/android/customization/picker/quickaffordance/data/repository/KeyguardQuickAffordancePickerRepository.kt b/src/com/android/customization/picker/quickaffordance/data/repository/KeyguardQuickAffordancePickerRepository.kt
index 080d70e..80c7b4e 100644
--- a/src/com/android/customization/picker/quickaffordance/data/repository/KeyguardQuickAffordancePickerRepository.kt
+++ b/src/com/android/customization/picker/quickaffordance/data/repository/KeyguardQuickAffordancePickerRepository.kt
@@ -23,8 +23,11 @@
 import com.android.customization.picker.quickaffordance.shared.model.KeyguardQuickAffordancePickerSlotModel as SlotModel
 import com.android.systemui.shared.customization.data.content.CustomizationProviderClient as Client
 import com.android.wallpaper.config.BaseFlags
+import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.SharingStarted
 import kotlinx.coroutines.flow.map
+import kotlinx.coroutines.flow.shareIn
 
 /**
  * Abstracts access to application state related to functionality for selecting, picking, or setting
@@ -32,6 +35,7 @@
  */
 class KeyguardQuickAffordancePickerRepository(
     private val client: Client,
+    private val scope: CoroutineScope,
     private val flags: BaseFlags,
     private val context: Context
 ) {
@@ -46,15 +50,17 @@
 
     /** List of all available quick affordances. */
     val affordances: Flow<List<AffordanceModel>> =
-        client.observeAffordances().map { affordances ->
-            affordances.map { affordance -> affordance.toModel() }
-        }
+        client
+            .observeAffordances()
+            .map { affordances -> affordances.map { affordance -> affordance.toModel() } }
+            .shareIn(scope, replay = 1, started = SharingStarted.Lazily)
 
     /** List of slot-affordance pairs, modeling what the user has currently chosen for each slot. */
     val selections: Flow<List<SelectionModel>> =
-        client.observeSelections().map { selections ->
-            selections.map { selection -> selection.toModel() }
-        }
+        client
+            .observeSelections()
+            .map { selections -> selections.map { selection -> selection.toModel() } }
+            .shareIn(scope, replay = 1, started = SharingStarted.Lazily)
 
     private fun Client.Slot.toModel(): SlotModel {
         return SlotModel(
diff --git a/tests/robotests/src/com/android/customization/model/picker/quickaffordance/data/repository/KeyguardQuickAffordancePickerRepositoryTest.kt b/tests/robotests/src/com/android/customization/model/picker/quickaffordance/data/repository/KeyguardQuickAffordancePickerRepositoryTest.kt
index 38c43ca..5f9d7fb 100644
--- a/tests/robotests/src/com/android/customization/model/picker/quickaffordance/data/repository/KeyguardQuickAffordancePickerRepositoryTest.kt
+++ b/tests/robotests/src/com/android/customization/model/picker/quickaffordance/data/repository/KeyguardQuickAffordancePickerRepositoryTest.kt
@@ -58,6 +58,7 @@
         underTest =
             KeyguardQuickAffordancePickerRepository(
                 client = client,
+                scope = testScope.backgroundScope,
                 flags =
                     object : BaseFlags() {
                         override fun getCachedFlags(
diff --git a/tests/robotests/src/com/android/customization/model/picker/quickaffordance/domain/interactor/KeyguardQuickAffordancePickerInteractorTest.kt b/tests/robotests/src/com/android/customization/model/picker/quickaffordance/domain/interactor/KeyguardQuickAffordancePickerInteractorTest.kt
index 2668064..ac60282 100644
--- a/tests/robotests/src/com/android/customization/model/picker/quickaffordance/domain/interactor/KeyguardQuickAffordancePickerInteractorTest.kt
+++ b/tests/robotests/src/com/android/customization/model/picker/quickaffordance/domain/interactor/KeyguardQuickAffordancePickerInteractorTest.kt
@@ -65,6 +65,7 @@
                 repository =
                     KeyguardQuickAffordancePickerRepository(
                         client = client,
+                        scope = testScope.backgroundScope,
                         flags = testFlags,
                         context = ApplicationProvider.getApplicationContext(),
                     ),
diff --git a/tests/robotests/src/com/android/customization/model/picker/quickaffordance/ui/viewmodel/KeyguardQuickAffordancePickerViewModelTest.kt b/tests/robotests/src/com/android/customization/model/picker/quickaffordance/ui/viewmodel/KeyguardQuickAffordancePickerViewModelTest.kt
index 8dc8534..93a84e6 100644
--- a/tests/robotests/src/com/android/customization/model/picker/quickaffordance/ui/viewmodel/KeyguardQuickAffordancePickerViewModelTest.kt
+++ b/tests/robotests/src/com/android/customization/model/picker/quickaffordance/ui/viewmodel/KeyguardQuickAffordancePickerViewModelTest.kt
@@ -87,6 +87,7 @@
                 repository =
                     KeyguardQuickAffordancePickerRepository(
                         client = client,
+                        scope = testScope.backgroundScope,
                         flags = injector.getFlags(),
                         context = context,
                     ),