diff --git a/tests/Android.bp b/tests/Android.bp
index 690a1b6..5b12a4a 100644
--- a/tests/Android.bp
+++ b/tests/Android.bp
@@ -41,13 +41,13 @@
     srcs: [
         ":ThemePickerTests_srcs",
         ":WallpaperPicker2Tests_srcs",
-        "module/src/com/android/customization/TestModule.kt",
     ],
     exclude_srcs: [":ThemePicker_src_prod"],
     static_libs: [
         "WallpaperPicker2TestLib",
         "WallpaperPicker2TestRunner",
         "ThemePickerTestLib",
+        "ThemePickerTestModule",
         "SystemUICustomizationTestUtils",
         "androidx.test.espresso.core",
         "androidx.test.espresso.contrib",
diff --git a/tests/module/Android.bp b/tests/module/Android.bp
index c174699..29d4965 100644
--- a/tests/module/Android.bp
+++ b/tests/module/Android.bp
@@ -19,6 +19,7 @@
 
 android_library {
     name: "ThemePickerTestModule",
+    defaults: ["ThemePicker_defaults_with_overrides"],
     srcs: ["src/com/android/customization/TestModule.kt"],
     static_libs: [
         "ThemePickerTestLib",