diff --git a/tests/Android.bp b/tests/Android.bp
new file mode 100644
index 0000000..9fed740
--- /dev/null
+++ b/tests/Android.bp
@@ -0,0 +1,47 @@
+// Copyright (C) 2022 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.
+//
+
+//
+// Build rule for WallpaperPicker2 tests
+//
+package {
+    default_applicable_licenses: ["Android-Apache-2.0"],
+}
+
+android_test {
+    name: "ThemePickerTests",
+
+    defaults: ["ThemePicker_defaults"],
+    srcs: [
+        "src/**/*.java",
+        "src/**/*.kt",
+    ],
+    static_libs: [
+        "WallpaperPicker2TestLib",
+        "androidx.test.rules",
+        "junit",
+        "kotlinx_coroutines_test",
+        "truth-prebuilt",
+    ],
+    libs: [
+        "android.test.runner",
+        "android.test.base",
+        "android.test.mock",
+    ],
+
+    kotlincflags: ["-Xjvm-default=enable"],
+    platform_apis: true,
+    test_suites: ["device-tests"],
+}
diff --git a/tests/AndroidManifest.xml b/tests/AndroidManifest.xml
new file mode 100644
index 0000000..2fd6b3f
--- /dev/null
+++ b/tests/AndroidManifest.xml
@@ -0,0 +1,27 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!-- Copyright (C) 2022 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.
+-->
+<manifest xmlns:android="http://schemas.android.com/apk/res/android"
+    package="com.android.wallpaper">
+
+    <application>
+        <uses-library android:name="android.test.runner" />
+    </application>
+
+    <instrumentation android:name="androidx.test.runner.AndroidJUnitRunner"
+        android:targetPackage="com.android.wallpaper"
+        android:label="Tests for ThemePicker" />
+
+</manifest>
diff --git a/tests/AndroidTest.xml b/tests/AndroidTest.xml
new file mode 100644
index 0000000..be2119b
--- /dev/null
+++ b/tests/AndroidTest.xml
@@ -0,0 +1,28 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!-- Copyright (C) 2022 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.
+-->
+<configuration description="Runs Tests for ThemePicker.">
+    <target_preparer class="com.android.tradefed.targetprep.TestAppInstallSetup">
+        <option name="test-file-name" value="ThemePickerTests.apk" />
+    </target_preparer>
+
+    <option name="test-suite-tag" value="apct" />
+    <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="hidden-api-checks" value="false"/>
+    </test>
+</configuration>
diff --git a/tests/src/com/android/customization/model/picker/quickaffordance/data/repository/KeyguardQuickAffordancePickerRepositoryTest.kt b/tests/src/com/android/customization/model/picker/quickaffordance/data/repository/KeyguardQuickAffordancePickerRepositoryTest.kt
index 4a88f3b..5e4fa6e 100644
--- a/tests/src/com/android/customization/model/picker/quickaffordance/data/repository/KeyguardQuickAffordancePickerRepositoryTest.kt
+++ b/tests/src/com/android/customization/model/picker/quickaffordance/data/repository/KeyguardQuickAffordancePickerRepositoryTest.kt
@@ -72,7 +72,7 @@
                 com.android.systemui.shared.quickaffordance.data.content
                     .KeyguardQuickAffordanceProviderContract
                     .FlagsTable
-                    .FLAG_NAME_FEATURE_ENABLED,
+                    .FLAG_NAME_CUSTOM_LOCK_SCREEN_QUICK_AFFORDANCES_ENABLED,
                 true,
             )
             val values = mutableListOf<Boolean>()
@@ -90,7 +90,7 @@
                 com.android.systemui.shared.quickaffordance.data.content
                     .KeyguardQuickAffordanceProviderContract
                     .FlagsTable
-                    .FLAG_NAME_FEATURE_ENABLED,
+                    .FLAG_NAME_CUSTOM_LOCK_SCREEN_QUICK_AFFORDANCES_ENABLED,
                 false,
             )
             val values = mutableListOf<Boolean>()