From 2d61dbdb93fa941c54b3990c2ee046074d2fce07 Mon Sep 17 00:00:00 2001
From: Steve Li <steve.li@planetinnovation.com.au>
Date: Tue, 5 Apr 2022 10:05:30 +1000
Subject: [PATCH] unittest: Add setUpClass and tearDownClass handling.

Supports setUp and tearDown functionality at Class level.
---
 python-stdlib/unittest/test_unittest.py | 19 +++++++++++++++++++
 python-stdlib/unittest/unittest.py      |  6 ++++++
 2 files changed, 25 insertions(+)

diff --git a/python-stdlib/unittest/test_unittest.py b/python-stdlib/unittest/test_unittest.py
index 1b02056d..8e108995 100644
--- a/python-stdlib/unittest/test_unittest.py
+++ b/python-stdlib/unittest/test_unittest.py
@@ -157,5 +157,24 @@ class TestUnittestAssertions(unittest.TestCase):
                 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
index 0fcb3293..cfb274a6 100644
--- 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_suite(c, test_result, suite_name=""):
             except AttributeError:
                 pass
 
+    set_up_class()
+
     if hasattr(o, "runTest"):
         name = str(o)
         run_one(o.runTest)
@@ -426,6 +430,8 @@ def run_suite(c, test_result, suite_name=""):
         name = o.__name__
         run_one(o)
 
+    tear_down_class()
+
     return exceptions