[tvm] branch main updated: make test_runtime_rpc use pytest.main() (#7482)

[tq...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git


The following commit(s) were added to refs/heads/main by this push:
     new 5a0573b  make test_runtime_rpc use pytest.main() (#7482)
5a0573b is described below

commit 5a0573b0fcc1bea2045f43a2fce6d85d05d5102a
Author: Andrew Reusch <ar...@octoml.ai>
AuthorDate: Sat Feb 20 06:21:17 2021 -0800

    make test_runtime_rpc use pytest.main() (#7482)
---
 tests/python/unittest/test_runtime_rpc.py | 28 +++++++++-------------------
 1 file changed, 9 insertions(+), 19 deletions(-)

diff --git a/tests/python/unittest/test_runtime_rpc.py b/tests/python/unittest/test_runtime_rpc.py
index e975a16..11c1098 100644
--- a/tests/python/unittest/test_runtime_rpc.py
+++ b/tests/python/unittest/test_runtime_rpc.py
@@ -17,11 +17,12 @@
 import tvm
 from tvm import te
 import tvm.testing
+import logging
+import multiprocessing
 import os
 import stat
-import logging
+import sys
 import time
-import multiprocessing
 
 import pytest
 import numpy as np
@@ -29,6 +30,12 @@ from tvm import rpc
 from tvm.contrib import utils, cc
 from tvm.rpc.tracker import Tracker
 
+
+if __name__ == "__main__":
+    # NOTE: must live here to avoid registering PackedFunc with libtvm.so twice.
+    sys.exit(pytest.main([__file__] + sys.argv[1:]))
+
+
 # tkonolige: The issue as I understand it is this: multiprocessing's spawn
 # method launches a new process and then imports the relevant modules. This
 # means that all registered functions must exist at the top level scope. In
@@ -526,20 +533,3 @@ def test_rpc_tracker_request():
     proc2.join()
     server.terminate()
     tracker.terminate()
-
-
-if __name__ == "__main__":
-    logging.basicConfig(level=logging.INFO)
-    test_rpc_echo()
-    test_rpc_session_constructor_args()
-    test_rpc_return_ndarray()
-    test_rpc_return_func()
-    test_bigendian_rpc()
-    test_rpc_remote_module()
-    test_rpc_file_exchange()
-    test_rpc_array()
-    test_rpc_simple()
-    test_local_func()
-    test_rpc_tracker_register()
-    test_rpc_tracker_request()
-    test_rpc_large_array()