diff --git a/src/com/android/customization/module/ThemePickerInjector.kt b/src/com/android/customization/module/ThemePickerInjector.kt
index da25950..c751178 100644
--- a/src/com/android/customization/module/ThemePickerInjector.kt
+++ b/src/com/android/customization/module/ThemePickerInjector.kt
@@ -70,8 +70,14 @@
 import com.android.wallpaper.config.BaseFlags
 import com.android.wallpaper.module.CustomizationSections
 import com.android.wallpaper.module.FragmentFactory
+import com.android.wallpaper.module.NetworkStatusNotifier
+import com.android.wallpaper.module.PartnerProvider
 import com.android.wallpaper.module.WallpaperPicker2Injector
+import com.android.wallpaper.module.WallpaperPreferences
+import com.android.wallpaper.module.logging.UserEventLogger
+import com.android.wallpaper.network.Requester
 import com.android.wallpaper.picker.CustomizationPickerActivity
+import com.android.wallpaper.picker.customization.data.content.WallpaperClient
 import com.android.wallpaper.picker.customization.data.content.WallpaperClientImpl
 import com.android.wallpaper.picker.customization.data.repository.WallpaperColorsRepository
 import com.android.wallpaper.picker.customization.data.repository.WallpaperRepository
@@ -79,6 +85,8 @@
 import com.android.wallpaper.picker.di.modules.BackgroundDispatcher
 import com.android.wallpaper.picker.di.modules.MainDispatcher
 import com.android.wallpaper.picker.undo.domain.interactor.SnapshotRestorer
+import com.android.wallpaper.system.UiModeManagerWrapper
+import com.android.wallpaper.util.DisplayUtils
 import com.android.wallpaper.util.ScreenSizeCalculator
 import dagger.Lazy
 import javax.inject.Inject
@@ -94,7 +102,32 @@
     @MainDispatcher private val mainDispatcher: CoroutineDispatcher,
     @BackgroundDispatcher private val bgScope: CoroutineScope,
     @BackgroundDispatcher private val bgDispatcher: CoroutineDispatcher,
-) : WallpaperPicker2Injector(mainScope, bgDispatcher), CustomizationInjector {
+    private val colorContrastSectionViewModelFactory: Lazy<ColorContrastSectionViewModel.Factory>,
+    private val themesUserEventLogger: Lazy<ThemesUserEventLogger>,
+    displayUtils: Lazy<DisplayUtils>,
+    requester: Lazy<Requester>,
+    networkStatusNotifier: Lazy<NetworkStatusNotifier>,
+    partnerProvider: Lazy<PartnerProvider>,
+    val uiModeManager: Lazy<UiModeManagerWrapper>,
+    userEventLogger: Lazy<UserEventLogger>,
+    injectedWallpaperClient: Lazy<WallpaperClient>,
+    private val injectedWallpaperInteractor: Lazy<WallpaperInteractor>,
+    prefs: Lazy<WallpaperPreferences>,
+) :
+    WallpaperPicker2Injector(
+        mainScope,
+        bgDispatcher,
+        displayUtils,
+        requester,
+        networkStatusNotifier,
+        partnerProvider,
+        uiModeManager,
+        userEventLogger,
+        injectedWallpaperClient,
+        injectedWallpaperInteractor,
+        prefs,
+    ),
+    CustomizationInjector {
     private var customizationSections: CustomizationSections? = null
     private var wallpaperInteractor: WallpaperInteractor? = null
     private var keyguardQuickAffordancePickerInteractor: KeyguardQuickAffordancePickerInteractor? =
@@ -127,9 +160,6 @@
     private var clockRegistryProvider: ClockRegistryProvider? = null
 
     // Injected objects, sorted by type
-    @Inject
-    lateinit var colorContrastSectionViewModelFactory: Lazy<ColorContrastSectionViewModel.Factory>
-    @Inject lateinit var themesUserEventLogger: Lazy<ThemesUserEventLogger>
 
     override fun getCustomizationSections(activity: ComponentActivity): CustomizationSections {
         val appContext = activity.applicationContext
@@ -180,9 +210,7 @@
         return fragmentFactory ?: ThemePickerFragmentFactory().also { fragmentFactory }
     }
 
-    override fun getSnapshotRestorers(
-        context: Context,
-    ): Map<Int, SnapshotRestorer> {
+    override fun getSnapshotRestorers(context: Context): Map<Int, SnapshotRestorer> {
         return super<WallpaperPicker2Injector>.getSnapshotRestorers(context).toMutableMap().apply {
             this[KEY_QUICK_AFFORDANCE_SNAPSHOT_RESTORER] =
                 getKeyguardQuickAffordanceSnapshotRestorer(context)
@@ -229,7 +257,7 @@
                             getColorCustomizationManager(appContext).currentColorSource,
                             COLOR_SOURCE_PRESET,
                         )
-                    }
+                    },
                 )
                 .also { wallpaperInteractor = it }
     }
