diff --git a/src/com/android/customization/module/ThemePickerInjector.kt b/src/com/android/customization/module/ThemePickerInjector.kt
index 8ab541e..10b71a5 100644
--- a/src/com/android/customization/module/ThemePickerInjector.kt
+++ b/src/com/android/customization/module/ThemePickerInjector.kt
@@ -36,8 +36,8 @@
 import com.android.customization.model.themedicon.data.repository.ThemeIconRepository
 import com.android.customization.model.themedicon.domain.interactor.ThemedIconInteractor
 import com.android.customization.model.themedicon.domain.interactor.ThemedIconSnapshotRestorer
-import com.android.customization.module.logging.StatsLogUserEventLogger
 import com.android.customization.module.logging.ThemesUserEventLogger
+import com.android.customization.module.logging.ThemesUserEventLoggerImpl
 import com.android.customization.picker.clock.data.repository.ClockPickerRepositoryImpl
 import com.android.customization.picker.clock.data.repository.ClockRegistryProvider
 import com.android.customization.picker.clock.domain.interactor.ClockPickerInteractor
@@ -168,7 +168,7 @@
     @Synchronized
     override fun getUserEventLogger(context: Context): ThemesUserEventLogger {
         return userEventLogger as? ThemesUserEventLogger
-            ?: StatsLogUserEventLogger(getPreferences(context.applicationContext)).also {
+            ?: ThemesUserEventLoggerImpl(getPreferences(context.applicationContext)).also {
                 userEventLogger = it
             }
     }
diff --git a/src/com/android/customization/module/SysUiStatsLogger.kt b/src/com/android/customization/module/logging/SysUiStatsLogger.kt
similarity index 80%
rename from src/com/android/customization/module/SysUiStatsLogger.kt
rename to src/com/android/customization/module/logging/SysUiStatsLogger.kt
index 142bd22..14f0be6 100644
--- a/src/com/android/customization/module/SysUiStatsLogger.kt
+++ b/src/com/android/customization/module/logging/SysUiStatsLogger.kt
@@ -13,7 +13,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package com.android.customization.module
+package com.android.customization.module.logging
 
 import com.android.systemui.shared.system.SysUiStatsLog
 import com.android.systemui.shared.system.SysUiStatsLog.STYLE_UICHANGED__CLOCK_SIZE__CLOCK_SIZE_UNSPECIFIED
@@ -38,9 +38,11 @@
     private var wallpaperCategoryHash = 0
     private var wallpaperIdHash = 0
     private var colorPreference = 0
-    private var locationPreference = STYLE_UICHANGED__LOCATION_PREFERENCE__LOCATION_PREFERENCE_UNSPECIFIED
+    private var locationPreference =
+        STYLE_UICHANGED__LOCATION_PREFERENCE__LOCATION_PREFERENCE_UNSPECIFIED
     private var datePreference = STYLE_UICHANGED__DATE_PREFERENCE__DATE_PREFERENCE_UNSPECIFIED
-    private var launchedPreference = STYLE_UICHANGED__LAUNCHED_PREFERENCE__LAUNCHED_PREFERENCE_UNSPECIFIED
+    private var launchedPreference =
+        STYLE_UICHANGED__LAUNCHED_PREFERENCE__LAUNCHED_PREFERENCE_UNSPECIFIED
     private var effectPreference = STYLE_UICHANGED__EFFECT_PREFERENCE__EFFECT_PREFERENCE_UNSPECIFIED
     private var effectIdHash = 0
     private var lockWallpaperCategoryHash = 0
@@ -52,8 +54,10 @@
     private var timeElapsedMillis = 0L
     private var effectResultCode = -1
     private var sessionId = 0
-    private var setWallpaperEntryPoint = STYLE_UICHANGED__SET_WALLPAPER_ENTRY_POINT__SET_WALLPAPER_ENTRY_POINT_UNSPECIFIED
-    private var wallpaperDestination = STYLE_UICHANGED__WALLPAPER_DESTINATION__WALLPAPER_DESTINATION_UNSPECIFIED
+    private var setWallpaperEntryPoint =
+        STYLE_UICHANGED__SET_WALLPAPER_ENTRY_POINT__SET_WALLPAPER_ENTRY_POINT_UNSPECIFIED
+    private var wallpaperDestination =
+        STYLE_UICHANGED__WALLPAPER_DESTINATION__WALLPAPER_DESTINATION_UNSPECIFIED
     private var colorSource = STYLE_UICHANGED__COLOR_SOURCE__COLOR_SOURCE_UNSPECIFIED
     private var seedColor = 0
     private var clockSize = STYLE_UICHANGED__CLOCK_SIZE__CLOCK_SIZE_UNSPECIFIED
@@ -65,9 +69,7 @@
         this.colorPackageHash = colorPackageHash
     }
 
