[BE] move more unittest.main() to run_tests() (#50923)

Rong Rong (AI Infra) · facebook-github-bot · commit 806010b75e2b · 2021-01-25T07:23:09.000-08:00
Summary: Relate to #50483. Everything except ONNX, detectron and release notes tests are moved to use common_utils.run_tests() to ensure CI reports XML correctly. Pull Request resolved: #50923 Reviewed By: samestep Differential Revision: D26027621 Pulled By: walterddr fbshipit-source-id: b04c03f10d1fe96181b720c4c3868e86e4c6281a
diff --git a/test/custom_backend/test_custom_backend.py b/test/custom_backend/test_custom_backend.py
@@ -1,12 +1,12 @@
 import os
 import tempfile
 import torch
-import unittest
 
 from backend import Model, to_custom_backend, get_custom_backend_library_path
+from torch.testing._internal.common_utils import TestCase, run_tests
 
 
-class TestCustomBackend(unittest.TestCase):
+class TestCustomBackend(TestCase):
     def setUp(self):
         # Load the library containing the custom backend.
         self.library_path = get_custom_backend_library_path()
@@ -51,4 +51,4 @@ def test_save_load(self):
 
 
 if __name__ == "__main__":
-    unittest.main()
+    run_tests()
diff --git a/test/custom_operator/test_custom_classes.py b/test/custom_operator/test_custom_classes.py
@@ -5,6 +5,9 @@
 import glob
 import os
 
+from torch.testing._internal.common_utils import TestCase, run_tests
+
+
 def get_custom_class_library_path():
     library_filename = glob.glob("build/*custom_class*")
     assert (len(library_filename) == 1)
@@ -18,7 +21,7 @@ def test_equality(f, cmp_key):
     obj2 = jit.script(f)()
     return (cmp_key(obj1), cmp_key(obj2))
 
-class TestCustomOperators(unittest.TestCase):
+class TestCustomOperators(TestCase):
     def setUp(self):
         ops.load_library(get_custom_class_library_path())
 
@@ -77,4 +80,4 @@ def f():
 
 
 if __name__ == "__main__":
-    unittest.main()
+    run_tests()
diff --git a/test/custom_operator/test_custom_ops.py b/test/custom_operator/test_custom_ops.py
@@ -1,14 +1,14 @@
 import os.path
 import tempfile
-import unittest
 
 import torch
 from torch import ops
 
 from model import Model, get_custom_op_library_path
+from torch.testing._internal.common_utils import TestCase, run_tests
 
 
-class TestCustomOperators(unittest.TestCase):
+class TestCustomOperators(TestCase):
     def setUp(self):
         self.library_path = get_custom_op_library_path()
         ops.load_library(self.library_path)
@@ -90,4 +90,4 @@ def test_saving_and_loading_script_module_with_custom_op(self):
 
 
 if __name__ == "__main__":
-    unittest.main()
+    run_tests()
diff --git a/test/mobile/test_lite_script_module.py b/test/mobile/test_lite_script_module.py
@@ -1,4 +1,3 @@
-import unittest
 import torch
 import torch.utils.bundled_inputs
 from torch.utils.mobile_optimizer import *
@@ -7,8 +6,9 @@
 from collections import namedtuple
 
 from torch.jit.mobile import _load_for_lite_interpreter
+from torch.testing._internal.common_utils import TestCase, run_tests
 
-class TestLiteScriptModule(unittest.TestCase):
+class TestLiteScriptModule(TestCase):
 
     def test_load_mobile_module(self):
         class MyTestModule(torch.nn.Module):
@@ -287,4 +287,4 @@ def forward(self):
             script_module._save_to_buffer_for_lite_interpreter()
 
 if __name__ == '__main__':
-    unittest.main()
+    run_tests()
diff --git a/test/test_tensorexpr_pybind.py b/test/test_tensorexpr_pybind.py
@@ -1,6 +1,6 @@
 import torch
-import unittest
 
+from torch.testing._internal.common_utils import run_tests
 from torch.testing._internal.jit_utils import JitTestCase
 
 class kernel_arena_scope(object):
@@ -37,4 +37,4 @@ def compute(i):
             torch.testing.assert_allclose(tA + tB, tC)
 
 if __name__ == '__main__':
-    unittest.main()
+    run_tests()
