diff --git a/tests/python/s_tir/schedule/test_tir_schedule_read_write_at.py b/tests/python/s_tir/schedule/test_tir_schedule_read_write_at.py
index 85e6a7a0e0ae..0db98fac2919 100644
--- a/tests/python/s_tir/schedule/test_tir_schedule_read_write_at.py
+++ b/tests/python/s_tir/schedule/test_tir_schedule_read_write_at.py
@@ -37,6 +37,7 @@
 import pytest
 
 import tvm
+import tvm.testing
 from tvm import tirx
 from tvm.s_tir.schedule.testing import (
     assert_structural_equal_ignore_global_symbol,
diff --git a/tests/python/s_tir/schedule/test_tir_schedule_storage_align.py b/tests/python/s_tir/schedule/test_tir_schedule_storage_align.py
index 2280292ec1c9..fd64f2ad8a22 100644
--- a/tests/python/s_tir/schedule/test_tir_schedule_storage_align.py
+++ b/tests/python/s_tir/schedule/test_tir_schedule_storage_align.py
@@ -19,6 +19,7 @@
 import pytest
 
 import tvm
+import tvm.testing
 from tvm import tirx
 from tvm.s_tir.schedule.testing import (
     assert_structural_equal_ignore_global_symbol,
@@ -173,11 +174,4 @@ def test_storage_align_invalid_annotation():
 
 
 if __name__ == "__main__":
-    test_storage_align()
-    test_storage_align_update()
-    test_storage_align_invalid_factor1()
-    test_storage_align_invalid_factor2()
-    test_storage_align_invalid_buffer()
-    test_storage_align_invalid_buffer_index()
-    test_storage_align_invalid_axis()
-    test_storage_align_invalid_annotation()
+    tvm.testing.main()