gh-134079: Add addCleanup, enterContext and doCleanups to unittest.subTest and tests

Lib/test/test_unittest/test_runner.py

@@ -1452,6 +1452,206 @@ def test_1(self):
         run(Foo('test_1'), expect_durations=False)
 
 
+class SubTestMockContextManager:
+    def __init__(self, events_list, name="cm"):
+        self.events = events_list
+        self.name = name
+        self.entered_value = f"{name}_entered_value"
+
+    def __enter__(self):
+        self.events.append(f"{self.name}_enter")
+        return self.entered_value
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.events.append((f"{self.name}_exit", exc_type.__name__ if exc_type else None))
+        return False
+
+
+class SubTestFailingEnterCM:
+    def __init__(self, events_list, name="cm_enter_fail"):
+        self.events = events_list
+        self.name = name
+
+    def __enter__(self):
+        self.events.append(f"{self.name}_enter_attempt")
+        raise RuntimeError("Simulated __enter__ failure")
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.events.append((f"{self.name}_exit_UNEXPECTED", exc_type.__name__ if exc_type else None))
+        return False
+
+
+class TestSubTestCleanups(unittest.TestCase):
+    def setUp(self):
+        self.events = []
+        TestSubTestCleanups._active_events_list = self.events
+        TestSubTestCleanups._active_instance = self
+
+    def tearDown(self):
+        if hasattr(TestSubTestCleanups, '_active_events_list'):
+            del TestSubTestCleanups._active_events_list
+        if hasattr(TestSubTestCleanups, '_active_instance'):
+            del TestSubTestCleanups._active_instance
+
+    @classmethod
+    def _get_events(cls):
+        return cls._active_events_list
+
+    def _record_event(self, event_name, *args, **kwargs):
+        TestSubTestCleanups._get_events().append(event_name)
+
+    def test_addCleanup_operation_and_LIFO_order(self):
+        class MyTests(unittest.TestCase):
+            def runTest(inner_self):
+                events = TestSubTestCleanups._get_events()
+                recorder = TestSubTestCleanups._active_instance._record_event
+                with inner_self.subTest() as sub:
+                    events.append("subtest_body_start")
+                    sub.addCleanup(recorder, "cleanup_2_args", "arg")
+                    sub.addCleanup(recorder, "cleanup_1")
+                    events.append("subtest_body_end")
+
+        MyTests().run()
+
+        expected_events = [
+            "subtest_body_start",
+            "subtest_body_end",
+            "cleanup_1",
+            "cleanup_2_args",
+        ]
+        self.assertEqual(self.events, expected_events)
+
+    def test_addCleanup_runs_on_subtest_failure(self):
+        class MyTests(unittest.TestCase):
+            def runTest(inner_self):
+                recorder = TestSubTestCleanups._active_instance._record_event
+                with inner_self.subTest() as sub:
+                    sub.addCleanup(recorder, "cleanup_on_fail")
+                    inner_self.fail("Intentional subtest failure")
+
+        test = MyTests()
+        result = unittest.TestResult()
+        test.run(result)
+
+        self.assertFalse(result.wasSuccessful())
+        self.assertEqual(len(result.failures), 1)
+        self.assertEqual(self.events, ["cleanup_on_fail"])
+
+    def test_enterContext_success(self):
+        class MyTests(unittest.TestCase):
+            def runTest(inner_self):
+                events = TestSubTestCleanups._get_events()
+                with inner_self.subTest() as sub:
+                    cm = SubTestMockContextManager(events, name="cm_ok")
+                    value = sub.enterContext(cm)
+                    inner_self.assertEqual(value, "cm_ok_entered_value")
+                    events.append(f"context_body_value_{value}")
+
+        test = MyTests()
+        result = unittest.TestResult()
+        test.run(result)
+
+        self.assertTrue(result.wasSuccessful())
+        expected_events = [
+            "cm_ok_enter",
+            "context_body_value_cm_ok_entered_value",
+            ("cm_ok_exit", None),
+        ]
+        self.assertEqual(self.events, expected_events)
+
+    def test_enterContext_body_fails_after_enter(self):
+        class MyTests(unittest.TestCase):
+            def runTest(inner_self):
+                events = TestSubTestCleanups._get_events()
+                with inner_self.subTest(description="inner_subtest_expected_to_fail") as sub:
+                    cm = SubTestMockContextManager(events, name="cm_body_fail")
+                    sub.enterContext(cm)
+                    events.append("before_body_raise")
+                    raise ValueError("Deliberate body failure after enterContext")
+
+        test_instance = MyTests()
+        internal_result = unittest.TestResult()
+        test_instance.run(internal_result)
+
+        self.assertFalse(internal_result.wasSuccessful())
+        self.assertEqual(len(internal_result.failures), 0)
+        self.assertEqual(len(internal_result.errors), 1)
+
+        errored_test_obj, tb = internal_result.errors[0]
+        self.assertIsInstance(errored_test_obj, unittest.case._SubTest)
+        self.assertIn("ValueError: Deliberate body failure after enterContext", tb)
+        self.assertEqual(errored_test_obj.params.get("description"), "inner_subtest_expected_to_fail")
+
+        expected_events = [
+            "cm_body_fail_enter",
+            "before_body_raise",
+            ("cm_body_fail_exit", None),
+        ]
+        self.assertEqual(self.events, expected_events)
+
+    def test_enterContext_manager_enter_fails(self):
+        class MyTests(unittest.TestCase):
+            def runTest(inner_self):
+                events = TestSubTestCleanups._get_events()
+                with inner_self.subTest() as sub:
+                    cm_enter_fail = SubTestFailingEnterCM(events)
+                    with inner_self.assertRaisesRegex(RuntimeError, "Simulated __enter__ failure"):
+                        sub.enterContext(cm_enter_fail)
+                    inner_self.assertEqual(len(sub._cleanups), 0)
+
+        test = MyTests()
+        result = unittest.TestResult()
+        test.run(result)
+
+        self.assertTrue(result.wasSuccessful())
+        self.assertEqual(self.events, ["cm_enter_fail_enter_attempt"])
+
+    def test_failing_cleanup_function_in_subtest(self):
+        def failing_cleanup_func():
+            TestSubTestCleanups._get_events().append("failing_cleanup_called")
+            raise ZeroDivisionError("cleanup boom")
+
+        class MyTests(unittest.TestCase):
+            def runTest(inner_self):
+                recorder = TestSubTestCleanups._active_instance._record_event
+                with inner_self.subTest() as sub:
+                    sub.addCleanup(recorder, "cleanup_after_fail_attempt")
+                    sub.addCleanup(failing_cleanup_func)
+                    sub.addCleanup(recorder, "cleanup_before_fail_attempt")
+
+        test = MyTests()
+        result = unittest.TestResult()
+        test.run(result)
+
+        self.assertFalse(result.wasSuccessful())
+        self.assertEqual(len(result.errors), 1)
+
+        expected_events_order = [
+            "cleanup_before_fail_attempt",
+            "failing_cleanup_called",
+            "cleanup_after_fail_attempt",
+        ]
+        self.assertEqual(self.events, expected_events_order)
+
+    def test_multiple_subtests_independent_cleanups(self):
+        class MyTests(unittest.TestCase):
+            def runTest(inner_self):
+                events = TestSubTestCleanups._get_events()
+                recorder = TestSubTestCleanups._active_instance._record_event
+                for i in range(2):
+                    with inner_self.subTest(i=i) as sub:
+                        cleanup_msg = f"cleanup_sub_{i}"
+                        sub.addCleanup(recorder, cleanup_msg)
+                        events.append(f"subtest_body_{i}")
+
+        MyTests().run()
+
+        expected_events = [
+            "subtest_body_0", "cleanup_sub_0",
+            "subtest_body_1", "cleanup_sub_1",
+        ]
+        self.assertEqual(self.events, expected_events)
+
 
 if __name__ == "__main__":
     unittest.main()

