Add utility to replace direct call to pytest.main

python/tvm/testing/utils.py

@@ -62,6 +62,7 @@ def test_something():
 `TVM_TEST_TARGETS` environment variable in the CI.
 
 """
+import inspect
 import copy
 import copyreg
 import ctypes
@@ -1513,3 +1514,8 @@ def identity_after(x, sleep):
 def terminate_self():
     """Testing function to terminate the process."""
     sys.exit(-1)
+
+
+def main():
+    test_file = inspect.getsourcefile(sys._getframe(1))
+    sys.exit(pytest.main([test_file] + sys.argv[1:]))

tests/micro/arduino/test_arduino_error_detection.py

@@ -21,6 +21,7 @@
 from tvm.micro.project_api.server import ServerError
 
 import test_utils
+import tvm.testing
 
 # A new project and workspace dir is created for EVERY test
 @pytest.fixture
@@ -46,4 +47,4 @@ def test_bugged_project_compile_fails(workspace_dir, project):
 
 
 if __name__ == "__main__":
-    sys.exit(pytest.main([__file__] + sys.argv[1:]))
+    tvm.testing.main()

tests/micro/arduino/test_arduino_rpc_server.py

@@ -29,6 +29,7 @@
 import pytest
 
 import tvm
+import tvm.testing
 from PIL import Image
 from tvm import relay
 from tvm.relay.testing import byoc
@@ -367,4 +368,4 @@ def test_tensors(sess):
 
 
 if __name__ == "__main__":
-    sys.exit(pytest.main([__file__] + sys.argv[1:]))
+    tvm.testing.main()

tests/micro/arduino/test_arduino_workflow.py

@@ -21,6 +21,8 @@
 import sys
 import pytest
 
+import tvm.testing
+
 import test_utils
 
 """
@@ -218,4 +220,4 @@ def test_project_inference_runtime(serial_output):
 
 
 if __name__ == "__main__":
-    sys.exit(pytest.main([__file__] + sys.argv[1:]))
+    tvm.testing.main()

tests/micro/common/test_tvmc.py

@@ -26,6 +26,7 @@
 import shutil
 
 import tvm
+import tvm.testing
 from tvm.contrib.download import download_testdata
 
 from ..zephyr.test_utils import ZEPHYR_BOARDS
@@ -217,4 +218,4 @@ def test_tvmc_model_run(board, output_dir):
 
 
 if __name__ == "__main__":
-    sys.exit(pytest.main([__file__] + sys.argv[1:]))
+    tvm.testing.main()

tests/micro/zephyr/test_zephyr.py

@@ -26,6 +26,7 @@
 from PIL import Image
 
 import tvm
+import tvm.testing
 import tvm.relay as relay
 from tvm.relay.backend import Executor, Runtime
 from tvm.relay.testing import byoc
@@ -504,4 +505,4 @@ def test_autotune_conv2d(temp_dir, board, west_cmd, tvm_debug):
 
 
 if __name__ == "__main__":
-    sys.exit(pytest.main([__file__] + sys.argv[1:]))
+    tvm.testing.main()

tests/micro/zephyr/test_zephyr_aot.py

@@ -135,4 +135,4 @@ def test_qemu_make_fail(temp_dir, board, west_cmd, tvm_debug):
 
 
 if __name__ == "__main__":
-    sys.exit(pytest.main([__file__] + sys.argv[1:]))
+    tvm.testing.main()

tests/micro/zephyr/test_zephyr_armv7m.py

@@ -187,4 +187,4 @@ def test_armv7m_intrinsic(temp_dir, board, west_cmd, tvm_debug):
 
 
 if __name__ == "__main__":
-    sys.exit(pytest.main([__file__] + sys.argv[1:]))
+    tvm.testing.main()

tests/python/ci/test_ci.py

@@ -20,6 +20,7 @@
 import json
 import textwrap
 import pytest
+import tvm.testing
 
 from test_utils import REPO_ROOT
 
@@ -732,4 +733,4 @@ def run(type, data, check):
 
 
 if __name__ == "__main__":
-    sys.exit(pytest.main([__file__] + sys.argv[1:]))
+    tvm.testing.main()

tests/python/contrib/test_cmsisnn/test_binary_ops.py

