Add addCleanup and doCleanups to unittest.TestCase.

Closes issue 5679.

Michael Foord

1 files changed, 52 insertions, 22 deletions

diff --git a/Lib/unittest.py b/Lib/unittest.py
index 43a2ced..d5c2927 100644
--- a/Lib/unittest.py
+++ b/Lib/unittest.py
@@ -340,12 +340,14 @@ class TestCase(object):
            not have a method with the specified name.
         """
         self._testMethodName = methodName
+        self._result = None
         try:
             testMethod = getattr(self, methodName)
         except AttributeError:
             raise ValueError("no such test method in %s: %s" % \
                   (self.__class__, methodName))
         self._testMethodDoc = testMethod.__doc__
+        self._cleanups = []
 
         # Map types to custom assertEqual functions that will compare
         # instances of said type in more detail to generate a more useful
@@ -372,6 +374,14 @@ class TestCase(object):
         """
         self._type_equality_funcs[typeobj] = _AssertWrapper(function)
 
+    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 setUp(self):
         "Hook method for setting up the test fixture before exercising it."
         pass
@@ -429,44 +439,60 @@ class TestCase(object):
     def run(self, result=None):
         if result is None:
             result = self.defaultTestResult()
+        self._result = result
         result.startTest(self)
         testMethod = getattr(self, self._testMethodName)
         try:
-            try:
-                self.setUp()
-            except SkipTest as e:
-                result.addSkip(self, str(e))
-                return
-            except Exception:
-                result.addError(self, sys.exc_info())
-                return
-
             success = False
             try:
-                testMethod()
-            except self.failureException:
-                result.addFailure(self, sys.exc_info())
-            except _ExpectedFailure as e:
-                result.addExpectedFailure(self, e.exc_info)
-            except _UnexpectedSuccess:
-                result.addUnexpectedSuccess(self)
+                self.setUp()
             except SkipTest as e:
                 result.addSkip(self, str(e))
             except Exception:
                 result.addError(self, sys.exc_info())
             else:
-                success = True
+                try:
+                    testMethod()
+                except self.failureException:
+                    result.addFailure(self, sys.exc_info())
+                except _ExpectedFailure as e:
+                    result.addExpectedFailure(self, e.exc_info)
+                except _UnexpectedSuccess:
+                    result.addUnexpectedSuccess(self)
+                except SkipTest as e:
+                    result.addSkip(self, str(e))
+                except Exception:
+                    result.addError(self, sys.exc_info())
+                else:
+                    success = True
 
-            try:
-                self.tearDown()
-            except Exception:
-                result.addError(self, sys.exc_info())
-                success = False
+                try:
+                    self.tearDown()
+                except Exception:
+                    result.addError(self, sys.exc_info())
+                    success = False
+
+            cleanUpSuccess = self.doCleanups()
+            success = success and cleanUpSuccess
             if success:
                 result.addSuccess(self)
         finally:
             result.stopTest(self)
 
+    def doCleanups(self):
+        """Execute all cleanup functions. Normally called for you after
+        tearDown."""
+        result = self._result
+        ok = True
+        while self._cleanups:
+            function, args, kwargs = self._cleanups.pop(-1)
+            try:
+                function(*args, **kwargs)
+            except Exception:
+                ok = False
+                result.addError(self, sys.exc_info())
+        return ok
+
     def __call__(self, *args, **kwds):
         return self.run(*args, **kwds)
 
@@ -1538,5 +1564,9 @@ Examples:
 main = TestProgram
 
 
+##############################################################################
+# Executing this module from the command line
+##############################################################################
+
 if __name__ == "__main__":
     main(module=None)