-    fun setFontPackageHash(fontPackageHash: Int) = apply {
-        this.fontPackageHash = fontPackageHash
-    }
+    fun setFontPackageHash(fontPackageHash: Int) = apply { this.fontPackageHash = fontPackageHash }
 
     fun setShapePackageHash(shapePackageHash: Int) = apply {
         this.shapePackageHash = shapePackageHash
@@ -83,13 +85,9 @@
         this.wallpaperCategoryHash = wallpaperCategoryHash
     }
 
-    fun setWallpaperIdHash(wallpaperIdHash: Int) = apply {
-        this.wallpaperIdHash = wallpaperIdHash
-    }
+    fun setWallpaperIdHash(wallpaperIdHash: Int) = apply { this.wallpaperIdHash = wallpaperIdHash }
 
-    fun setColorPreference(colorPreference: Int) = apply {
-        this.colorPreference = colorPreference
-    }
+    fun setColorPreference(colorPreference: Int) = apply { this.colorPreference = colorPreference }
 
     fun setLocationPreference(locationPreference: Int) = apply {
         this.locationPreference = locationPreference
@@ -128,16 +126,14 @@
     fun setColorVariant(colorVariant: Int) = apply { this.colorVariant = colorVariant }
 
     fun setTimeElapsed(timeElapsedMillis: Long) = apply {
-      this.timeElapsedMillis = timeElapsedMillis
+        this.timeElapsedMillis = timeElapsedMillis
     }
 
     fun setEffectResultCode(effectResultCode: Int) = apply {
         this.effectResultCode = effectResultCode
     }
 
-    fun setSessionId(sessionId: Int) = apply {
-        this.sessionId = sessionId
-    }
+    fun setSessionId(sessionId: Int) = apply { this.sessionId = sessionId }
 
     fun setSetWallpaperEntryPoint(@SetWallpaperEntryPoint setWallpaperEntryPoint: Int) = apply {
         this.setWallpaperEntryPoint = setWallpaperEntryPoint
@@ -147,29 +143,17 @@
         this.wallpaperDestination = wallpaperDestination
     }
 
-    fun setColorSource(colorSource: Int) = apply {
-        this.colorSource = colorSource
-    }
+    fun setColorSource(colorSource: Int) = apply { this.colorSource = colorSource }
 
-    fun setSeedColor(seedColor: Int) = apply {
-        this.seedColor = seedColor
-    }
+    fun setSeedColor(seedColor: Int) = apply { this.seedColor = seedColor }
 
-    fun setClockSize(clockSize: Int) = apply {
-        this.clockSize = clockSize
-    }
+    fun setClockSize(clockSize: Int) = apply { this.clockSize = clockSize }
 
-    fun setToggleOn(toggleOn: Boolean) = apply {
-        this.toggleOn = toggleOn
-    }
+    fun setToggleOn(toggleOn: Boolean) = apply { this.toggleOn = toggleOn }
 
-    fun setShortcut(shortcut: String) = apply {
-        this.shortcut = shortcut
-    }
+    fun setShortcut(shortcut: String) = apply { this.shortcut = shortcut }
 
-    fun setShortcutSlotId(shortcutSlotId: String) = apply {
-        this.shortcutSlotId = shortcutSlotId
-    }
+    fun setShortcutSlotId(shortcutSlotId: String) = apply { this.shortcutSlotId = shortcutSlotId }
 
     fun log() {
         SysUiStatsLog.write(
diff --git a/src/com/android/customization/module/logging/ThemesUserEventLogger.kt b/src/com/android/customization/module/logging/ThemesUserEventLogger.kt
index 4fd5334..1210343 100644
--- a/src/com/android/customization/module/logging/ThemesUserEventLogger.kt
+++ b/src/com/android/customization/module/logging/ThemesUserEventLogger.kt
@@ -15,19 +15,38 @@
  */
 package com.android.customization.module.logging
 
+import android.stats.style.StyleEnums
+import androidx.annotation.IntDef
 import com.android.customization.model.color.ColorOption
 import com.android.customization.model.grid.GridOption
 import com.android.wallpaper.module.logging.UserEventLogger
 
 /** Extension of [UserEventLogger] that adds ThemePicker specific events. */
 interface ThemesUserEventLogger : UserEventLogger {
-    /**
-     * Logs the color usage while color is applied.
-     *
-     * @param action color applied action.
-     * @param colorOption applied color option.
-     */
-    fun logColorApplied(action: Int, colorOption: ColorOption)
+
+    fun logThemeColorApplied(colorOption: ColorOption)
 
     fun logGridApplied(grid: GridOption)
+
+    fun logClockApplied(clockId: String)
+
+    fun logClockColorApplied(seedColor: Int)
+
+    fun logClockSizeApplied(@ClockSize clockSize: Int)
+
+    fun logThemedIconApplied(useThemeIcon: Boolean)
+
+    fun logLockScreenNotificationApplied(showLockScreenNotifications: Boolean)
+
+    fun logShortcutApplied(shortcut: String, shortcutSlotId: String)
+
+    fun logDarkThemeApplied(useDarkTheme: Boolean)
+
+    @IntDef(
+        StyleEnums.CLOCK_SIZE_UNSPECIFIED,
+        StyleEnums.CLOCK_SIZE_DYNAMIC,
+        StyleEnums.CLOCK_SIZE_SMALL,
+    )
+    @Retention(AnnotationRetention.SOURCE)
+    annotation class ClockSize
 }
diff --git a/src/com/android/customization/module/logging/StatsLogUserEventLogger.kt b/src/com/android/customization/module/logging/ThemesUserEventLoggerImpl.kt
similarity index 71%
rename from src/com/android/customization/module/logging/StatsLogUserEventLogger.kt
rename to src/com/android/customization/module/logging/ThemesUserEventLoggerImpl.kt
index 90a1c6f..8eee709 100644
--- a/src/com/android/customization/module/logging/StatsLogUserEventLogger.kt
+++ b/src/com/android/customization/module/logging/ThemesUserEventLoggerImpl.kt
@@ -21,33 +21,20 @@
 import android.text.TextUtils
 import com.android.customization.model.color.ColorOption
 import com.android.customization.model.grid.GridOption
-import com.android.customization.module.SysUiStatsLogger
+import com.android.customization.module.logging.ThemesUserEventLogger.ClockSize
 import com.android.systemui.shared.system.SysUiStatsLog
 import com.android.wallpaper.module.WallpaperPersister.DEST_BOTH
 import com.android.wallpaper.module.WallpaperPersister.DEST_HOME_SCREEN
 import com.android.wallpaper.module.WallpaperPersister.DEST_LOCK_SCREEN
 import com.android.wallpaper.module.WallpaperPreferences
-import com.android.wallpaper.module.logging.NoOpUserEventLogger
 import com.android.wallpaper.module.logging.UserEventLogger.EffectStatus
 import com.android.wallpaper.module.logging.UserEventLogger.SetWallpaperEntryPoint
 import com.android.wallpaper.module.logging.UserEventLogger.WallpaperDestination
 import com.android.wallpaper.util.LaunchSourceUtils
 
 /** StatsLog-backed implementation of [ThemesUserEventLogger]. */
-class StatsLogUserEventLogger(private val preferences: WallpaperPreferences) :
-    NoOpUserEventLogger(), ThemesUserEventLogger {
-
-    override fun logAppLaunched(launchSource: Intent) {
-        SysUiStatsLogger(SysUiStatsLog.STYLE_UICHANGED__ACTION__APP_LAUNCHED)
-            .setLaunchedPreference(getAppLaunchSource(launchSource))
-            .log()
-    }
-
-    override fun logActionClicked(collectionId: String, actionLabelResId: Int) {
-        SysUiStatsLogger(StyleEnums.WALLPAPER_EXPLORE)
-            .setWallpaperCategoryHash(getIdHashCode(collectionId))
-            .log()
-    }
+class ThemesUserEventLoggerImpl(private val preferences: WallpaperPreferences) :
+    ThemesUserEventLogger {
 
     override fun logSnapshot() {
         SysUiStatsLogger(StyleEnums.SNAPSHOT)
@@ -59,6 +46,12 @@
             .log()
     }
 
+    override fun logAppLaunched(launchSource: Intent) {
+        SysUiStatsLogger(StyleEnums.APP_LAUNCHED)
+            .setLaunchedPreference(launchSource.getAppLaunchSource())
+            .log()
+    }
+
     override fun logWallpaperApplied(
         collectionId: String?,
         wallpaperId: String?,
@@ -114,20 +107,69 @@
             .log()
     }
 
-    override fun logColorApplied(action: Int, colorOption: ColorOption) {
-        SysUiStatsLogger(action)
+    override fun logResetApplied() {
+        SysUiStatsLogger(StyleEnums.RESET_APPLIED).log()
+    }
+
+    override fun logWallpaperExploreButtonClicked() {
+        SysUiStatsLogger(StyleEnums.WALLPAPER_EXPLORE).log()
+    }
+
+    override fun logThemeColorApplied(colorOption: ColorOption) {
+        SysUiStatsLogger(StyleEnums.THEME_COLOR_APPLIED)
             .setColorPreference(colorOption.index)
             .setColorVariant(colorOption.style.ordinal + 1)
             .log()
     }
 
     override fun logGridApplied(grid: GridOption) {
-        SysUiStatsLogger(StyleEnums.PICKER_APPLIED).setLauncherGrid(grid.cols).log()
+        SysUiStatsLogger(StyleEnums.GRID_APPLIED).setLauncherGrid(grid.getLauncherGridInt()).log()
     }
 
-    private fun getAppLaunchSource(launchSource: Intent): Int {
-        return if (launchSource.hasExtra(LaunchSourceUtils.WALLPAPER_LAUNCH_SOURCE)) {
-            when (launchSource.getStringExtra(LaunchSourceUtils.WALLPAPER_LAUNCH_SOURCE)) {
+    override fun logClockApplied(clockId: String) {
+        SysUiStatsLogger(StyleEnums.CLOCK_APPLIED).setClockPackageHash(getIdHashCode(clockId)).log()
+    }
+
+    override fun logClockColorApplied(seedColor: Int) {
+        SysUiStatsLogger(StyleEnums.CLOCK_COLOR_APPLIED).setSeedColor(seedColor).log()
+    }
+
+    override fun logClockSizeApplied(@ClockSize clockSize: Int) {
+        SysUiStatsLogger(StyleEnums.CLOCK_SIZE_APPLIED).setClockSize(clockSize).log()
+    }
+
+    override fun logThemedIconApplied(useThemeIcon: Boolean) {
+        SysUiStatsLogger(StyleEnums.THEMED_ICON_APPLIED).setToggleOn(useThemeIcon).log()
+    }
+
+    override fun logLockScreenNotificationApplied(showLockScreenNotifications: Boolean) {
+        SysUiStatsLogger(StyleEnums.LOCK_SCREEN_NOTIFICATION_APPLIED)
+            .setToggleOn(showLockScreenNotifications)
+            .log()
+    }
+
+    override fun logShortcutApplied(shortcut: String, shortcutSlotId: String) {
+        SysUiStatsLogger(StyleEnums.SHORTCUT_APPLIED)
+            .setShortcut(shortcut)
+            .setShortcutSlotId(shortcutSlotId)
+            .log()
+    }
+
+    override fun logDarkThemeApplied(useDarkTheme: Boolean) {
+        SysUiStatsLogger(StyleEnums.DARK_THEME_APPLIED).setToggleOn(useDarkTheme).log()
+    }
+
+    /**
+     * The grid integer depends on the column and row numbers. For example: 4x5 is 405 13x37 is 1337
+     * The upper limit for the column / row count is 99.
+     */
+    private fun GridOption.getLauncherGridInt(): Int {
+        return cols * 100 + rows
+    }
+
+    private fun Intent.getAppLaunchSource(): Int {
+        return if (hasExtra(LaunchSourceUtils.WALLPAPER_LAUNCH_SOURCE)) {
+            when (getStringExtra(LaunchSourceUtils.WALLPAPER_LAUNCH_SOURCE)) {
                 LaunchSourceUtils.LAUNCH_SOURCE_LAUNCHER ->
                     SysUiStatsLog.STYLE_UICHANGED__LAUNCHED_PREFERENCE__LAUNCHED_LAUNCHER
                 LaunchSourceUtils.LAUNCH_SOURCE_SETTINGS ->
@@ -142,17 +184,11 @@
                     SysUiStatsLog
                         .STYLE_UICHANGED__LAUNCHED_PREFERENCE__LAUNCHED_PREFERENCE_UNSPECIFIED
             }
-        } else if (launchSource.hasExtra(LaunchSourceUtils.LAUNCH_SETTINGS_SEARCH)) {
+        } else if (hasExtra(LaunchSourceUtils.LAUNCH_SETTINGS_SEARCH)) {
             SysUiStatsLog.STYLE_UICHANGED__LAUNCHED_PREFERENCE__LAUNCHED_SETTINGS_SEARCH
-        } else if (
-            launchSource.action != null &&
-                launchSource.action == WallpaperManager.ACTION_CROP_AND_SET_WALLPAPER
-        ) {
+        } else if (action != null && action == WallpaperManager.ACTION_CROP_AND_SET_WALLPAPER) {
             SysUiStatsLog.STYLE_UICHANGED__LAUNCHED_PREFERENCE__LAUNCHED_CROP_AND_SET_ACTION
-        } else if (
-            launchSource.categories != null &&
-                launchSource.categories.contains(Intent.CATEGORY_LAUNCHER)
-        ) {
+        } else if (categories != null && categories.contains(Intent.CATEGORY_LAUNCHER)) {
             SysUiStatsLog.STYLE_UICHANGED__LAUNCHED_PREFERENCE__LAUNCHED_LAUNCH_ICON
         } else {
             SysUiStatsLog.STYLE_UICHANGED__LAUNCHED_PREFERENCE__LAUNCHED_PREFERENCE_UNSPECIFIED
diff --git a/src/com/android/customization/picker/clock/ui/fragment/ClockSettingsFragment.kt b/src/com/android/customization/picker/clock/ui/fragment/ClockSettingsFragment.kt
index 068f44f..dc70633 100644
--- a/src/com/android/customization/picker/clock/ui/fragment/ClockSettingsFragment.kt
+++ b/src/com/android/customization/picker/clock/ui/fragment/ClockSettingsFragment.kt
@@ -88,14 +88,14 @@
                             injector
                                 .getCurrentWallpaperInfoFactory(context)
                                 .createCurrentWallpaperInfos(
-                                    { homeWallpaper, lockWallpaper, _ ->
-                                        continuation.resume(
-                                            lockWallpaper ?: homeWallpaper,
-                                            null,
-                                        )
-                                    },
+                                    context,
                                     forceReload,
-                                )
+                                ) { homeWallpaper, lockWallpaper, _ ->
+                                    continuation.resume(
+                                        lockWallpaper ?: homeWallpaper,
+                                        null,
+                                    )
+                                }
                         }
                     },
                     onWallpaperColorChanged = { colors ->
diff --git a/src/com/android/customization/picker/color/ui/fragment/ColorPickerFragment.kt b/src/com/android/customization/picker/color/ui/fragment/ColorPickerFragment.kt
index 911b821..941d2c6 100644
--- a/src/com/android/customization/picker/color/ui/fragment/ColorPickerFragment.kt
+++ b/src/com/android/customization/picker/color/ui/fragment/ColorPickerFragment.kt
@@ -114,23 +114,23 @@
                         wallpaperInfoProvider = { forceReload ->
                             suspendCancellableCoroutine { continuation ->
                                 wallpaperInfoFactory.createCurrentWallpaperInfos(
-                                    { homeWallpaper, lockWallpaper, _ ->
-                                        lifecycleScope.launch {
-                                            if (
-                                                wallpaperColorsRepository.lockWallpaperColors.value
-                                                    is WallpaperColorsModel.Loading
-                                            ) {
-                                                loadInitialColors(
-                                                    wallpaperManager,
-                                                    wallpaperColorsRepository,
-                                                    CustomizationSections.Screen.LOCK_SCREEN
-                                                )
-                                            }
-                                        }
-                                        continuation.resume(lockWallpaper ?: homeWallpaper, null)
-                                    },
+                                    context,
                                     forceReload,
-                                )
+                                ) { homeWallpaper, lockWallpaper, _ ->
+                                    lifecycleScope.launch {
+                                        if (
+                                            wallpaperColorsRepository.lockWallpaperColors.value
+                                                is WallpaperColorsModel.Loading
+                                        ) {
+                                            loadInitialColors(
+                                                wallpaperManager,
+                                                wallpaperColorsRepository,
+                                                CustomizationSections.Screen.LOCK_SCREEN
+                                            )
+                                        }
+                                    }
+                                    continuation.resume(lockWallpaper ?: homeWallpaper, null)
+                                }
                             }
                         },
                         onWallpaperColorChanged = { colors ->
@@ -165,23 +165,23 @@
                     wallpaperInfoProvider = { forceReload ->
                         suspendCancellableCoroutine { continuation ->
                             wallpaperInfoFactory.createCurrentWallpaperInfos(
-                                { homeWallpaper, lockWallpaper, _ ->
-                                    lifecycleScope.launch {
-                                        if (
-                                            wallpaperColorsRepository.homeWallpaperColors.value
-                                                is WallpaperColorsModel.Loading
-                                        ) {
-                                            loadInitialColors(
-                                                wallpaperManager,
-                                                wallpaperColorsRepository,
-                                                CustomizationSections.Screen.HOME_SCREEN
-                                            )
-                                        }
-                                    }
-                                    continuation.resume(homeWallpaper ?: lockWallpaper, null)
-                                },
+                                context,
                                 forceReload,
-                            )
+                            ) { homeWallpaper, lockWallpaper, _ ->
+                                lifecycleScope.launch {
+                                    if (
+                                        wallpaperColorsRepository.homeWallpaperColors.value
+                                            is WallpaperColorsModel.Loading
+                                    ) {
+                                        loadInitialColors(
+                                            wallpaperManager,
+                                            wallpaperColorsRepository,
+                                            CustomizationSections.Screen.HOME_SCREEN
+                                        )
+                                    }
+                                }
+                                continuation.resume(homeWallpaper ?: lockWallpaper, null)
+                            }
                         }
                     },
                     onWallpaperColorChanged = { colors ->
diff --git a/src/com/android/customization/picker/grid/ui/fragment/GridFragment.kt b/src/com/android/customization/picker/grid/ui/fragment/GridFragment.kt
index 52bc5e1..2a301b4 100644
--- a/src/com/android/customization/picker/grid/ui/fragment/GridFragment.kt
+++ b/src/com/android/customization/picker/grid/ui/fragment/GridFragment.kt
@@ -185,11 +185,11 @@
                     wallpaperInfoProvider = {
                         suspendCancellableCoroutine { continuation ->
                             wallpaperInfoFactory.createCurrentWallpaperInfos(
-                                { homeWallpaper, lockWallpaper, _ ->
-                                    continuation.resume(homeWallpaper ?: lockWallpaper, null)
-                                },
+                                context,
                                 /* forceRefresh= */ true,
-                            )
+                            ) { homeWallpaper, lockWallpaper, _ ->
+                                continuation.resume(homeWallpaper ?: lockWallpaper, null)
+                            }
                         }
                     },
                     wallpaperInteractor = wallpaperInteractor,
diff --git a/src/com/android/customization/picker/preview/ui/section/PreviewWithThemeSectionController.kt b/src/com/android/customization/picker/preview/ui/section/PreviewWithThemeSectionController.kt
index 0edc216..b3e778b 100644
--- a/src/com/android/customization/picker/preview/ui/section/PreviewWithThemeSectionController.kt
+++ b/src/com/android/customization/picker/preview/ui/section/PreviewWithThemeSectionController.kt
@@ -35,6 +35,7 @@
 import com.android.wallpaper.picker.customization.ui.viewmodel.ScreenPreviewViewModel
 import com.android.wallpaper.util.DisplayUtils
 import com.android.wallpaper.util.PreviewUtils
+import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.suspendCancellableCoroutine
 
 /**
@@ -69,6 +70,7 @@
         isTwoPaneAndSmallWidth,
         customizationPickerViewModel,
     ) {
+    @OptIn(ExperimentalCoroutinesApi::class)
     override fun createScreenPreviewViewModel(context: Context): ScreenPreviewViewModel {
         return PreviewWithThemeViewModel(
             previewUtils =
@@ -92,21 +94,21 @@
             wallpaperInfoProvider = { forceReload ->
                 suspendCancellableCoroutine { continuation ->
                     wallpaperInfoFactory.createCurrentWallpaperInfos(
-                        { homeWallpaper, lockWallpaper, _ ->
-                            val wallpaper =
-                                if (isOnLockScreen) {
-                                    lockWallpaper ?: homeWallpaper
-                                } else {
-                                    homeWallpaper ?: lockWallpaper
-                                }
-                            loadInitialColors(
-                                context = context,
-                                screen = screen,
-                            )
-                            continuation.resume(wallpaper, null)
-                        },
+                        context,
                         forceReload,
-                    )
+                    ) { homeWallpaper, lockWallpaper, _ ->
+                        val wallpaper =
+                            if (isOnLockScreen) {
+                                lockWallpaper ?: homeWallpaper
+                            } else {
+                                homeWallpaper ?: lockWallpaper
+                            }
+                        loadInitialColors(
+                            context = context,
+                            screen = screen,
+                        )
+                        continuation.resume(wallpaper, null)
+                    }
                 }
             },
             onWallpaperColorChanged = { colors ->
diff --git a/src/com/android/customization/picker/quickaffordance/ui/viewmodel/KeyguardQuickAffordancePickerViewModel.kt b/src/com/android/customization/picker/quickaffordance/ui/viewmodel/KeyguardQuickAffordancePickerViewModel.kt
index f832cde..809d20e 100644
--- a/src/com/android/customization/picker/quickaffordance/ui/viewmodel/KeyguardQuickAffordancePickerViewModel.kt
+++ b/src/com/android/customization/picker/quickaffordance/ui/viewmodel/KeyguardQuickAffordancePickerViewModel.kt
@@ -92,11 +92,11 @@
             wallpaperInfoProvider = { forceReload ->
                 suspendCancellableCoroutine { continuation ->
                     wallpaperInfoFactory.createCurrentWallpaperInfos(
-                        { homeWallpaper, lockWallpaper, _ ->
-                            continuation.resume(lockWallpaper ?: homeWallpaper, null)
-                        },
+                        context,
                         forceReload,
-                    )
+                    ) { homeWallpaper, lockWallpaper, _ ->
+                        continuation.resume(lockWallpaper ?: homeWallpaper, null)
+                    }
                 }
             },
             wallpaperInteractor = wallpaperInteractor,
diff --git a/tests/common/Android.bp b/tests/common/Android.bp
index d488ee4..0961886 100644
--- a/tests/common/Android.bp
+++ b/tests/common/Android.bp
@@ -25,8 +25,8 @@
 
     defaults: ["ThemePicker_common_defaults"],
     srcs: [
-        "src/com/android/customization/testing/**/*.java",
-        "src/com/android/customization/testing/**/*.kt",
+        "src/**/*.java",
+        "src/**/*.kt",
     ],
     static_libs: [
         "WallpaperPicker2TestLib",
diff --git a/tests/common/src/com/android/customization/module/logging/TestThemesUserEventLogger.kt b/tests/common/src/com/android/customization/module/logging/TestThemesUserEventLogger.kt
new file mode 100644
index 0000000..b2835f4
--- /dev/null
+++ b/tests/common/src/com/android/customization/module/logging/TestThemesUserEventLogger.kt
@@ -0,0 +1,42 @@
+/*
+ * Copyright (C) 2019 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.customization.model.color.ColorOption
+import com.android.customization.model.grid.GridOption
+import com.android.wallpaper.module.logging.TestUserEventLogger
+
+/** Test implementation of [ThemesUserEventLogger]. */
+class TestThemesUserEventLogger : TestUserEventLogger(), ThemesUserEventLogger {
+
+    override fun logThemeColorApplied(colorOption: ColorOption) {}
+
+    override fun logGridApplied(grid: GridOption) {}
+
+    override fun logClockApplied(clockId: String) {}
+
+    override fun logClockColorApplied(seedColor: Int) {}
+
+    override fun logClockSizeApplied(clockSize: Int) {}
+
+    override fun logThemedIconApplied(useThemeIcon: Boolean) {}
+
+    override fun logLockScreenNotificationApplied(showLockScreenNotifications: Boolean) {}
+
+    override fun logShortcutApplied(shortcut: String, shortcutSlotId: String) {}
+
+    override fun logDarkThemeApplied(useDarkTheme: Boolean) {}
+}
diff --git a/tests/common/src/com/android/customization/testing/TestCustomizationInjector.kt b/tests/common/src/com/android/customization/testing/TestCustomizationInjector.kt
index 8b898f8..e08209d 100644
--- a/tests/common/src/com/android/customization/testing/TestCustomizationInjector.kt
+++ b/tests/common/src/com/android/customization/testing/TestCustomizationInjector.kt
@@ -5,6 +5,7 @@
 import androidx.activity.ComponentActivity
 import com.android.customization.module.CustomizationInjector
 import com.android.customization.module.CustomizationPreferences
+import com.android.customization.module.logging.TestThemesUserEventLogger
 import com.android.customization.module.logging.ThemesUserEventLogger
 import com.android.customization.picker.clock.domain.interactor.ClockPickerInteractor
 import com.android.customization.picker.clock.ui.view.ClockViewFactory
diff --git a/tests/common/src/com/android/customization/testing/TestThemesUserEventLogger.kt b/tests/common/src/com/android/customization/testing/TestThemesUserEventLogger.kt
deleted file mode 100644
index 2236b62..0000000
--- a/tests/common/src/com/android/customization/testing/TestThemesUserEventLogger.kt
+++ /dev/null
@@ -1,29 +0,0 @@
-/*
- * Copyright (C) 2019 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.testing
-
-import com.android.customization.model.color.ColorOption
-import com.android.customization.model.grid.GridOption
-import com.android.customization.module.logging.ThemesUserEventLogger
-import com.android.wallpaper.testing.TestUserEventLogger
-
-/** Test implementation of [ThemesUserEventLogger]. */
-class TestThemesUserEventLogger : TestUserEventLogger(), ThemesUserEventLogger {
-
-    override fun logColorApplied(action: Int, colorOption: ColorOption) {}
-
-    override fun logGridApplied(grid: GridOption) {}
-}