@@ -275,4 +275,4 @@ def test_invalid_parameters(
 
 
 if __name__ == "__main__":
-    sys.exit(pytest.main([__file__] + sys.argv[1:]))
+    tvm.testing.main()

tests/python/contrib/test_cmsisnn/test_conv2d.py

@@ -512,4 +512,4 @@ def test_invalid_parameters(
 
 
 if __name__ == "__main__":
-    sys.exit(pytest.main([__file__] + sys.argv[1:]))
+    tvm.testing.main()

tests/python/contrib/test_cmsisnn/test_extract_constants.py

@@ -21,6 +21,7 @@
 import numpy as np
 import pytest
 import tvm
+import tvm.testing
 from tvm import relay
 
 tvm._ffi._init_api("relay.ext.cmsisnn.transform", __name__)
@@ -224,4 +225,4 @@ def test_multiple_functions_non_cmsisnn_compiler(external_compiler):
 
 
 if __name__ == "__main__":
-    sys.exit(pytest.main([__file__] + sys.argv[1:]))
+    tvm.testing.main()

tests/python/contrib/test_cmsisnn/test_fully_connected.py

@@ -234,4 +234,4 @@ def test_invalid_parameters(
 
 
 if __name__ == "__main__":
-    sys.exit(pytest.main([__file__] + sys.argv[1:]))
+    tvm.testing.main()

tests/python/contrib/test_cmsisnn/test_generate_constants.py

@@ -21,6 +21,7 @@
 import numpy as np
 import pytest
 import tvm
+import tvm.testing
 from tvm import relay
 from tvm.relay.op.contrib import cmsisnn
 
@@ -225,4 +226,4 @@ def test_op_int8(
 
 
 if __name__ == "__main__":
-    sys.exit(pytest.main([__file__] + sys.argv[1:]))
+    tvm.testing.main()

tests/python/contrib/test_cmsisnn/test_networks.py

@@ -117,4 +117,4 @@ def test_cnn_small(test_runner):
 
 
 if __name__ == "__main__":
-    sys.exit(pytest.main([__file__] + sys.argv[1:]))
+    tvm.testing.main()

tests/python/contrib/test_cmsisnn/test_pooling.py

@@ -169,6 +169,4 @@ def test_invalid_layout(op):
 
 
 if __name__ == "__main__":
-    import sys
-
-    sys.exit(pytest.main([__file__] + sys.argv[1:]))
+    tvm.testing.main()

tests/python/contrib/test_cmsisnn/test_scalar_to_tensor_constant.py

@@ -21,6 +21,7 @@
 import numpy as np
 import pytest
 import tvm
+import tvm.testing
 from tvm import relay
 
 tvm._ffi._init_api("relay.ext.cmsisnn.transform", __name__)
@@ -286,4 +287,4 @@ def get_mod():
 
 
 if __name__ == "__main__":
-    sys.exit(pytest.main([__file__] + sys.argv[1:]))
+    tvm.testing.main()

tests/python/contrib/test_cmsisnn/test_softmax.py

@@ -129,4 +129,4 @@ def test_invalid_parameters(in_dtype, out_dtype, zero_point, scale, out_zero_poi
 
 
 if __name__ == "__main__":
-    sys.exit(pytest.main([__file__] + sys.argv[1:]))
+    tvm.testing.main()

tests/python/contrib/test_cudnn.py

@@ -20,6 +20,7 @@
 import pytest
 
 import tvm
+import tvm.testing
 from tvm import te
 from tvm import relay
 from tvm.contrib import cudnn
@@ -623,4 +624,4 @@ def test_relay_cudnn_conv2d_bias_act(
 
 
 if __name__ == "__main__":
-    sys.exit(pytest.main(sys.argv))
+    tvm.testing.main()

tests/python/contrib/test_dnnl.py

@@ -985,6 +985,4 @@ def get_graph():
 
 
 if __name__ == "__main__":
-    import sys
-
-    sys.exit(pytest.main([__file__] + sys.argv[1:]))
+    tvm.testing.main()

tests/python/contrib/test_hexagon/conv2d/test_conv2d_blocked.py

@@ -19,6 +19,7 @@
 
 import platform
 import tvm
+import tvm.testing
 from tvm import te
 from tvm import topi
 from tvm.topi import testing
@@ -191,4 +192,4 @@ def test_conv2d(
 
 
 if __name__ == "__main__":
-    sys.exit(pytest.main(sys.argv))
+    tvm.testing.main()

tests/python/contrib/test_hexagon/conv2d/test_conv2d_conv2d.py

@@ -19,6 +19,7 @@
 
 import platform
 import tvm
+import tvm.testing
 from tvm import te
 from tvm import topi
 from tvm.topi import testing
@@ -236,4 +237,4 @@ def test_conv2d(
 
 
 if __name__ == "__main__":
-    sys.exit(pytest.main(sys.argv))
+    tvm.testing.main()

tests/python/contrib/test_hexagon/test_2d_physical_buffers.py

@@ -336,4 +336,4 @@ def test_execute(
 
 
 if __name__ == "__main__":
-    sys.exit(pytest.main(sys.argv))
+    tvm.testing.main()

tests/python/contrib/test_hexagon/test_launcher.py

@@ -388,4 +388,4 @@ def test_aot_executor_multiple_conv2d(hexagon_session: Session, aot_host_target,
 
 
 if __name__ == "__main__":
-    sys.exit(pytest.main(sys.argv))
+    tvm.testing.main()

tests/python/contrib/test_hexagon/test_maxpool2d_blocked.py

@@ -18,6 +18,7 @@
 import sys
 
 import tvm
+import tvm.testing
 from tvm import te
 from tvm import topi
 from tvm.topi import testing
@@ -150,4 +151,4 @@ def test_maxpool(self, shape_nhwc, window_size, stride, pad, dtype, target):
 
 
 if __name__ == "__main__":
-    sys.exit(pytest.main(sys.argv))
+    tvm.testing.main()

tests/python/contrib/test_hexagon/test_models.py

@@ -135,4 +135,4 @@ def test_mobilenet_aot(hexagon_session: Session, aot_host_target, aot_target, en
 
 
 if __name__ == "__main__":
-    sys.exit(pytest.main(sys.argv))
+    tvm.testing.main()

tests/python/contrib/test_hexagon/test_usmp.py

@@ -109,4 +109,4 @@ def test_conv2d(hexagon_session: Session, aot_host_target, aot_target, usmp_enab
 
 
 if __name__ == "__main__":
-    sys.exit(pytest.main(sys.argv))
+    tvm.testing.main()

tests/python/contrib/test_hexagon/topi/test_batch_matmul.py

@@ -20,6 +20,7 @@
 import sys
 
 import tvm
+import tvm.testing
 from tvm import topi
 from tvm import te
 from tvm.contrib.hexagon.session import Session
@@ -139,4 +140,4 @@ def get_ref_data():
 
 
 if __name__ == "__main__":
-    sys.exit(pytest.main(sys.argv))
+    tvm.testing.main()

tests/python/contrib/test_hexagon/topi/test_conv2d_nchw.py

@@ -20,6 +20,7 @@
 import sys
 
 import tvm
+import tvm.testing
 from tvm import topi
 from tvm import te
 from tvm.contrib.hexagon.session import Session
@@ -242,4 +243,4 @@ class TestAsymmetricPadding(BaseConv2DTests):
 
 
 if __name__ == "__main__":
-    sys.exit(pytest.main(sys.argv))
+    tvm.testing.main()

tests/python/contrib/test_hexagon/topi/test_conv2d_nhwc.py

@@ -20,6 +20,7 @@
 import sys
 
 import tvm
+import tvm.testing
 from tvm import topi
 from tvm import te
 from tvm.contrib.hexagon.session import Session
@@ -121,4 +122,4 @@ class TestConv2dNHWC(BaseConv2DTests):
 
 
 if __name__ == "__main__":
-    sys.exit(pytest.main(sys.argv))
+    tvm.testing.main()

tests/python/contrib/test_hexagon/topi/test_dense.py

@@ -20,6 +20,7 @@
 import sys
 
 import tvm
+import tvm.testing
 from tvm import topi
 from tvm import te
 from tvm.contrib.hexagon.session import Session
@@ -115,4 +116,4 @@ def test_dense(
 
 
 if __name__ == "__main__":
-    sys.exit(pytest.main(sys.argv))
+    tvm.testing.main()

tests/python/contrib/test_hexagon/topi/test_pooling.py

@@ -20,6 +20,7 @@
 import sys
 
 import tvm
+import tvm.testing
 from tvm import topi
 from tvm import te
 from tvm.contrib.hexagon.session import Session
@@ -736,4 +737,4 @@ def test_pool3d(
 
 
 if __name__ == "__main__":
-    sys.exit(pytest.main(sys.argv))
+    tvm.testing.main()

tests/python/contrib/test_hexagon/topi/test_reduce.py

@@ -20,6 +20,7 @@
 import sys
 
 import tvm
+import tvm.testing
 from tvm import topi
 from tvm import te
 from tvm.contrib.hexagon.session import Session
@@ -163,4 +164,4 @@ def test_reduce_map(
 
 
 if __name__ == "__main__":
-    sys.exit(pytest.main(sys.argv))
+    tvm.testing.main()

tests/python/contrib/test_hexagon/topi/test_softmax.py

@@ -20,6 +20,7 @@
 import sys
 
 import tvm
+import tvm.testing
 from tvm import topi
 from tvm import te
 from tvm.contrib.hexagon.session import Session
@@ -97,4 +98,4 @@ def get_ref_data(shape):
 
 
 if __name__ == "__main__":
-    sys.exit(pytest.main(sys.argv))
+    tvm.testing.main()

tests/python/contrib/test_rpc_server_device.py

@@ -440,4 +440,4 @@ def test_check_auto_schedule_tuning(host, port):  # pylint: disable=too-many-loc
 
 
 if __name__ == "__main__":
-    sys.exit(pytest.main([__file__] + sys.argv[1:]))
+    tvm.testing.main()

tests/python/contrib/test_tensorrt.py

@@ -1402,6 +1402,4 @@ def test_empty_subgraph(run_module):
 
 
 if __name__ == "__main__":
-    import sys
-
-    sys.exit(pytest.main([__file__] + sys.argv[1:]))
+    tvm.testing.main()