diff --git a/Android.bp b/Android.bp
index 4dbe399..9f7dbdb 100644
--- a/Android.bp
+++ b/Android.bp
@@ -66,8 +66,9 @@
         + "$(location soong_zip) -o $(out) -C $$RES_DIR -D $$RES_DIR"
 }
 
+// Common defaults that doesn't have ThemePicker specifics.
 java_defaults {
-    name: "ThemePicker_defaults",
+    name: "ThemePicker_common_defaults",
 
     static_libs: [
         "guava",
@@ -92,7 +93,6 @@
     srcs: [
         ":WallpaperPicker2_srcs",
         ":ThemePicker_srcs",
-        ":ThemePicker_src_overrides",
     ],
 
     required: ["android.software.theme_picker.xml"],
@@ -111,6 +111,12 @@
     system_ext_specific: true,
 }
 
+java_defaults {
+    name: "ThemePicker_defaults",
+    defaults: ["ThemePicker_common_defaults"],
+    srcs: [":ThemePicker_src_overrides"],
+}
+
 prebuilt_etc {
     name: "android.software.theme_picker.xml",
     system_ext_specific: true,
diff --git a/src/com/android/customization/module/SysUiStatsLogger.kt b/src/com/android/customization/module/SysUiStatsLogger.kt
index 8e97b0b..ee821ee 100644
--- a/src/com/android/customization/module/SysUiStatsLogger.kt
+++ b/src/com/android/customization/module/SysUiStatsLogger.kt
@@ -15,8 +15,15 @@
  */
 package com.android.customization.module
 
-import android.stats.style.StyleEnums
 import com.android.systemui.shared.system.SysUiStatsLog
+import com.android.systemui.shared.system.SysUiStatsLog.STYLE_UICHANGED__CLOCK_SIZE__CLOCK_SIZE_UNSPECIFIED
+import com.android.systemui.shared.system.SysUiStatsLog.STYLE_UICHANGED__COLOR_SOURCE__COLOR_SOURCE_UNSPECIFIED
+import com.android.systemui.shared.system.SysUiStatsLog.STYLE_UICHANGED__DATE_PREFERENCE__DATE_PREFERENCE_UNSPECIFIED
+import com.android.systemui.shared.system.SysUiStatsLog.STYLE_UICHANGED__EFFECT_PREFERENCE__EFFECT_PREFERENCE_UNSPECIFIED
+import com.android.systemui.shared.system.SysUiStatsLog.STYLE_UICHANGED__LAUNCHED_PREFERENCE__LAUNCHED_PREFERENCE_UNSPECIFIED
+import com.android.systemui.shared.system.SysUiStatsLog.STYLE_UICHANGED__LOCATION_PREFERENCE__LOCATION_PREFERENCE_UNSPECIFIED
+import com.android.systemui.shared.system.SysUiStatsLog.STYLE_UICHANGED__SET_WALLPAPER_ENTRY_POINT__SET_WALLPAPER_ENTRY_POINT_UNSPECIFIED
+import com.android.systemui.shared.system.SysUiStatsLog.STYLE_UICHANGED__WALLPAPER_DESTINATION__WALLPAPER_DESTINATION_UNSPECIFIED
 import com.android.systemui.shared.system.SysUiStatsLog.STYLE_UI_CHANGED
 
 /** The builder for [SysUiStatsLog]. */
@@ -30,10 +37,10 @@
     private var wallpaperCategoryHash = 0
     private var wallpaperIdHash = 0
     private var colorPreference = 0
-    private var locationPreference = StyleEnums.EFFECT_PREFERENCE_UNSPECIFIED
-    private var datePreference = StyleEnums.DATE_PREFERENCE_UNSPECIFIED
-    private var launchedPreference = StyleEnums.LAUNCHED_PREFERENCE_UNSPECIFIED
-    private var effectPreference = StyleEnums.EFFECT_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 effectPreference = STYLE_UICHANGED__EFFECT_PREFERENCE__EFFECT_PREFERENCE_UNSPECIFIED
     private var effectIdHash = 0
     private var lockWallpaperCategoryHash = 0
     private var lockWallpaperIdHash = 0
@@ -43,6 +50,15 @@
     private var colorVariant = 0
     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 colorSource = STYLE_UICHANGED__COLOR_SOURCE__COLOR_SOURCE_UNSPECIFIED
+    private var seedColor = 0
+    private var clockSize = STYLE_UICHANGED__CLOCK_SIZE__CLOCK_SIZE_UNSPECIFIED
+    private var toggleOn = false
+    private var shortcut = ""
+    private var shortcutSlotId = ""
 
     fun setColorPackageHash(colorPackageHash: Int) = apply {
         this.colorPackageHash = colorPackageHash
@@ -118,6 +134,42 @@
         this.effectResultCode = effectResultCode
     }
 
+    fun setSessionId(sessionId: Int) = apply {
+        this.sessionId = sessionId
+    }
+
+    fun setSetWallpaperEntryPoint(setWallpaperEntryPoint: Int) = apply {
+        this.setWallpaperEntryPoint = setWallpaperEntryPoint
+    }
+
+    fun setWallpaperDestination(wallpaperDestination: Int) = apply {
+        this.wallpaperDestination = wallpaperDestination
+    }
+
+    fun setColorSource(colorSource: Int) = apply {
+        this.colorSource = colorSource
+    }
+
+    fun setSeedColor(seedColor: Int) = apply {
+        this.seedColor = seedColor
+    }
+
+    fun setClockSize(clockSize: Int) = apply {
+        this.clockSize = clockSize
+    }
+
+    fun setToggleOn(toggleOn: Boolean) = apply {
+        this.toggleOn = toggleOn
+    }
+
+    fun setShortcut(shortcut: String) = apply {
+        this.shortcut = shortcut
+    }
+
+    fun setShortcutSlotId(shortcutSlotId: String) = apply {
+        this.shortcutSlotId = shortcutSlotId
+    }
+
     fun log() {
         SysUiStatsLog.write(
             STYLE_UI_CHANGED,
@@ -143,6 +195,15 @@
             colorVariant,
             timeElapsedMillis,
             effectResultCode,
+            sessionId,
+            setWallpaperEntryPoint,
+            wallpaperDestination,
+            colorSource,
+            seedColor,
+            clockSize,
+            toggleOn,
+            shortcut,
+            shortcutSlotId,
         )
     }
 }
diff --git a/tests/common/Android.bp b/tests/common/Android.bp
index 2904fad..d488ee4 100644
--- a/tests/common/Android.bp
+++ b/tests/common/Android.bp
@@ -23,7 +23,7 @@
 android_library {
     name: "ThemePickerTestLib",
 
-    defaults: ["ThemePicker_defaults"],
+    defaults: ["ThemePicker_common_defaults"],
     srcs: [
         "src/com/android/customization/testing/**/*.java",
         "src/com/android/customization/testing/**/*.kt",
diff --git a/tests/common/src/com/android/customization/testing/TestCustomizationInjector.kt b/tests/common/src/com/android/customization/testing/TestCustomizationInjector.kt
index 445286f..cd587d6 100644
--- a/tests/common/src/com/android/customization/testing/TestCustomizationInjector.kt
+++ b/tests/common/src/com/android/customization/testing/TestCustomizationInjector.kt
@@ -23,7 +23,7 @@
 import javax.inject.Singleton
 
 @Singleton
-class TestCustomizationInjector @Inject constructor() : TestInjector(), CustomizationInjector {
+open class TestCustomizationInjector @Inject constructor() : TestInjector(), CustomizationInjector {
     private var customizationPrefs: CustomizationPreferences? = null
     private var themesUserEventLogger: ThemesUserEventLogger? = null