@@ -264,7 +292,7 @@
         val appContext = context.applicationContext
         return KeyguardQuickAffordancePickerInteractor(
             KeyguardQuickAffordancePickerRepository(client, getApplicationCoroutineScope()),
-            client
+            client,
         ) {
             getKeyguardQuickAffordanceSnapshotRestorer(appContext)
         }
@@ -285,13 +313,13 @@
         return keyguardQuickAffordanceSnapshotRestorer
             ?: KeyguardQuickAffordanceSnapshotRestorer(
                     getKeyguardQuickAffordancePickerInteractor(context),
-                    getKeyguardQuickAffordancePickerProviderClient(context)
+                    getKeyguardQuickAffordancePickerProviderClient(context),
                 )
                 .also { keyguardQuickAffordanceSnapshotRestorer = it }
     }
 
     fun getNotificationSectionViewModelFactory(
-        context: Context,
+        context: Context
     ): NotificationSectionViewModel.Factory {
         return notificationSectionViewModelFactory
             ?: NotificationSectionViewModel.Factory(
@@ -301,9 +329,7 @@
                 .also { notificationSectionViewModelFactory = it }
     }
 
-    private fun getNotificationsInteractor(
-        context: Context,
-    ): NotificationSettingsInteractor {
+    private fun getNotificationsInteractor(context: Context): NotificationSettingsInteractor {
         return notificationSettingsInteractor
             ?: NotificationSettingsInteractor(
                     repository =
@@ -311,7 +337,7 @@
                             scope = getApplicationCoroutineScope(),
                             backgroundDispatcher = bgDispatcher,
                             secureSettingsRepository = getSecureSettingsRepository(context),
-                        ),
+                        )
                 )
                 .also { notificationSettingsInteractor = it }
     }
@@ -319,10 +345,7 @@
     private fun getNotificationsSnapshotRestorer(context: Context): NotificationsSnapshotRestorer {
         return notificationsSnapshotRestorer
             ?: NotificationsSnapshotRestorer(
-                    interactor =
-                        getNotificationsInteractor(
-                            context = context,
-                        ),
+                    interactor = getNotificationsInteractor(context = context),
                     backgroundScope = bgScope,
                 )
                 .also { notificationsSnapshotRestorer = it }
@@ -340,9 +363,7 @@
             .get()
     }
 
-    override fun getClockPickerInteractor(
-        context: Context,
-    ): ClockPickerInteractor {
+    override fun getClockPickerInteractor(context: Context): ClockPickerInteractor {
         val appContext = context.applicationContext
         return clockPickerInteractor
             ?: ClockPickerInteractor(
@@ -397,9 +418,7 @@
                 }
     }
 
-    private fun getClockPickerSnapshotRestorer(
-        context: Context,
-    ): ClockPickerSnapshotRestorer {
+    private fun getClockPickerSnapshotRestorer(context: Context): ClockPickerSnapshotRestorer {
         return clockPickerSnapshotRestorer
             ?: ClockPickerSnapshotRestorer(getClockPickerInteractor(context)).also {
                 clockPickerSnapshotRestorer = it
@@ -408,7 +427,7 @@
 
     override fun getWallpaperColorResources(
         wallpaperColors: WallpaperColors,
-        context: Context
+        context: Context,
     ): WallpaperColorResources {
         return ThemedWallpaperColorResources(wallpaperColors, getSecureSettingsRepository(context))
     }
@@ -423,11 +442,11 @@
                     repository =
                         ColorPickerRepositoryImpl(
                             wallpaperColorsRepository,
-                            getColorCustomizationManager(appContext)
+                            getColorCustomizationManager(appContext),
                         ),
                     snapshotRestorer = {
                         getColorPickerSnapshotRestorer(appContext, wallpaperColorsRepository)
-                    }
+                    },
                 )
                 .also { colorPickerInteractor = it }
     }
@@ -463,9 +482,7 @@
             }
     }
 
