diff --git a/tests/module/Android.bp b/tests/module/Android.bp
index 29d4965..f865832 100644
--- a/tests/module/Android.bp
+++ b/tests/module/Android.bp
@@ -20,7 +20,7 @@
 android_library {
     name: "ThemePickerTestModule",
     defaults: ["ThemePicker_defaults_with_overrides"],
-    srcs: ["src/com/android/customization/TestModule.kt"],
+    srcs: ["src/com/android/wallpaper/ThemePickerTestModule.kt"],
     static_libs: [
         "ThemePickerTestLib",
         "hilt_android",
diff --git a/tests/module/src/com/android/customization/TestModule.kt b/tests/module/src/com/android/wallpaper/ThemePickerTestModule.kt
similarity index 81%
rename from tests/module/src/com/android/customization/TestModule.kt
rename to tests/module/src/com/android/wallpaper/ThemePickerTestModule.kt
index c39d3e3..687ec9b 100644
--- a/tests/module/src/com/android/customization/TestModule.kt
+++ b/tests/module/src/com/android/wallpaper/ThemePickerTestModule.kt
@@ -1,4 +1,19 @@
-package com.android.customization
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.android.wallpaper
 
 import androidx.test.core.app.ApplicationProvider
 import com.android.customization.model.color.ColorCustomizationManager
@@ -9,12 +24,12 @@
 import com.android.customization.module.logging.ThemesUserEventLogger
 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.PartnerProvider
 import com.android.wallpaper.module.WallpaperPreferences
 import com.android.wallpaper.module.logging.TestUserEventLogger
 import com.android.wallpaper.module.logging.UserEventLogger
+import com.android.wallpaper.modules.ThemePickerAppModule
 import com.android.wallpaper.picker.preview.data.util.DefaultLiveWallpaperDownloader
 import com.android.wallpaper.picker.preview.data.util.LiveWallpaperDownloader
 import com.android.wallpaper.testing.TestInjector
@@ -30,8 +45,8 @@
 import javax.inject.Singleton
 
 @Module
-@TestInstallIn(components = [SingletonComponent::class], replaces = [AppModule::class])
-abstract class TestModule {
+@TestInstallIn(components = [SingletonComponent::class], replaces = [ThemePickerAppModule::class])
+abstract class ThemePickerTestModule {
     //// WallpaperPicker2 prod
 
     @Binds @Singleton abstract fun bindInjector(impl: TestCustomizationInjector): Injector