Lib/unittest/case.py

@@ -559,10 +559,12 @@ def subTest(self, msg=_subtest_msg_sentinel, **params):
             params_map = _OrderedChainMap(params)
         else:
             params_map = parent.params.new_child(params)
-        self._subtest = _SubTest(self, msg, params_map)
+        subtest = _SubTest(self, msg, params_map)
+        self._subtest = subtest
+        cleanup_helper = _SubTestCleanupHelper(subtest, self._outcome)
         try:
             with self._outcome.testPartExecutor(self._subtest, subTest=True):
-                yield
+                yield cleanup_helper
             if not self._outcome.success:
                 result = self._outcome.result
                 if result is not None and result.failfast:
@@ -572,6 +574,7 @@ def subTest(self, msg=_subtest_msg_sentinel, **params):
                 # stop now and register the expected failure.
                 raise _ShouldStop
         finally:
+            cleanup_helper.doCleanups()
             self._subtest = parent
 
     def _addExpectedFailure(self, result, exc_info):
@@ -1626,3 +1629,47 @@ def shortDescription(self):
 
     def __str__(self):
         return "{} {}".format(self.test_case, self._subDescription())
+
+
+class _SubTestCleanupHelper:
+    """
+    Helper class to manage cleanups and context managers inside subTest blocks,
+    without exposing full TestCase functionality.
+    """
+
+    def __init__(self, subtest, outcome=None):
+        self._cleanups = []
+        self._subtest = subtest
+        self._outcome = outcome
+
+    @staticmethod
+    def _callCleanup(function, /, *args, **kwargs):
+        function(*args, **kwargs)
+
+    def addCleanup(self, function, /, *args, **kwargs):
+        """Add a function, with arguments, to be called when the test is
+        completed. Functions added are called on a LIFO basis and are
+        called after tearDown on test failure or success.
+
+        Cleanup items are called even if setUp fails (unlike tearDown)."""
+        self._cleanups.append((function, args, kwargs))
+
+    def enterContext(self, cm):
+        """Enters the supplied context manager.
+
+        If successful, also adds its __exit__ method as a cleanup
+        function and returns the result of the __enter__ method.
+        """
+        return _enter_context(cm, self.addCleanup)
+
+    def doCleanups(self):
+        """Execute all cleanup functions registered for this subtest."""
+        outcome = self._outcome or _Outcome()
+        while self._cleanups:
+            function, args, kwargs = self._cleanups.pop()
+            if hasattr(outcome, 'testPartExecutor'):
+                with outcome.testPartExecutor(self._subtest, subTest=True):
+                    self._callCleanup(function, *args, **kwargs)
+            else:
+                self._callCleanup(function, *args, **kwargs)
+        return outcome.success

Misc/NEWS.d/next/Library/2025-05-22-18-45-14.gh-issue-134079.n89xN5.rst

@@ -0,0 +1 @@
+unittest.subTest now supports addCleanup(), enterContext(), and doCleanups(), enabling resource management and cleanup inside subTest blocks.