From f88aefb2cabf99914dbc9aeb0c23fd7ee5fea9db Mon Sep 17 00:00:00 2001
From: Pairshoe <wtw18@mails.tsinghua.edu.cn>
Date: Wed, 26 Oct 2022 23:57:29 +0800
Subject: [PATCH] Add: pytest for import_onnx

---
 python/infinitensor/test_import_onnx.py | 12 ------------
 python/test/run_test.py                 |  5 +++++
 python/test/test_import_onnx.py         |  9 +++++++++
 3 files changed, 14 insertions(+), 12 deletions(-)
 delete mode 100644 python/infinitensor/test_import_onnx.py
 create mode 100644 python/test/run_test.py
 create mode 100644 python/test/test_import_onnx.py

diff --git a/python/infinitensor/test_import_onnx.py b/python/infinitensor/test_import_onnx.py
deleted file mode 100644
index 581fb206..00000000
--- a/python/infinitensor/test_import_onnx.py
+++ /dev/null
@@ -1,12 +0,0 @@
-from pyinfinitensor import *
-from import_onnx import *
-
-import sys
-
-def main(netPath):
-    runtime = CpuRuntimeObj.getInstance()
-    graphBuilder = GraphBuilderObj(runtime)
-    import_onnx(graphBuilder, netPath)
-
-if __name__ == "__main__":
-    main(sys.argv[1])
diff --git a/python/test/run_test.py b/python/test/run_test.py
new file mode 100644
index 00000000..1521c421
--- /dev/null
+++ b/python/test/run_test.py
@@ -0,0 +1,5 @@
+import pytest
+
+
+if __name__ == "__main__":
+    retcode = pytest.main()
diff --git a/python/test/test_import_onnx.py b/python/test/test_import_onnx.py
new file mode 100644
index 00000000..de4d6914
--- /dev/null
+++ b/python/test/test_import_onnx.py
@@ -0,0 +1,9 @@
+from pyinfinitensor import *
+from infinitensor import import_onnx
+
+
+class Test_ImportOnnx:
+    def test_Netname(self):
+        runtime = CpuRuntimeObj.getInstance()
+        graphBuilder = GraphBuilderObj(runtime)
+        import_onnx(graphBuilder, '/path/to/net')