-    fun getDarkModeSnapshotRestorer(
-        context: Context,
-    ): DarkModeSnapshotRestorer {
+    fun getDarkModeSnapshotRestorer(context: Context): DarkModeSnapshotRestorer {
         val appContext = context.applicationContext
         return darkModeSnapshotRestorer
             ?: DarkModeSnapshotRestorer(
@@ -476,9 +493,7 @@
                 .also { darkModeSnapshotRestorer = it }
     }
 
-    protected fun getThemedIconSnapshotRestorer(
-        context: Context,
-    ): ThemedIconSnapshotRestorer {
+    protected fun getThemedIconSnapshotRestorer(context: Context): ThemedIconSnapshotRestorer {
         val optionProvider = ThemedIconSwitchProvider.getInstance(context)
         return themedIconSnapshotRestorer
             ?: ThemedIconSnapshotRestorer(
@@ -493,10 +508,9 @@
 
     protected fun getThemedIconInteractor(): ThemedIconInteractor {
         return themedIconInteractor
-            ?: ThemedIconInteractor(
-                    repository = ThemeIconRepository(),
-                )
-                .also { themedIconInteractor = it }
+            ?: ThemedIconInteractor(repository = ThemeIconRepository()).also {
+                themedIconInteractor = it
+            }
     }
 
     override fun getClockSettingsViewModelFactory(
@@ -508,10 +522,7 @@
             ?: ClockSettingsViewModel.Factory(
                     context.applicationContext,
                     getClockPickerInteractor(context),
-                    getColorPickerInteractor(
-                        context,
-                        wallpaperColorsRepository,
-                    ),
+                    getColorPickerInteractor(context, wallpaperColorsRepository),
                     getUserEventLogger(),
                 ) { clockId ->
                     clockId?.let { clockViewFactory.getController(clockId).config.isReactiveToTone }
@@ -520,9 +531,7 @@
                 .also { clockSettingsViewModelFactory = it }
     }
 
-    fun getGridScreenViewModelFactory(
-        context: Context,
-    ): ViewModelProvider.Factory {
+    fun getGridScreenViewModelFactory(context: Context): ViewModelProvider.Factory {
         return gridScreenViewModelFactory
             ?: GridScreenViewModel.Factory(
                     context = context,
@@ -549,14 +558,11 @@
                 .also { gridInteractor = it }
     }
 
-    private fun getGridSnapshotRestorer(
-        context: Context,
-    ): GridSnapshotRestorer {
+    private fun getGridSnapshotRestorer(context: Context): GridSnapshotRestorer {
         return gridSnapshotRestorer
-            ?: GridSnapshotRestorer(
-                    interactor = getGridInteractor(context),
-                )
-                .also { gridSnapshotRestorer = it }
+            ?: GridSnapshotRestorer(interactor = getGridInteractor(context)).also {
+                gridSnapshotRestorer = it
+            }
     }
 
     override fun isCurrentSelectedColorPreset(context: Context): Boolean {
diff --git a/tests/common/src/com/android/customization/testing/TestCustomizationInjector.kt b/tests/common/src/com/android/customization/testing/TestCustomizationInjector.kt
index caa5029..108d85f 100644
--- a/tests/common/src/com/android/customization/testing/TestCustomizationInjector.kt
+++ b/tests/common/src/com/android/customization/testing/TestCustomizationInjector.kt
@@ -16,9 +16,16 @@
 import com.android.customization.picker.color.ui.viewmodel.ColorPickerViewModel
 import com.android.customization.picker.quickaffordance.domain.interactor.KeyguardQuickAffordancePickerInteractor
 import com.android.systemui.shared.clocks.ClockRegistry
+import com.android.wallpaper.module.NetworkStatusNotifier
+import com.android.wallpaper.module.PartnerProvider
+import com.android.wallpaper.module.WallpaperPreferences
 import com.android.wallpaper.module.logging.UserEventLogger
+import com.android.wallpaper.network.Requester
 import com.android.wallpaper.picker.customization.data.repository.WallpaperColorsRepository
+import com.android.wallpaper.picker.customization.domain.interactor.WallpaperInteractor
+import com.android.wallpaper.testing.FakeWallpaperClient
 import com.android.wallpaper.testing.TestInjector
+import com.android.wallpaper.util.DisplayUtils
 import javax.inject.Inject
 import javax.inject.Singleton
 
@@ -27,8 +34,26 @@
 @Inject
 constructor(
     private val customPrefs: TestDefaultCustomizationPreferences,
-    private val themesUserEventLogger: ThemesUserEventLogger
-) : TestInjector(themesUserEventLogger), CustomizationInjector {
+    private val themesUserEventLogger: ThemesUserEventLogger,
+    displayUtils: DisplayUtils,
+    requester: Requester,
+    networkStatusNotifier: NetworkStatusNotifier,
+    partnerProvider: PartnerProvider,
+    wallpaperClient: FakeWallpaperClient,
+    injectedWallpaperInteractor: WallpaperInteractor,
+    prefs: WallpaperPreferences,
+) :
+    TestInjector(
+        themesUserEventLogger,
+        displayUtils,
+        requester,
+        networkStatusNotifier,
+        partnerProvider,
+        wallpaperClient,
+        injectedWallpaperInteractor,
+        prefs,
+    ),
+    CustomizationInjector {
     /////////////////
     // CustomizationInjector implementations
     /////////////////
@@ -53,7 +78,7 @@
 
     override fun getWallpaperColorResources(
         wallpaperColors: WallpaperColors,
-        context: Context
+        context: Context,
     ): WallpaperColorResources {
         throw UnsupportedOperationException("not implemented")
     }
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 53ade86..8c57ab8 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
@@ -33,6 +33,10 @@
 import com.android.systemui.shared.keyguard.shared.model.KeyguardQuickAffordanceSlots
 import com.android.themepicker.R
 import com.android.wallpaper.module.InjectorProvider
+import com.android.wallpaper.module.NetworkStatusNotifier
+import com.android.wallpaper.module.PartnerProvider
+import com.android.wallpaper.module.WallpaperPreferences
+import com.android.wallpaper.network.Requester
 import com.android.wallpaper.picker.common.icon.ui.viewmodel.Icon
 import com.android.wallpaper.picker.common.text.ui.viewmodel.Text
 import com.android.wallpaper.picker.customization.data.repository.WallpaperRepository
@@ -44,6 +48,8 @@
 import com.android.wallpaper.testing.TestInjector
 import com.android.wallpaper.testing.TestWallpaperPreferences
 import com.android.wallpaper.testing.collectLastValue
+import com.android.wallpaper.util.DisplayUtils
+import com.android.wallpaper.util.DisplaysProvider
 import com.google.common.truth.Truth.assertThat
 import com.google.common.truth.Truth.assertWithMessage
 import kotlinx.coroutines.Dispatchers
@@ -58,6 +64,7 @@
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
+import org.mockito.Mockito.mock
 import org.robolectric.RobolectricTestRunner
 
 @OptIn(ExperimentalCoroutinesApi::class)
@@ -77,7 +84,6 @@
 
     @Before
     fun setUp() {
-        InjectorProvider.setInjector(TestInjector(logger))
         context = ApplicationProvider.getApplicationContext()
         val testDispatcher = StandardTestDispatcher()
         testScope = TestScope(testDispatcher)
@@ -108,8 +114,20 @@
                         client = FakeWallpaperClient(),
                         wallpaperPreferences = TestWallpaperPreferences(),
                         backgroundDispatcher = testDispatcher,
-                    ),
+                    )
             )
+        InjectorProvider.setInjector(
+            TestInjector(
+                logger,
+                DisplayUtils(context, mock(DisplaysProvider::class.java)),
+                mock(Requester::class.java),
+                mock(NetworkStatusNotifier::class.java),
+                mock(PartnerProvider::class.java),
+                FakeWallpaperClient(),
+                wallpaperInteractor,
+                mock(WallpaperPreferences::class.java),
+            )
+        )
         underTest =
             KeyguardQuickAffordancePickerViewModel.Factory(
                     context = context,
@@ -356,12 +374,12 @@
                         icon1 =
                             Icon.Loaded(
                                 FakeCustomizationProviderClient.ICON_1,
-                                Text.Loaded("Left shortcut")
+                                Text.Loaded("Left shortcut"),
                             ),
                         icon2 =
                             Icon.Loaded(
                                 FakeCustomizationProviderClient.ICON_3,
-                                Text.Loaded("Right shortcut")
+                                Text.Loaded("Right shortcut"),
                             ),
                     )
                 )
@@ -384,7 +402,7 @@
                         icon1 =
                             Icon.Loaded(
                                 FakeCustomizationProviderClient.ICON_1,
-                                Text.Loaded("Left shortcut")
+                                Text.Loaded("Left shortcut"),
                             ),
                         icon2 = null,
                     )
@@ -412,7 +430,7 @@
                         icon2 =
                             Icon.Loaded(
                                 FakeCustomizationProviderClient.ICON_3,
-                                Text.Loaded("Right shortcut")
+                                Text.Loaded("Right shortcut"),
                             ),
                     )
                 )
@@ -473,11 +491,7 @@
         assertThat(affordances).isNotNull()
         affordances?.forEach { affordance ->
             val nameMatchesSelectedName =
-                Text.evaluationEquals(
-                    context,
-                    affordance.text,
-                    Text.Loaded(selectedAffordanceText),
-                )
+                Text.evaluationEquals(context, affordance.text, Text.Loaded(selectedAffordanceText))
             val isSelected: Boolean? = collectLastValue(affordance.isSelected).invoke()
             assertWithMessage(
                     "Expected affordance with name \"${affordance.text}\" to have" +