[SPARK-54950][PYTHON][TEST] Move unittest magic to unittestutils

python/pyspark/testing/__init__.py

@@ -14,33 +14,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
-import sys
-import typing
-import unittest
-
-
-_unittest_main = None
-
-if sys.version_info >= (3, 12) and _unittest_main is None:
-    _unittest_main = unittest.main
-
-    def unittest_main(*args, **kwargs):
-        exit = kwargs.pop("exit", True)
-        kwargs["exit"] = False
-        res = _unittest_main(*args, **kwargs)
-
-        if exit:
-            if not res.result.wasSuccessful():
-                sys.exit(1)
-            elif res.result.testsRun == 0 and len(res.result.skipped) == 0:
-                sys.exit(5)
-            else:
-                sys.exit(0)
-
-        return res
-
-    unittest.main = unittest_main
-
 
 from pyspark.testing.unittestutils import main
 from pyspark.testing.utils import assertDataFrameEqual, assertSchemaEqual

python/pyspark/testing/unittestutils.py

@@ -15,6 +15,7 @@
 # limitations under the License.
 #
 
+import sys
 import unittest
 
 
@@ -25,4 +26,18 @@ def main(module="__main__", output="target/test-reports"):
         testRunner = xmlrunner.XMLTestRunner(output=output, verbosity=2)
     except ImportError:
         testRunner = None
-    unittest.main(module=module, testRunner=testRunner, verbosity=2)
+
+    # Python 3.12+ incorrectly exits with status code 5 if setUpClass fails.
+    # We need to check the result and exit with the correct status code.
+    # This is fixed in latest 3.13+ but is not backported to 3.12.
+    # To be safe, we apply this on all 3.12 - 3.14.
+    if sys.version_info >= (3, 12) and sys.version_info < (3, 15):
+        res = unittest.main(module=module, testRunner=testRunner, verbosity=2, exit=False)
+        if not res.result.wasSuccessful():
+            sys.exit(1)
+        elif res.result.testsRun == 0 and len(res.result.skipped) == 0:
+            sys.exit(5)
+        else:
+            sys.exit(0)
+    else:
+        unittest.main(module=module, testRunner=testRunner, verbosity=2)