Implemented MockerTestCase.addCleanup(). It allows one to

register cleanup functions to be called after the test is
complete.

3 files changed, 47 insertions, 21 deletions

diff --git a/NEWS b/NEWS
index 7cc912e..78ffccf 100644
--- a/NEWS
+++ b/NEWS
@@ -1,6 +1,10 @@
-0.X.X (2007-11-XX)
+0.10 (2007-11-XX)
 ==================
 
+- Implemented MockerTestCase.addCleanup().  It allows one to
+  register cleanup functions to be called after the test is
+  complete.
+
 - MockerTestCase now verifies if the mocker is put in replay
   mode in cases where events were recorded.
 
diff --git a/mocker.py b/mocker.py
index 6189ccd..e546266 100644
--- a/mocker.py
+++ b/mocker.py
@@ -106,15 +106,16 @@ class MockerTestCase(unittest.TestCase):
                             raise RuntimeError("Mocker must be put in replay "
                                                "mode with self.mocker.replay()")
                 except:
-                    self.__cleanup()
                     raise
                 else:
                     if (hasattr(result, "addCallback") and
                         hasattr(result, "addErrback")):
-                        result.addErrback(self.__cleanup)
-                        result.addCallback(self.__cleanup_verify)
+                        def verify(result):
+                            self.mocker.verify()
+                            return result
+                        result.addCallback(verify)
                     else:
-                        self.__cleanup_verify()
+                        self.mocker.verify()
                     return result
             # Copy all attributes from the original method..
             for attr in dir(test_method):
@@ -124,25 +125,37 @@ class MockerTestCase(unittest.TestCase):
                             getattr(test_method, attr))
             setattr(self, methodName, test_method_wrapper)
 
+        # We could overload run() normally, but other well-known testing
+        # frameworks do it as well, and some of them won't call the super,
+        # which might mean that cleanup wouldn't happen.  With that in mind,
+        # we make integration easier by using the following trick.
+        run_method = self.run
+        def run_wrapper(*args, **kwargs):
+            try:
+                return run_method(*args, **kwargs)
+            finally:
+                self.__cleanup()
+        self.run = run_wrapper
+
         self.mocker = Mocker()
 
+        self.__cleanup_funcs = []
         self.__cleanup_paths = []
 
         super(MockerTestCase, self).__init__(methodName)
 
-    def __cleanup_verify(self, result=None):
-        self.__cleanup()
-        self.mocker.verify()
-        return result
-
-    def __cleanup(self, result=None):
+    def __cleanup(self):
         for path in self.__cleanup_paths:
             if os.path.isfile(path):
                 os.unlink(path)
             elif os.path.isdir(path):
                 shutil.rmtree(path)
         self.mocker.restore()
-        return result
+        for func, args, kwargs in self.__cleanup_funcs:
+            func(*args, **kwargs)
+
+    def addCleanup(self, func, *args, **kwargs):
+        self.__cleanup_funcs.append((func, args, kwargs))
 
     def makeFile(self, content=None, suffix="", prefix="tmp", basename=None,
                  dirname=None):
diff --git a/test.py b/test.py
index 92fb8a3..3d27d54 100755
--- a/test.py
+++ b/test.py
@@ -373,7 +373,7 @@ class MockerTestCaseTest(TestCase):
         result = unittest.TestResult()
         MyTest("test_method").run(result)
 
-        self.assertEquals(calls, ["restore", "verify"])
+        self.assertEquals(calls, ["verify", "restore"])
         self.assertTrue(result.wasSuccessful())
 
         del calls[:]
@@ -400,6 +400,22 @@ class MockerTestCaseTest(TestCase):
         self.assertEquals(len(result.failures), 1)
         self.assertTrue("mock.x" in result.failures[0][1])
 
+    def test_add_cleanup(self):
+        stash = []
+        def func(a, b):
+            stash.append((a, b))
+
+        class MyTest(MockerTestCase):
+            def tearDown(self):
+                self.addCleanup(func, 3, b=4)
+            def test_method(self):
+                self.addCleanup(func, 1, b=2)
+                stash.append(stash[:])
+
+        MyTest("test_method").run()
+
+        self.assertEquals(stash, [[], (1, 2), (3, 4)])
+
     def test_twisted_trial_deferred_support(self):
         calls = []
         callbacks = []
@@ -431,14 +447,7 @@ class MockerTestCaseTest(TestCase):
         self.assertEquals(calls, [])
         self.assertEquals(len(callbacks), 1)
         self.assertEquals(callbacks[-1]("foo"), "foo")
-        self.assertEquals(calls, ["restore", "verify"])
-
-        test.mocker.replay()
-        del calls[:]
-
-        self.assertEquals(len(errbacks), 1)
-        self.assertEquals(errbacks[-1]("foo"), "foo")
-        self.assertEquals(calls, ["restore"])
+        self.assertEquals(calls, ["verify"])
 
 
     def test_fail_unless_is_raises_on_mismatch(self):