diff --git a/src_override/com/android/wallpaper/modules/ThemePickerAppModule.kt b/src_override/com/android/wallpaper/modules/ThemePickerAppModule.kt
index 4426b49..31b4cd8 100644
--- a/src_override/com/android/wallpaper/modules/ThemePickerAppModule.kt
+++ b/src_override/com/android/wallpaper/modules/ThemePickerAppModule.kt
@@ -51,6 +51,8 @@
 import com.android.wallpaper.picker.category.domain.interactor.implementations.DefaultCategoriesLoadingStatusInteractor
 import com.android.wallpaper.picker.category.ui.view.providers.IndividualPickerFactory
 import com.android.wallpaper.picker.category.ui.view.providers.implementation.DefaultIndividualPickerFactory
+import com.android.wallpaper.picker.category.wrapper.DefaultWallpaperCategoryWrapper
+import com.android.wallpaper.picker.category.wrapper.WallpaperCategoryWrapper
 import com.android.wallpaper.picker.common.preview.ui.binder.ThemePickerWorkspaceCallbackBinder
 import com.android.wallpaper.picker.common.preview.ui.binder.WorkspaceCallbackBinder
 import com.android.wallpaper.picker.customization.ui.binder.CustomizationOptionsBinder
@@ -86,11 +88,17 @@
     @Binds
     @Singleton
     abstract fun bindCreativeCategoryInteractor(
-        impl: CreativeCategoryInteractorImpl,
+        impl: CreativeCategoryInteractorImpl
     ): CreativeCategoryInteractor
 
     @Binds
     @Singleton
+    abstract fun bindWallpaperCategoryWrapper(
+        impl: DefaultWallpaperCategoryWrapper
+    ): WallpaperCategoryWrapper
+
+    @Binds
+    @Singleton
     abstract fun bindCustomizationInjector(impl: ThemePickerInjector): CustomizationInjector
 
     @Binds
@@ -122,7 +130,7 @@
     @Binds
     @Singleton
     abstract fun bindLoadingStatusInteractor(
-        impl: DefaultCategoriesLoadingStatusInteractor,
+        impl: DefaultCategoriesLoadingStatusInteractor
     ): CategoriesLoadingStatusInteractor
 
     @Binds
diff --git a/tests/common/src/com/android/customization/testing/TestCustomizationInjector.kt b/tests/common/src/com/android/customization/testing/TestCustomizationInjector.kt
index c988424..4e97599 100644
--- a/tests/common/src/com/android/customization/testing/TestCustomizationInjector.kt
+++ b/tests/common/src/com/android/customization/testing/TestCustomizationInjector.kt
@@ -15,6 +15,7 @@
 import com.android.customization.picker.color.ui.viewmodel.ColorPickerViewModel
 import com.android.customization.picker.quickaffordance.domain.interactor.KeyguardQuickAffordancePickerInteractor
 import com.android.wallpaper.module.logging.UserEventLogger
+import com.android.wallpaper.picker.category.wrapper.WallpaperCategoryWrapper
 import com.android.wallpaper.picker.customization.data.repository.WallpaperColorsRepository
 import com.android.wallpaper.testing.TestInjector
 import javax.inject.Inject
@@ -25,7 +26,7 @@
 @Inject
 constructor(
     private val customPrefs: TestDefaultCustomizationPreferences,
-    private val themesUserEventLogger: ThemesUserEventLogger
+    private val themesUserEventLogger: ThemesUserEventLogger,
 ) : TestInjector(themesUserEventLogger), CustomizationInjector {
     /////////////////
     // CustomizationInjector implementations
@@ -43,7 +44,7 @@
 
     override fun getWallpaperColorResources(
         wallpaperColors: WallpaperColors,
-        context: Context
+        context: Context,
     ): WallpaperColorResources {
         throw UnsupportedOperationException("not implemented")
     }
@@ -79,4 +80,8 @@
     override fun getUserEventLogger(): UserEventLogger {
         return themesUserEventLogger
     }
+
+    override fun getWallpaperCategoryWrapper(): WallpaperCategoryWrapper {
+        return super.fakeWallpaperCategoryWrapper
+    }
 }
diff --git a/tests/module/src/com/android/wallpaper/ThemePickerTestModule.kt b/tests/module/src/com/android/wallpaper/ThemePickerTestModule.kt
index ac49994..8f09d51 100644
--- a/tests/module/src/com/android/wallpaper/ThemePickerTestModule.kt
+++ b/tests/module/src/com/android/wallpaper/ThemePickerTestModule.kt
@@ -45,6 +45,7 @@
 import com.android.wallpaper.module.logging.UserEventLogger
 import com.android.wallpaper.modules.ThemePickerAppModule
 import com.android.wallpaper.network.Requester
+import com.android.wallpaper.picker.category.wrapper.WallpaperCategoryWrapper
 import com.android.wallpaper.picker.common.preview.ui.binder.ThemePickerWorkspaceCallbackBinder
 import com.android.wallpaper.picker.common.preview.ui.binder.WorkspaceCallbackBinder
 import com.android.wallpaper.picker.customization.ui.binder.CustomizationOptionsBinder
@@ -55,6 +56,7 @@
 import com.android.wallpaper.picker.preview.ui.util.DefaultImageEffectDialogUtil
 import com.android.wallpaper.picker.preview.ui.util.ImageEffectDialogUtil
 import com.android.wallpaper.testing.FakeDefaultRequester
+import com.android.wallpaper.testing.FakeWallpaperCategoryWrapper
 import com.android.wallpaper.testing.TestPartnerProvider
 import com.android.wallpaper.util.converter.DefaultWallpaperModelFactory
 import com.android.wallpaper.util.converter.WallpaperModelFactory
@@ -102,6 +104,12 @@
 
     @Binds
     @Singleton
+    abstract fun bindWallpaperCategoryWrapper(
+        impl: FakeWallpaperCategoryWrapper
+    ): WallpaperCategoryWrapper
+
+    @Binds
+    @Singleton
     abstract fun bindImageEffectDialogUtil(
         impl: DefaultImageEffectDialogUtil
     ): ImageEffectDialogUtil
@@ -164,7 +172,7 @@
         fun provideColorCustomizationManager(): ColorCustomizationManager {
             return ColorCustomizationManager.getInstance(
                 ApplicationProvider.getApplicationContext(),
-                OverlayManagerCompat(ApplicationProvider.getApplicationContext())
+                OverlayManagerCompat(ApplicationProvider.getApplicationContext()),
             )
         }