Merge "[WPP logging] Wire logShortcutApplied" into main
diff --git a/src/com/android/customization/module/CustomizationInjector.kt b/src/com/android/customization/module/CustomizationInjector.kt
index 4bb2b72..d761598 100644
--- a/src/com/android/customization/module/CustomizationInjector.kt
+++ b/src/com/android/customization/module/CustomizationInjector.kt
@@ -21,7 +21,6 @@
 import com.android.customization.picker.clock.domain.interactor.ClockPickerInteractor
 import com.android.customization.picker.clock.ui.view.ClockViewFactory
 import com.android.customization.picker.clock.ui.viewmodel.ClockCarouselViewModel
-import com.android.customization.picker.clock.ui.viewmodel.ClockSectionViewModel
 import com.android.customization.picker.clock.ui.viewmodel.ClockSettingsViewModel
 import com.android.customization.picker.color.domain.interactor.ColorPickerInteractor
 import com.android.customization.picker.color.ui.viewmodel.ColorPickerViewModel
@@ -41,10 +40,6 @@
 
     fun getClockPickerInteractor(context: Context): ClockPickerInteractor
 
-    fun getClockSectionViewModel(
-        context: Context,
-    ): ClockSectionViewModel
-
     fun getColorPickerInteractor(
         context: Context,
         wallpaperColorsRepository: WallpaperColorsRepository,
diff --git a/src/com/android/customization/module/ThemePickerInjector.kt b/src/com/android/customization/module/ThemePickerInjector.kt
index 0c19b7b..74e28b2 100644
--- a/src/com/android/customization/module/ThemePickerInjector.kt
+++ b/src/com/android/customization/module/ThemePickerInjector.kt
@@ -15,6 +15,7 @@
  */
 package com.android.customization.module
 
+import android.app.Activity
 import android.app.UiModeManager
 import android.app.WallpaperManager
 import android.content.Context
@@ -44,7 +45,6 @@
 import com.android.customization.picker.clock.ui.view.ClockViewFactory
 import com.android.customization.picker.clock.ui.view.ClockViewFactoryImpl
 import com.android.customization.picker.clock.ui.viewmodel.ClockCarouselViewModel
-import com.android.customization.picker.clock.ui.viewmodel.ClockSectionViewModel
 import com.android.customization.picker.clock.ui.viewmodel.ClockSettingsViewModel
 import com.android.customization.picker.color.data.repository.ColorPickerRepositoryImpl
 import com.android.customization.picker.color.domain.interactor.ColorPickerInteractor
@@ -105,9 +105,8 @@
         null
     private var notificationsSnapshotRestorer: NotificationsSnapshotRestorer? = null
     private var clockPickerInteractor: ClockPickerInteractor? = null
-    private var clockSectionViewModel: ClockSectionViewModel? = null
     private var clockCarouselViewModelFactory: ClockCarouselViewModel.Factory? = null
-    private var clockViewFactories: MutableMap<Int, ClockViewFactory> = HashMap()
+    private var clockViewFactory: ClockViewFactory? = null
     private var clockPickerSnapshotRestorer: ClockPickerSnapshotRestorer? = null
     private var notificationsInteractor: NotificationsInteractor? = null
     private var notificationSectionViewModelFactory: NotificationSectionViewModel.Factory? = null
@@ -250,7 +249,7 @@
         val client = getKeyguardQuickAffordancePickerProviderClient(context)
         val appContext = context.applicationContext
         return KeyguardQuickAffordancePickerInteractor(
-            KeyguardQuickAffordancePickerRepository(client),
+            KeyguardQuickAffordancePickerRepository(client, getApplicationCoroutineScope()),
             client
         ) {
             getKeyguardQuickAffordanceSnapshotRestorer(appContext)
@@ -345,17 +344,6 @@
                 .also { clockPickerInteractor = it }
     }
 
-    override fun getClockSectionViewModel(
-        context: Context,
-    ): ClockSectionViewModel {
-        return clockSectionViewModel
-            ?: ClockSectionViewModel(
-                    context.applicationContext,
-                    getClockPickerInteractor(context.applicationContext)
-                )
-                .also { clockSectionViewModel = it }
-    }
-
     override fun getClockCarouselViewModelFactory(
         interactor: ClockPickerInteractor,
         clockViewFactory: ClockViewFactory,
@@ -367,8 +355,7 @@
     }
 
     override fun getClockViewFactory(activity: ComponentActivity): ClockViewFactory {
-        val activityHashCode = activity.hashCode()
-        return clockViewFactories[activityHashCode]
+        return clockViewFactory
             ?: ClockViewFactoryImpl(
                     activity.applicationContext,
                     ScreenSizeCalculator.getInstance()
@@ -377,13 +364,13 @@
                     getClockRegistry(activity.applicationContext),
                 )
                 .also {
-                    clockViewFactories[activityHashCode] = it
+                    clockViewFactory = it
                     activity.lifecycle.addObserver(
                         object : DefaultLifecycleObserver {
                             override fun onDestroy(owner: LifecycleOwner) {
                                 super.onDestroy(owner)
-                                clockViewFactories[activityHashCode]?.onDestroy()
-                                clockViewFactories.remove(activityHashCode)
+                                if ((owner as Activity).isChangingConfigurations()) return
+                                clockViewFactory?.onDestroy()
                             }
                         }
                     )
diff --git a/src/com/android/customization/module/logging/AppSessionId.kt b/src/com/android/customization/module/logging/AppSessionId.kt
new file mode 100644
index 0000000..c831f22
--- /dev/null
+++ b/src/com/android/customization/module/logging/AppSessionId.kt
@@ -0,0 +1,43 @@
+/*
+ * Copyright (C) 2023 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.android.customization.module.logging
+
+import com.android.internal.logging.InstanceId
+import com.android.internal.logging.InstanceIdSequence
+import javax.inject.Inject
+import javax.inject.Singleton
+
+@Singleton
+class AppSessionId @Inject constructor() {
+
+    private var sessionId: InstanceId = newInstanceId()
+
+    fun createNewId(): AppSessionId {
+        sessionId = newInstanceId()
+        return this
+    }
+
+    fun getId(): Int {
+        return sessionId.hashCode()
+    }
+
+    private fun newInstanceId(): InstanceId = InstanceIdSequence(INSTANCE_ID_MAX).newInstanceId()
+
+    companion object {
+        // At most 20 bits: ~1m possibilities, ~0.5% probability of collision in 100 values
+        private const val INSTANCE_ID_MAX = 1 shl 20
+    }
+}
diff --git a/src/com/android/customization/module/logging/SysUiStatsLogger.kt b/src/com/android/customization/module/logging/SysUiStatsLogger.kt
index 14f0be6..6c55df8 100644
--- a/src/com/android/customization/module/logging/SysUiStatsLogger.kt
+++ b/src/com/android/customization/module/logging/SysUiStatsLogger.kt
@@ -53,7 +53,7 @@
     private var colorVariant = 0
     private var timeElapsedMillis = 0L
     private var effectResultCode = -1
-    private var sessionId = 0
+    private var appSessionId = 0
     private var setWallpaperEntryPoint =
         STYLE_UICHANGED__SET_WALLPAPER_ENTRY_POINT__SET_WALLPAPER_ENTRY_POINT_UNSPECIFIED
     private var wallpaperDestination =
@@ -133,7 +133,7 @@
         this.effectResultCode = effectResultCode
     }
 
-    fun setSessionId(sessionId: Int) = apply { this.sessionId = sessionId }
+    fun setAppSessionId(sessionId: Int) = apply { this.appSessionId = sessionId }
 
     fun setSetWallpaperEntryPoint(@SetWallpaperEntryPoint setWallpaperEntryPoint: Int) = apply {
         this.setWallpaperEntryPoint = setWallpaperEntryPoint
@@ -180,7 +180,7 @@
             colorVariant,
             timeElapsedMillis,
             effectResultCode,
-            sessionId,
+            appSessionId,
             setWallpaperEntryPoint,
             wallpaperDestination,
             colorSource,
diff --git a/src/com/android/customization/module/logging/ThemesUserEventLoggerImpl.kt b/src/com/android/customization/module/logging/ThemesUserEventLoggerImpl.kt
index 8998dee..e4b8c8b 100644
--- a/src/com/android/customization/module/logging/ThemesUserEventLoggerImpl.kt
+++ b/src/com/android/customization/module/logging/ThemesUserEventLoggerImpl.kt
@@ -40,6 +40,7 @@
 @Inject
 constructor(
     private val preferences: WallpaperPreferences,
+    private val appSessionId: AppSessionId,
 ) : ThemesUserEventLogger {
 
     override fun logSnapshot() {
@@ -54,6 +55,7 @@
 
     override fun logAppLaunched(launchSource: Intent) {
         SysUiStatsLogger(StyleEnums.APP_LAUNCHED)
+            .setAppSessionId(appSessionId.createNewId().getId())
             .setLaunchedPreference(launchSource.getAppLaunchSource())
             .log()
     }
@@ -70,6 +72,7 @@
         val isHomeWallpaperSet = destination == DEST_HOME_SCREEN || destination == DEST_BOTH
         val isLockWallpaperSet = destination == DEST_LOCK_SCREEN || destination == DEST_BOTH
         SysUiStatsLogger(StyleEnums.WALLPAPER_APPLIED)
+            .setAppSessionId(appSessionId.getId())
             .setWallpaperCategoryHash(if (isHomeWallpaperSet) categoryHash else 0)
             .setWallpaperIdHash(if (isHomeWallpaperSet) wallpaperIdHash else 0)
             .setLockWallpaperCategoryHash(if (isLockWallpaperSet) categoryHash else 0)
@@ -87,6 +90,7 @@
         resultCode: Int
     ) {
         SysUiStatsLogger(StyleEnums.WALLPAPER_EFFECT_APPLIED)
+            .setAppSessionId(appSessionId.getId())
             .setEffectPreference(status)
             .setEffectIdHash(getIdHashCode(effect))
             .setTimeElapsed(timeElapsedMillis)
@@ -96,6 +100,7 @@
 
     override fun logEffectProbe(effect: String, @EffectStatus status: Int) {
         SysUiStatsLogger(StyleEnums.WALLPAPER_EFFECT_PROBE)
+            .setAppSessionId(appSessionId.getId())
             .setEffectPreference(status)
             .setEffectIdHash(getIdHashCode(effect))
             .log()
@@ -107,6 +112,7 @@
         timeElapsedMillis: Long
     ) {
         SysUiStatsLogger(StyleEnums.WALLPAPER_EFFECT_FG_DOWNLOAD)
+            .setAppSessionId(appSessionId.getId())
             .setEffectPreference(status)
             .setEffectIdHash(getIdHashCode(effect))
             .setTimeElapsed(timeElapsedMillis)
@@ -114,11 +120,11 @@
     }
 
     override fun logResetApplied() {
-        SysUiStatsLogger(StyleEnums.RESET_APPLIED).log()
+        SysUiStatsLogger(StyleEnums.RESET_APPLIED).setAppSessionId(appSessionId.getId()).log()
     }
 
     override fun logWallpaperExploreButtonClicked() {
-        SysUiStatsLogger(StyleEnums.WALLPAPER_EXPLORE).log()
+        SysUiStatsLogger(StyleEnums.WALLPAPER_EXPLORE).setAppSessionId(appSessionId.getId()).log()
     }
 
     override fun logThemeColorApplied(
@@ -127,6 +133,7 @@
         seedColor: Int,
     ) {
         SysUiStatsLogger(StyleEnums.THEME_COLOR_APPLIED)
+            .setAppSessionId(appSessionId.getId())
             .setColorSource(source)
             .setColorVariant(variant)
             .setSeedColor(seedColor)
@@ -134,40 +141,60 @@
     }
 
     override fun logGridApplied(grid: GridOption) {
-        SysUiStatsLogger(StyleEnums.GRID_APPLIED).setLauncherGrid(grid.getLauncherGridInt()).log()
+        SysUiStatsLogger(StyleEnums.GRID_APPLIED)
+            .setAppSessionId(appSessionId.getId())
+            .setLauncherGrid(grid.getLauncherGridInt())
+            .log()
     }
 
     override fun logClockApplied(clockId: String) {
-        SysUiStatsLogger(StyleEnums.CLOCK_APPLIED).setClockPackageHash(getIdHashCode(clockId)).log()
+        SysUiStatsLogger(StyleEnums.CLOCK_APPLIED)
+            .setAppSessionId(appSessionId.getId())
+            .setClockPackageHash(getIdHashCode(clockId))
+            .log()
     }
 
     override fun logClockColorApplied(seedColor: Int) {
-        SysUiStatsLogger(StyleEnums.CLOCK_COLOR_APPLIED).setSeedColor(seedColor).log()
+        SysUiStatsLogger(StyleEnums.CLOCK_COLOR_APPLIED)
+            .setAppSessionId(appSessionId.getId())
+            .setSeedColor(seedColor)
+            .log()
     }
 
     override fun logClockSizeApplied(@ClockSize clockSize: Int) {
-        SysUiStatsLogger(StyleEnums.CLOCK_SIZE_APPLIED).setClockSize(clockSize).log()
+        SysUiStatsLogger(StyleEnums.CLOCK_SIZE_APPLIED)
+            .setAppSessionId(appSessionId.getId())
+            .setClockSize(clockSize)
+            .log()
     }
 
     override fun logThemedIconApplied(useThemeIcon: Boolean) {
-        SysUiStatsLogger(StyleEnums.THEMED_ICON_APPLIED).setToggleOn(useThemeIcon).log()
+        SysUiStatsLogger(StyleEnums.THEMED_ICON_APPLIED)
+            .setAppSessionId(appSessionId.getId())
+            .setToggleOn(useThemeIcon)
+            .log()
     }
 
     override fun logLockScreenNotificationApplied(showLockScreenNotifications: Boolean) {
         SysUiStatsLogger(StyleEnums.LOCK_SCREEN_NOTIFICATION_APPLIED)
+            .setAppSessionId(appSessionId.getId())
             .setToggleOn(showLockScreenNotifications)
             .log()
     }
 
     override fun logShortcutApplied(shortcut: String, shortcutSlotId: String) {
         SysUiStatsLogger(StyleEnums.SHORTCUT_APPLIED)
+            .setAppSessionId(appSessionId.getId())
             .setShortcut(shortcut)
             .setShortcutSlotId(shortcutSlotId)
             .log()
     }
 
     override fun logDarkThemeApplied(useDarkTheme: Boolean) {
-        SysUiStatsLogger(StyleEnums.DARK_THEME_APPLIED).setToggleOn(useDarkTheme).log()
+        SysUiStatsLogger(StyleEnums.DARK_THEME_APPLIED)
+            .setAppSessionId(appSessionId.getId())
+            .setToggleOn(useDarkTheme)
+            .log()
     }
 
     /**
diff --git a/src/com/android/customization/picker/clock/ui/binder/ClockSectionViewBinder.kt b/src/com/android/customization/picker/clock/ui/binder/ClockSectionViewBinder.kt
deleted file mode 100644
index 7dc0d0c..0000000
--- a/src/com/android/customization/picker/clock/ui/binder/ClockSectionViewBinder.kt
+++ /dev/null
@@ -1,53 +0,0 @@
-/*
- * Copyright (C) 2023 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- *
- */
-
-package com.android.customization.picker.clock.ui.binder
-
-import android.view.View
-import android.widget.TextView
-import androidx.lifecycle.Lifecycle
-import androidx.lifecycle.LifecycleOwner
-import androidx.lifecycle.lifecycleScope
-import androidx.lifecycle.repeatOnLifecycle
-import com.android.customization.picker.clock.ui.viewmodel.ClockSectionViewModel
-import com.android.wallpaper.R
-import kotlinx.coroutines.flow.collectLatest
-import kotlinx.coroutines.launch
-
-object ClockSectionViewBinder {
-    fun bind(
-        view: View,
-        viewModel: ClockSectionViewModel,
-        lifecycleOwner: LifecycleOwner,
-        onClicked: () -> Unit,
-    ) {
-        view.setOnClickListener { onClicked() }
-
-        val selectedClockColorAndSize: TextView =
-            view.requireViewById(R.id.selected_clock_color_and_size)
-
-        lifecycleOwner.lifecycleScope.launch {
-            lifecycleOwner.repeatOnLifecycle(Lifecycle.State.STARTED) {
-                launch {
-                    viewModel.selectedClockColorAndSizeText.collectLatest {
-                        selectedClockColorAndSize.text = it
-                    }
-                }
-            }
-        }
-    }
-}
diff --git a/src/com/android/customization/picker/clock/ui/section/ClockSectionController.kt b/src/com/android/customization/picker/clock/ui/section/ClockSectionController.kt
deleted file mode 100644
index b47c243..0000000
--- a/src/com/android/customization/picker/clock/ui/section/ClockSectionController.kt
+++ /dev/null
@@ -1,62 +0,0 @@
-/*
- * Copyright (C) 2022 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package com.android.customization.picker.clock.ui.section
-
-import android.content.Context
-import android.view.LayoutInflater
-import androidx.lifecycle.LifecycleOwner
-import androidx.lifecycle.lifecycleScope
-import com.android.customization.picker.clock.ui.binder.ClockSectionViewBinder
-import com.android.customization.picker.clock.ui.fragment.ClockSettingsFragment
-import com.android.customization.picker.clock.ui.view.ClockSectionView
-import com.android.customization.picker.clock.ui.viewmodel.ClockSectionViewModel
-import com.android.wallpaper.R
-import com.android.wallpaper.config.BaseFlags
-import com.android.wallpaper.model.CustomizationSectionController
-import com.android.wallpaper.model.CustomizationSectionController.CustomizationSectionNavigationController
-import kotlinx.coroutines.launch
-
-/** A [CustomizationSectionController] for clock customization. */
-class ClockSectionController(
-    private val navigationController: CustomizationSectionNavigationController,
-    private val lifecycleOwner: LifecycleOwner,
-    private val flag: BaseFlags,
-    private val viewModel: ClockSectionViewModel,
-) : CustomizationSectionController<ClockSectionView> {
-
-    override fun isAvailable(context: Context): Boolean {
-        return flag.isCustomClocksEnabled(context!!)
-    }
-
-    override fun createView(context: Context): ClockSectionView {
-        val view =
-            LayoutInflater.from(context)
-                .inflate(
-                    R.layout.clock_section_view,
-                    null,
-                ) as ClockSectionView
-        lifecycleOwner.lifecycleScope.launch {
-            ClockSectionViewBinder.bind(
-                view = view,
-                viewModel = viewModel,
-                lifecycleOwner = lifecycleOwner
-            ) {
-                navigationController.navigateTo(ClockSettingsFragment())
-            }
-        }
-        return view
-    }
-}
diff --git a/src/com/android/customization/picker/clock/ui/viewmodel/ClockSectionViewModel.kt b/src/com/android/customization/picker/clock/ui/viewmodel/ClockSectionViewModel.kt
deleted file mode 100644
index 8a65522..0000000
--- a/src/com/android/customization/picker/clock/ui/viewmodel/ClockSectionViewModel.kt
+++ /dev/null
@@ -1,51 +0,0 @@
-/*
- * Copyright (C) 2023 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- *
- */
-package com.android.customization.picker.clock.ui.viewmodel
-
-import android.content.Context
-import com.android.customization.picker.clock.domain.interactor.ClockPickerInteractor
-import com.android.customization.picker.clock.shared.ClockSize
-import com.android.wallpaper.R
-import java.util.Locale
-import kotlinx.coroutines.flow.Flow
-import kotlinx.coroutines.flow.combine
-import kotlinx.coroutines.flow.map
-
-/** View model for the clock section view on the lockscreen customization surface. */
-class ClockSectionViewModel(context: Context, interactor: ClockPickerInteractor) {
-    val appContext: Context = context.applicationContext
-    val clockColorMap: Map<String, ClockColorViewModel> =
-        ClockColorViewModel.getPresetColorMap(appContext.resources)
-    val selectedClockColorAndSizeText: Flow<String> =
-        combine(interactor.selectedColorId, interactor.selectedClockSize, ::Pair).map {
-            (selectedColorId, selectedClockSize) ->
-            val colorText =
-                clockColorMap[selectedColorId]?.colorName
-                    ?: appContext.getString(R.string.default_theme_title)
-            val sizeText =
-                when (selectedClockSize) {
-                    ClockSize.SMALL -> appContext.getString(R.string.clock_size_small)
-                    ClockSize.DYNAMIC -> appContext.getString(R.string.clock_size_dynamic)
-                }
-            appContext
-                .getString(R.string.clock_color_and_size_description, colorText, sizeText)
-                .lowercase()
-                .replaceFirstChar {
-                    if (it.isLowerCase()) it.titlecase(Locale.getDefault()) else it.toString()
-                }
-        }
-}
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 10473a2..6bfe348 100644
--- a/src/com/android/customization/picker/quickaffordance/data/repository/KeyguardQuickAffordancePickerRepository.kt
+++ b/src/com/android/customization/picker/quickaffordance/data/repository/KeyguardQuickAffordancePickerRepository.kt
@@ -21,8 +21,11 @@
 import com.android.customization.picker.quickaffordance.shared.model.KeyguardQuickAffordancePickerSelectionModel as SelectionModel
 import com.android.customization.picker.quickaffordance.shared.model.KeyguardQuickAffordancePickerSlotModel as SlotModel
 import com.android.systemui.shared.customization.data.content.CustomizationProviderClient as Client
+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
@@ -30,6 +33,7 @@
  */
 class KeyguardQuickAffordancePickerRepository(
     private val client: Client,
+    private val scope: CoroutineScope
 ) {
     /** List of slots available on the device. */
     val slots: Flow<List<SlotModel>> =
@@ -37,15 +41,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/common/Android.bp b/tests/common/Android.bp
index 0961886..9215ba6 100644
--- a/tests/common/Android.bp
+++ b/tests/common/Android.bp
@@ -32,7 +32,7 @@
         "WallpaperPicker2TestLib",
         "androidx.annotation_annotation",
         "kotlinx_coroutines_test",
-        "truth-prebuilt",
+        "truth",
     ],
 
     platform_apis: true,
diff --git a/tests/common/src/com/android/customization/testing/TestCustomizationInjector.kt b/tests/common/src/com/android/customization/testing/TestCustomizationInjector.kt
index d1d9af8..065a1f0 100644
--- a/tests/common/src/com/android/customization/testing/TestCustomizationInjector.kt
+++ b/tests/common/src/com/android/customization/testing/TestCustomizationInjector.kt
@@ -9,7 +9,6 @@
 import com.android.customization.picker.clock.domain.interactor.ClockPickerInteractor
 import com.android.customization.picker.clock.ui.view.ClockViewFactory
 import com.android.customization.picker.clock.ui.viewmodel.ClockCarouselViewModel
-import com.android.customization.picker.clock.ui.viewmodel.ClockSectionViewModel
 import com.android.customization.picker.clock.ui.viewmodel.ClockSettingsViewModel
 import com.android.customization.picker.color.domain.interactor.ColorPickerInteractor
 import com.android.customization.picker.color.ui.viewmodel.ColorPickerViewModel
@@ -50,10 +49,6 @@
         throw UnsupportedOperationException("not implemented")
     }
 
-    override fun getClockSectionViewModel(context: Context): ClockSectionViewModel {
-        throw UnsupportedOperationException("not implemented")
-    }
-
     override fun getColorPickerInteractor(
         context: Context,
         wallpaperColorsRepository: WallpaperColorsRepository,
diff --git a/tests/robotests/Android.bp b/tests/robotests/Android.bp
index ed684f4..c34ed2c 100644
--- a/tests/robotests/Android.bp
+++ b/tests/robotests/Android.bp
@@ -1,4 +1,3 @@
-
 package {
     default_applicable_licenses: ["Android-Apache-2.0"],
 }
@@ -17,7 +16,7 @@
         "androidx.test.rules",
         "junit",
         "kotlinx_coroutines_test",
-        "truth-prebuilt",
+        "truth",
     ],
     libs: [
         "androidx.test.core",
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 8a5d582..8687b30 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
@@ -20,6 +20,7 @@
 import androidx.test.filters.SmallTest
 import com.android.customization.picker.quickaffordance.data.repository.KeyguardQuickAffordancePickerRepository
 import com.android.systemui.shared.customization.data.content.FakeCustomizationProviderClient
+import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.test.TestScope
@@ -28,6 +29,7 @@
 import kotlinx.coroutines.test.setMain
 import org.junit.After
 import org.junit.Before
+import org.junit.Test
 import org.junit.runner.RunWith
 import org.robolectric.RobolectricTestRunner
 
@@ -51,9 +53,16 @@
         underTest =
             KeyguardQuickAffordancePickerRepository(
                 client = client,
+                scope = testScope.backgroundScope,
             )
     }
 
+    // We need at least one test to prevent Studio errors
+    @Test
+    fun creationSucceeds() {
+        assertThat(underTest).isNotNull()
+    }
+
     @After
     fun tearDown() {
         Dispatchers.resetMain()
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 11098ec..bf53f61 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
@@ -62,6 +62,7 @@
                 repository =
                     KeyguardQuickAffordancePickerRepository(
                         client = client,
+                        scope = testScope.backgroundScope,
                     ),
                 client = client,
                 snapshotRestorer = {
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 a33c407..a82cc00 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
@@ -89,6 +89,7 @@
                 repository =
                     KeyguardQuickAffordancePickerRepository(
                         client = client,
+                        scope = testScope.backgroundScope,
                     ),
                 client = client,
                 snapshotRestorer = {
diff --git a/tests/robotests/src/com/android/customization/picker/clock/ui/viewmodel/ClockSectionViewModelTest.kt b/tests/robotests/src/com/android/customization/picker/clock/ui/viewmodel/ClockSectionViewModelTest.kt
deleted file mode 100644
index 19a704c..0000000
--- a/tests/robotests/src/com/android/customization/picker/clock/ui/viewmodel/ClockSectionViewModelTest.kt
+++ /dev/null
@@ -1,95 +0,0 @@
-/*
- * Copyright (C) 2023 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package com.android.customization.picker.clock.ui.viewmodel
-
-import androidx.test.filters.SmallTest
-import androidx.test.platform.app.InstrumentationRegistry
-import com.android.customization.picker.clock.data.repository.FakeClockPickerRepository
-import com.android.customization.picker.clock.domain.interactor.ClockPickerInteractor
-import com.android.customization.picker.clock.domain.interactor.ClockPickerSnapshotRestorer
-import com.android.customization.picker.clock.shared.ClockSize
-import com.android.customization.picker.clock.shared.model.ClockMetadataModel
-import com.android.wallpaper.testing.FakeSnapshotStore
-import com.android.wallpaper.testing.collectLastValue
-import com.google.common.truth.Truth.assertThat
-import kotlinx.coroutines.Dispatchers
-import kotlinx.coroutines.ExperimentalCoroutinesApi
-import kotlinx.coroutines.runBlocking
-import kotlinx.coroutines.test.StandardTestDispatcher
-import kotlinx.coroutines.test.resetMain
-import kotlinx.coroutines.test.runTest
-import kotlinx.coroutines.test.setMain
-import org.junit.After
-import org.junit.Before
-import org.junit.Test
-import org.junit.runner.RunWith
-import org.robolectric.RobolectricTestRunner
-
-@OptIn(ExperimentalCoroutinesApi::class)
-@SmallTest
-@RunWith(RobolectricTestRunner::class)
-class ClockSectionViewModelTest {
-
-    private lateinit var clockColorMap: Map<String, ClockColorViewModel>
-    private lateinit var interactor: ClockPickerInteractor
-    private lateinit var underTest: ClockSectionViewModel
-
-    @Before
-    fun setUp() {
-        val testDispatcher = StandardTestDispatcher()
-        Dispatchers.setMain(testDispatcher)
-        val context = InstrumentationRegistry.getInstrumentation().targetContext
-        clockColorMap = ClockColorViewModel.getPresetColorMap(context.resources)
-        interactor =
-            ClockPickerInteractor(
-                repository = FakeClockPickerRepository(),
-                snapshotRestorer = {
-                    ClockPickerSnapshotRestorer(interactor = interactor).apply {
-                        runBlocking { setUpSnapshotRestorer(store = FakeSnapshotStore()) }
-                    }
-                },
-            )
-        underTest =
-            ClockSectionViewModel(
-                context,
-                interactor,
-            )
-    }
-
-    @After
-    fun tearDown() {
-        Dispatchers.resetMain()
-    }
-
-    @Test
-    fun setSelectedClock() = runTest {
-        val colorGrey = clockColorMap.values.first()
-        val observedSelectedClockColorAndSizeText =
-            collectLastValue(underTest.selectedClockColorAndSizeText)
-
-        interactor.setClockColor(
-            colorGrey.colorId,
-            ClockMetadataModel.DEFAULT_COLOR_TONE_PROGRESS,
-            ClockSettingsViewModel.blendColorWithTone(
-                colorGrey.color,
-                colorGrey.getColorTone(ClockMetadataModel.DEFAULT_COLOR_TONE_PROGRESS),
-            )
-        )
-        interactor.setClockSize(ClockSize.DYNAMIC)
-
-        assertThat(observedSelectedClockColorAndSizeText()).isEqualTo("Grey, dynamic")
-    }
-}