diff --git a/tests/Android.bp b/tests/Android.bp
index de86002..1ce577b 100644
--- a/tests/Android.bp
+++ b/tests/Android.bp
@@ -15,16 +15,24 @@
     },
 }
 
-cc_prebuilt_binary {
-    name: "policy.py",
-    srcs: ["policy.py"],
-    host_supported: true,
-    required: ["libsepolwrap"],
+python_defaults {
+    name: "py2_only",
+    version: {
+        py2: {
+            enabled: true,
+        },
+        py3: {
+            enabled: false,
+        },
+    },
 }
 
-cc_prebuilt_binary {
-    name: "treble_sepolicy_tests.py",
-    srcs: ["treble_sepolicy_tests.py"],
-    host_supported: true,
-    required: ["policy.py"],
+python_binary_host {
+    name: "treble_sepolicy_tests",
+    srcs: [
+        "treble_sepolicy_tests.py",
+        "policy.py",
+    ],
+    required: ["libsepolwrap"],
+    defaults: ["py2_only"],
 }