Add addCleanup and doCleanups to unittest.TestCase.

Closes issue 5679.

Michael Foord

Doc/library/unittest.rst

@@ -978,6 +978,36 @@ Test cases
       .. versionadded:: 2.7
 
 
+   .. method:: addCleanup(function[, *args[, **kwargs]])
+
+      Add a function to be called after :meth:`tearDown` to cleanup resources
+      used during the test. Functions will be called in reverse order to the
+      order they are added (LIFO). They are called with any arguments and
+      keyword arguments passed into :meth:`addCleanup` when they are
+      added.
+
+      If :meth:`setUp` fails, meaning that :meth:`tearDown` is not called,
+      then any cleanup functions added will still be called.
+
+      .. versionadded:: 2.7
+
+
+   .. method:: doCleanups()
+
+      This method is called uncoditionally after :meth:`tearDown`, or
+      after :meth:`setUp` if :meth:`setUp` raises an exception.
+
+      It is responsible for calling all the cleanup functions added by
+      :meth:`addCleanup`. If you need cleanup functions to be called
+      *prior* to :meth:`tearDown` then you can call :meth:`doCleanups`
+      yourself.
+
+      :meth:`doCleanups` pops methods off the stack of cleanup
+      functions one at a time, so it can be called at any time.
+
+      .. versionadded:: 2.7
+
+
 .. class:: FunctionTestCase(testFunc[, setUp[, tearDown[, description]]])
 
    This class implements the portion of the :class:`TestCase` interface which

Doc/whatsnew/2.7.rst

@@ -452,6 +452,13 @@ changes, or look through the Subversion logs for all the details.
 
   (Implemented by Antoine Pitrou; :issue:`4444`.)
 
+  The methods :meth:`addCleanup` and :meth:`doCleanups` were added.
+  :meth:`addCleanup` allows you to add cleanup functions that
+  will be called unconditionally (after :meth:`setUp` if
+  :meth:`setUp` fails, otherwise after :meth:`tearDown`). This allows
+  for much simpler resource allocation and deallocation during tests.
+  :issue:`5679`
+
   A number of new methods were added that provide more specialized
   tests.  Many of these methods were written by Google engineers
   for use in their test suites; Gregory P. Smith, Michael Foord, and

Lib/test/test_unittest.py

@@ -3041,6 +3041,111 @@ class TestLongMessage(TestCase):
                              "^unexpectedly identical: None : oops$"])
 
 
+class TestCleanUp(TestCase):
+
+    def testCleanUp(self):
+        class TestableTest(TestCase):
+            def testNothing(self):
+                pass
+
+        test = TestableTest('testNothing')
+        self.assertEqual(test._cleanups, [])
+
+        cleanups = []
+
+        def cleanup1(*args, **kwargs):
+            cleanups.append((1, args, kwargs))
+
+        def cleanup2(*args, **kwargs):
+            cleanups.append((2, args, kwargs))
+
+        test.addCleanup(cleanup1, 1, 2, 3, four='hello', five='goodbye')
+        test.addCleanup(cleanup2)
+
+        self.assertEqual(test._cleanups,
+                         [(cleanup1, (1, 2, 3), dict(four='hello', five='goodbye')),
+                          (cleanup2, (), {})])
+
+        result = test.doCleanups()
+        self.assertTrue(result)
+
+        self.assertEqual(cleanups, [(2, (), {}), (1, (1, 2, 3), dict(four='hello', five='goodbye'))])
+
+    def testCleanUpWithErrors(self):
+        class TestableTest(TestCase):
+            def testNothing(self):
+                pass
+
+        class MockResult(object):
+            errors = []
+            def addError(self, test, exc_info):
+                self.errors.append((test, exc_info))
+
+        result = MockResult()
+        test = TestableTest('testNothing')
+        test._result = result
+
+        exc1 = Exception('foo')
+        exc2 = Exception('bar')
+        def cleanup1():
+            raise exc1
+
+        def cleanup2():
+            raise exc2
+
+        test.addCleanup(cleanup1)
+        test.addCleanup(cleanup2)
+
+        self.assertFalse(test.doCleanups())
+
+        (test1, (Type1, instance1, _)), (test2, (Type2, instance2, _)) = reversed(MockResult.errors)
+        self.assertEqual((test1, Type1, instance1), (test, Exception, exc1))
+        self.assertEqual((test2, Type2, instance2), (test, Exception, exc2))
+
+    def testCleanupInRun(self):
+        blowUp = False
+        ordering = []
+
+        class TestableTest(TestCase):
+            def setUp(self):
+                ordering.append('setUp')
+                if blowUp:
+                    raise Exception('foo')
+
+            def testNothing(self):
+                ordering.append('test')
+
+            def tearDown(self):
+                ordering.append('tearDown')
+
+        test = TestableTest('testNothing')
+
+        def cleanup1():
+            ordering.append('cleanup1')
+        def cleanup2():
+            ordering.append('cleanup2')
+        test.addCleanup(cleanup1)
+        test.addCleanup(cleanup2)
+
+        def success(some_test):
+            self.assertEqual(some_test, test)
+            ordering.append('success')
+
+        result = unittest.TestResult()
+        result.addSuccess = success
+
+        test.run(result)
+        self.assertEqual(ordering, ['setUp', 'test', 'tearDown',
+                                    'cleanup2', 'cleanup1', 'success'])
+
+        blowUp = True
+        ordering = []
+        test = TestableTest('testNothing')
+        test.addCleanup(cleanup1)
+        test.run(result)
+        self.assertEqual(ordering, ['setUp', 'cleanup1'])
+
+
 class Test_TestProgram(TestCase):
 
     # Horrible white box test
@@ -3110,7 +3215,6 @@ class Test_TestProgram(TestCase):
             testLoader=self.FooBarLoader())
 
 
-
 ######################################################################
 ## Main
 ######################################################################
@@ -3119,7 +3223,7 @@ def test_main():
     test_support.run_unittest(Test_TestCase, Test_TestLoader,
         Test_TestSuite, Test_TestResult, Test_FunctionTestCase,
         Test_TestSkipping, Test_Assertions, TestLongMessage,
-        Test_TestProgram)
+        Test_TestProgram, TestCleanUp)
 
 if __name__ == "__main__":
     test_main()

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,19 +439,18 @@ 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:
+            success = False
             try:
                 self.setUp()
             except SkipTest as e:
                 result.addSkip(self, str(e))
-                return
             except Exception:
                 result.addError(self, sys.exc_info())
-                return
-
-            success = False
+            else:
                 try:
                     testMethod()
                 except self.failureException:
@@ -462,11 +471,28 @@ class TestCase(object):
                 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)