Implemented support for Deferred results as understood by

Twisted Trial's TestCase, so that coexistence by multiple
inheritance is possible and trivial.

3 files changed, 65 insertions, 6 deletions

diff --git a/NEWS b/NEWS
index 4aeccfe..85ed71e 100644
--- a/NEWS
+++ b/NEWS
@@ -1,3 +1,11 @@
+0.X.X (2007-11-XX)
+==================
+
+- Implemented support for Deferred results as understood by
+  Twisted Trial's TestCase, so that coexistence by multiple
+  inheritance is possible and trivial.
+
+
 0.9.3 (2007-11-24)
 ==================
 
diff --git a/mocker.py b/mocker.py
index 8929b91..79f7191 100644
--- a/mocker.py
+++ b/mocker.py
@@ -98,15 +98,18 @@ class MockerTestCase(unittest.TestCase):
         if test_method is not None:
             def test_method_wrapper():
                 try:
-                    test_method()
+                    result = test_method()
                 except:
                     self.__cleanup()
-                    self.mocker.restore()
                     raise
                 else:
-                    self.__cleanup()
-                    self.mocker.restore()
-                    self.mocker.verify()
+                    if (hasattr(result, "addCallback") and
+                        hasattr(result, "addErrback")):
+                        result.addErrback(self.__cleanup)
+                        result.addCallback(self.__cleanup_verify)
+                    else:
+                        self.__cleanup_verify()
+                    return result
             # Copy all attributes from the original method..
             for attr in dir(test_method):
                 # .. unless they're present in our wrapper already.
@@ -121,12 +124,19 @@ class MockerTestCase(unittest.TestCase):
 
         super(MockerTestCase, self).__init__(methodName)
 
-    def __cleanup(self):
+    def __cleanup_verify(self, result=None):
+        self.__cleanup()
+        self.mocker.verify()
+        return result
+
+    def __cleanup(self, result=None):
         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
 
     def makeFile(self, content=None, suffix="", prefix="tmp", basename=None,
                  dirname=None):
diff --git a/test.py b/test.py
index 8e3eb36..a427c64 100755
--- a/test.py
+++ b/test.py
@@ -371,6 +371,47 @@ class MockerTestCaseTest(TestCase):
         self.assertEquals(len(result.failures), 1)
         self.assertTrue("mock.x" in result.failures[0][1])
 
+    def test_twisted_trial_deferred_support(self):
+        calls = []
+        callbacks = []
+        errbacks = []
+        deferreds = []
+        class Deferred(object):
+            def addCallback(self, callback):
+                callbacks.append(callback)
+            def addErrback(self, errback):
+                errbacks.append(errback)
+        class MyEvent(Event):
+            def verify(self):
+                calls.append("verify")
+            def restore(self):
+                calls.append("restore")
+        class MyTest(MockerTestCase):
+            def test_method(self):
+                self.mocker.add_event(MyEvent())
+                self.mocker.replay()
+                deferred = Deferred()
+                deferreds.append(deferred)
+                return deferred
+
+        result = unittest.TestResult()
+        test = MyTest("test_method")
+        deferred = test.test_method()
+
+        self.assertEquals(deferred, deferreds[-1])
+        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"])
+
+
     def test_fail_unless_is_raises_on_mismatch(self):
         try:
             self.test.failUnlessIs([], [])