diff --git a/tests/Android.bp b/tests/Android.bp
index 9df2720..2ab3ebc 100644
--- a/tests/Android.bp
+++ b/tests/Android.bp
@@ -20,17 +20,43 @@
     default_applicable_licenses: ["Android-Apache-2.0"],
 }
 
+filegroup {
+    name: "ThemePickerTests_srcs",
+
+    visibility: [
+        "//visibility:public"
+    ],
+
+    srcs: [
+        "src/**/*.java",
+        "src/**/*.kt",
+    ],
+}
+
 android_test {
     name: "ThemePickerTests",
 
     defaults: ["ThemePicker_defaults"],
     srcs: [
-        "src/**/*.java",
-        "src/**/*.kt",
+        ":ThemePickerTests_srcs",
+        ":WallpaperPicker2Tests_srcs",
+        "module/src/com/android/customization/TestModule.kt",
     ],
+    exclude_srcs: [":ThemePicker_src_prod"],
     static_libs: [
         "WallpaperPicker2TestLib",
+        "WallpaperPicker2TestRunner",
+        "ThemePickerTestLib",
+        "androidx.test.espresso.core",
+        "androidx.test.espresso.contrib",
+        "androidx.test.espresso.intents",
+        "androidx.test.ext.junit",
         "androidx.test.rules",
+        "hamcrest-library",
+        "hamcrest",
+        "hilt_android",
+        "hilt_android_testing",
+        "mockito-target-minus-junit4",
         "junit",
         "kotlinx_coroutines_test",
         "truth-prebuilt",
diff --git a/tests/AndroidManifest.xml b/tests/AndroidManifest.xml
index 2fd6b3f..f1c0cf2 100644
--- a/tests/AndroidManifest.xml
+++ b/tests/AndroidManifest.xml
@@ -18,9 +18,17 @@
 
     <application>
         <uses-library android:name="android.test.runner" />
+
+        <activity android:name="com.android.wallpaper.picker.PreviewActivity"
+            android:resizeableActivity="false">
+        </activity>
+
+        <activity android:name="com.android.wallpaper.picker.preview.ui.WallpaperPreviewActivity"
+            android:resizeableActivity="false">
+        </activity>
     </application>
 
-    <instrumentation android:name="androidx.test.runner.AndroidJUnitRunner"
+    <instrumentation android:name="com.android.wallpaper.CustomTestRunner"
         android:targetPackage="com.android.wallpaper"
         android:label="Tests for ThemePicker" />
 
diff --git a/tests/AndroidTest.xml b/tests/AndroidTest.xml
index be2119b..18e95ef 100644
--- a/tests/AndroidTest.xml
+++ b/tests/AndroidTest.xml
@@ -22,7 +22,7 @@
     <option name="test-tag" value="ThemePickerTests" />
     <test class="com.android.tradefed.testtype.AndroidJUnitTest" >
         <option name="package" value="com.android.wallpaper" />
-        <option name="runner" value="androidx.test.runner.AndroidJUnitRunner" />
+        <option name="runner" value="com.android.wallpaper.CustomTestRunner" />
         <option name="hidden-api-checks" value="false"/>
     </test>
 </configuration>
diff --git a/tests/common/src/com/android/customization/testing/TestCustomizationInjector.kt b/tests/common/src/com/android/customization/testing/TestCustomizationInjector.kt
index cd587d6..18175b3 100644
--- a/tests/common/src/com/android/customization/testing/TestCustomizationInjector.kt
+++ b/tests/common/src/com/android/customization/testing/TestCustomizationInjector.kt
@@ -23,8 +23,10 @@
 import javax.inject.Singleton
 
 @Singleton
-open class TestCustomizationInjector @Inject constructor() : TestInjector(), CustomizationInjector {
-    private var customizationPrefs: CustomizationPreferences? = null
+open class TestCustomizationInjector
+@Inject
+constructor(private val customPrefs: TestDefaultCustomizationPreferences) :
+    TestInjector(), CustomizationInjector {
     private var themesUserEventLogger: ThemesUserEventLogger? = null
 
     /////////////////
@@ -32,8 +34,7 @@
     /////////////////
 
     override fun getCustomizationPreferences(context: Context): CustomizationPreferences {
-        return customizationPrefs
-            ?: TestDefaultCustomizationPreferences(context).also { customizationPrefs = it }
+        return customPrefs
     }
 
     override fun getKeyguardQuickAffordancePickerInteractor(
diff --git a/tests/common/src/com/android/customization/testing/TestDefaultCustomizationPreferences.java b/tests/common/src/com/android/customization/testing/TestDefaultCustomizationPreferences.java
index 81890f0..c8e31a3 100644
--- a/tests/common/src/com/android/customization/testing/TestDefaultCustomizationPreferences.java
+++ b/tests/common/src/com/android/customization/testing/TestDefaultCustomizationPreferences.java
@@ -17,21 +17,32 @@
 
 import android.content.Context;
 
+import com.android.customization.module.CustomizationPreferences;
 import com.android.customization.module.DefaultCustomizationPreferences;
+import com.android.wallpaper.testing.TestWallpaperPreferences;
+
+import dagger.hilt.android.qualifiers.ApplicationContext;
 
 import java.util.HashSet;
 import java.util.Set;
 
+import javax.inject.Inject;
+import javax.inject.Singleton;
+
 /**
  * Test implementation of {@link DefaultCustomizationPreferences}.
  */
-public class TestDefaultCustomizationPreferences extends DefaultCustomizationPreferences {
+@Singleton
+public class TestDefaultCustomizationPreferences extends TestWallpaperPreferences implements
+        CustomizationPreferences {
 
     private String mCustomThemes;
     private final Set<String> mTabVisited = new HashSet<>();
+    private boolean mThemedIconEnabled = false;
 
-    public TestDefaultCustomizationPreferences(Context context) {
-        super(context);
+    @Inject
+    public TestDefaultCustomizationPreferences(@ApplicationContext Context context) {
+        super();
     }
 
     @Override
@@ -53,4 +64,14 @@
     public void setTabVisited(String id) {
         mTabVisited.add(id);
     }
+
+    @Override
+    public boolean getThemedIconEnabled() {
+        return mThemedIconEnabled;
+    }
+
+    @Override
+    public void setThemedIconEnabled(boolean enabled) {
+        mThemedIconEnabled = enabled;
+    }
 }
diff --git a/tests/module/src/com/android/customization/TestModule.kt b/tests/module/src/com/android/customization/TestModule.kt
new file mode 100644
index 0000000..d35c89a
--- /dev/null
+++ b/tests/module/src/com/android/customization/TestModule.kt
@@ -0,0 +1,50 @@
+package com.android.customization
+
+import com.android.customization.module.CustomizationInjector
+import com.android.customization.module.CustomizationPreferences
+import com.android.customization.testing.TestCustomizationInjector
+import com.android.customization.testing.TestDefaultCustomizationPreferences
+import com.android.wallpaper.module.AppModule
+import com.android.wallpaper.module.Injector
+import com.android.wallpaper.module.WallpaperPreferences
+import com.android.wallpaper.testing.TestInjector
+import com.android.wallpaper.testing.TestWallpaperPreferences
+import dagger.Binds
+import dagger.Module
+import dagger.hilt.components.SingletonComponent
+import dagger.hilt.testing.TestInstallIn
+import javax.inject.Singleton
+
+@Module
+@TestInstallIn(components = [SingletonComponent::class], replaces = [AppModule::class])
+abstract class TestModule {
+    //// WallpaperPicker2 prod
+
+    @Binds @Singleton abstract fun bindInjector(impl: TestCustomizationInjector): Injector
+
+    @Binds
+    @Singleton
+    abstract fun bindWallpaperPrefs(impl: TestDefaultCustomizationPreferences): WallpaperPreferences
+
+    //// WallpaperPicker2 test
+
+    @Binds @Singleton abstract fun bindTestInjector(impl: TestCustomizationInjector): TestInjector
+
+    @Binds
+    @Singleton
+    abstract fun bindTestWallpaperPrefs(
+        impl: TestDefaultCustomizationPreferences
+    ): TestWallpaperPreferences
+
+    //// ThemePicker prod
+
+    @Binds
+    @Singleton
+    abstract fun bindCustomizationInjector(impl: TestCustomizationInjector): CustomizationInjector
+
+    @Binds
+    @Singleton
+    abstract fun bindCustomizationPrefs(
+        impl: TestDefaultCustomizationPreferences
+    ): CustomizationPreferences
+}