unittest: Add setUpClass and tearDownClass handling.

Steve Li · pi-anl · commit 2d61dbdb93fa · 2022-05-06T10:47:37.000+10:00
Supports setUp and tearDown functionality at Class level.
diff --git a/python-stdlib/unittest/test_unittest.py b/python-stdlib/unittest/test_unittest.py
@@ -157,5 +157,24 @@ def test_subtest_even(self):
                 self.assertEqual(i % 2, 0)
 
 
+class TestUnittestSetup(unittest.TestCase):
+    class_setup_var = 0
+
+    def setUpClass(self):
+        TestUnittestSetup.class_setup_var += 1
+
+    def tearDownClass(self):
+        # Not sure how to actually test this, but we can check (in the test case below)
+        # that it hasn't been run already at least.
+        TestUnittestSetup.class_setup_var = -1
+
+    def testSetUpTearDownClass_1(self):
+        assert TestUnittestSetup.class_setup_var == 1, TestUnittestSetup.class_setup_var
+
+    def testSetUpTearDownClass_2(self):
+        # Test this twice, as if setUpClass() gets run like setUp() it would be run twice
+        assert TestUnittestSetup.class_setup_var == 1, TestUnittestSetup.class_setup_var
+
+
 if __name__ == "__main__":
     unittest.main()
diff --git a/python-stdlib/unittest/unittest.py b/python-stdlib/unittest/unittest.py
@@ -366,6 +366,8 @@ def run_suite(c, test_result, suite_name=""):
         o = c()
     else:
         o = c
+    set_up_class = getattr(o, "setUpClass", lambda: None)
+    tear_down_class = getattr(o, "tearDownClass", lambda: None)
     set_up = getattr(o, "setUp", lambda: None)
     tear_down = getattr(o, "tearDown", lambda: None)
     exceptions = []
@@ -410,6 +412,8 @@ def run_one(test_function):
             except AttributeError:
                 pass
 
+    set_up_class()
+
     if hasattr(o, "runTest"):
         name = str(o)
         run_one(o.runTest)
@@ -426,6 +430,8 @@ def run_one(test_function):
         name = o.__name__
         run_one(o)
 
+    tear_down_class()
+
     return exceptions
 
 
