diff --git a/Android.mk b/Android.mk
index 54fcfd05d21e074ec64e96f64cbc078e12b0d662..965c7a7c2c7cbc277bf42af57c91f39c6afb5fcd 100644
--- a/Android.mk
+++ b/Android.mk
@@ -1152,10 +1152,10 @@ sepolicy_tests := $(intermediates)/sepolicy_tests
 $(sepolicy_tests): PRIVATE_PLAT_FC := $(built_plat_fc)
 $(sepolicy_tests): PRIVATE_NONPLAT_FC := $(built_nonplat_fc)
 $(sepolicy_tests): PRIVATE_SEPOLICY := $(built_sepolicy)
-$(sepolicy_tests): $(HOST_OUT_EXECUTABLES)/sepolicy_tests.py \
+$(sepolicy_tests): $(HOST_OUT_EXECUTABLES)/sepolicy_tests \
 $(built_plat_fc) $(built_nonplat_fc) $(built_sepolicy)
 	@mkdir -p $(dir $@)
-	$(hide) python $(HOST_OUT_EXECUTABLES)/sepolicy_tests.py -l $(HOST_OUT)/lib64 -f $(PRIVATE_PLAT_FC) -f $(PRIVATE_NONPLAT_FC) -p $(PRIVATE_SEPOLICY)
+	$(hide) python $(HOST_OUT_EXECUTABLES)/sepolicy_tests -l $(HOST_OUT)/lib64 -f $(PRIVATE_PLAT_FC) -f $(PRIVATE_NONPLAT_FC) -p $(PRIVATE_SEPOLICY)
 	$(hide) touch $@
 
 ##################################
@@ -1262,11 +1262,11 @@ $(treble_sepolicy_tests): PRIVATE_SEPOLICY := $(built_sepolicy)
 $(treble_sepolicy_tests): PRIVATE_SEPOLICY_OLD := $(built_26.0_plat_sepolicy)
 $(treble_sepolicy_tests): PRIVATE_COMBINED_MAPPING := $(26.0_mapping.combined.cil)
 $(treble_sepolicy_tests): PRIVATE_PLAT_SEPOLICY := $(built_plat_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) $(built_plat_sepolicy) \
 $(built_26.0_plat_sepolicy) $(26.0_compat) $(26.0_mapping.combined.cil)
 	@mkdir -p $(dir $@)
-	$(hide) python $(HOST_OUT_EXECUTABLES)/treble_sepolicy_tests.py -l \
+	$(hide) python $(HOST_OUT_EXECUTABLES)/treble_sepolicy_tests -l \
 		$(HOST_OUT)/lib64 -f $(PRIVATE_PLAT_FC) -f $(PRIVATE_NONPLAT_FC) \
 		-b $(PRIVATE_PLAT_SEPOLICY) -m $(PRIVATE_COMBINED_MAPPING) \
 		-o $(PRIVATE_SEPOLICY_OLD) -p $(PRIVATE_SEPOLICY)
diff --git a/tests/Android.bp b/tests/Android.bp
index 5b2f78d2b6284fcacaadde292ec8bbfa32af2f86..ec4b465cf35f142e40399e98c2c6fb5741e187b2 100644
--- a/tests/Android.bp
+++ b/tests/Android.bp
@@ -15,29 +15,35 @@ cc_library_host_shared {
     },
 }
 
-cc_prebuilt_binary {
-    name: "mini_parser.py",
-    srcs: ["mini_parser.py"],
-    host_supported: true,
+python_defaults {
+    name: "py2_only",
+    version: {
+        py2: {
+            enabled: true,
+        },
+        py3: {
+            enabled: false,
+        },
+    },
 }
 
-cc_prebuilt_binary {
-    name: "policy.py",
-    srcs: ["policy.py"],
-    host_supported: true,
+python_binary_host {
+    name: "treble_sepolicy_tests",
+    srcs: [
+        "treble_sepolicy_tests.py",
+        "mini_parser.py",
+        "policy.py",
+    ],
     required: ["libsepolwrap"],
+    defaults: ["py2_only"],
 }
 
-cc_prebuilt_binary {
-    name: "treble_sepolicy_tests.py",
-    srcs: ["treble_sepolicy_tests.py"],
-    host_supported: true,
-    required: ["mini_parser.py", "policy.py"],
-}
-
-cc_prebuilt_binary {
-    name: "sepolicy_tests.py",
-    srcs: ["sepolicy_tests.py"],
-    host_supported: true,
-    required: ["policy.py"],
+python_binary_host {
+    name: "sepolicy_tests",
+    srcs: [
+        "sepolicy_tests.py",
+        "policy.py",
+    ],
+    required: ["libsepolwrap"],
+    defaults: ["py2_only"],
 }