diff --git a/Android.mk b/Android.mk
index 93ce6b876..b6c4ffee6 100644
--- a/Android.mk
+++ b/Android.mk
@@ -1150,10 +1150,10 @@ treble_sepolicy_tests := $(intermediates)/treble_sepolicy_tests
 $(treble_sepolicy_tests): PRIVATE_PLAT_FC := $(built_plat_fc)
 $(treble_sepolicy_tests): PRIVATE_NONPLAT_FC := $(built_nonplat_fc)
 $(treble_sepolicy_tests): PRIVATE_SEPOLICY := $(built_sepolicy)
-$(treble_sepolicy_tests): $(HOST_OUT_EXECUTABLES)/treble_sepolicy_tests.py \
+$(treble_sepolicy_tests): $(HOST_OUT_EXECUTABLES)/treble_sepolicy_tests \
 $(built_plat_fc) $(built_nonplat_fc) $(built_sepolicy)
 	@mkdir -p $(dir $@)
-	$(hide) python $(HOST_OUT_EXECUTABLES)/treble_sepolicy_tests.py -l $(HOST_OUT)/lib64 -f $(PRIVATE_PLAT_FC) -f $(PRIVATE_NONPLAT_FC) -p $(PRIVATE_SEPOLICY)
+	$(hide) python $(HOST_OUT_EXECUTABLES)/treble_sepolicy_tests -l $(HOST_OUT)/lib64 -f $(PRIVATE_PLAT_FC) -f $(PRIVATE_NONPLAT_FC) -p $(PRIVATE_SEPOLICY)
 	$(hide) touch $@
 endif # ($(PRODUCT_FULL_TREBLE),true)
 #################################
diff --git a/tests/Android.bp b/tests/Android.bp
index de8600200..1ce577b88 100644
--- a/tests/Android.bp
+++ b/tests/Android.bp
@@ -15,16 +15,24 @@ cc_library_host_shared {
     },
 }
 
-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"],
 }