Add addCleanup and doCleanups to unittest.TestCase.

Closes issue 5679.

Michael Foord

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()