Propagate exception info in _patch.__exit__

3 files changed, 39 insertions, 4 deletions

diff --git a/docs/changelog.txt b/docs/changelog.txt
index 6357fdd..5d81145 100644
--- a/docs/changelog.txt
+++ b/docs/changelog.txt
@@ -18,6 +18,8 @@ The standard library version!
 * Addition of `PropertyMock`, for mocking properties
 * `MagicMocks` made unorderable by default (in Python 3). The comparison
   methods (other than equality and inequality) now return `NotImplemented`
+* Propagate traceback info to support subclassing of `_patch` by other
+  libraries
 
 
 2012/02/13 Version 0.8.0
diff --git a/mock.py b/mock.py
index 9b9a513..a41c33d 100644
--- a/mock.py
+++ b/mock.py
@@ -1153,6 +1153,7 @@ class _patch(object):
 
             # can't use try...except...finally because of Python 2.4
             # compatibility
+            exc_info = tuple()
             try:
                 try:
                     for patching in patched.patchings:
@@ -1171,11 +1172,13 @@ class _patch(object):
                         # the patcher may have been started, but an exception
                         # raised whilst entering one of its additional_patchers
                         entered_patchers.append(patching)
+                    # Pass the exception to __exit__
+                    exc_info = sys.exc_info()
                     # re-raise the exception
                     raise
             finally:
                 for patching in reversed(entered_patchers):
-                    patching.__exit__()
+                    patching.__exit__(*exc_info)
 
         patched.patchings = [self]
         if hasattr(func, 'func_code'):
@@ -1317,7 +1320,7 @@ class _patch(object):
         del self.target
         for patcher in reversed(self.additional_patchers):
             if _is_started(patcher):
-                patcher.__exit__()
+                patcher.__exit__(*_)
 
     start = __enter__
     stop = __exit__
diff --git a/tests/testpatch.py b/tests/testpatch.py
index 5c08bfc..54d3878 100644
--- a/tests/testpatch.py
+++ b/tests/testpatch.py
@@ -10,8 +10,8 @@ from tests.support import unittest2, inPy3k, SomeClass, is_instance
 
 from mock import (
     NonCallableMock, CallableMixin, patch, sentinel,
-    MagicMock, Mock, NonCallableMagicMock, patch,
-    DEFAULT, call
+    MagicMock, Mock, NonCallableMagicMock, patch, _patch,
+    DEFAULT, call, _get_target
 )
 
 builtin_string = '__builtin__'
@@ -1652,6 +1652,36 @@ class PatchTest(unittest2.TestCase):
         self.assertEqual(squizz.squozz, 3)
 
 
+    def test_patch_propogrates_exc_on_exit(self):
+
+        class holder:
+            exc_info = None, None, None
+
+        class custom_patch(_patch):
+            def __exit__(self, etype=None, val=None, tb=None):
+                _patch.__exit__(self, etype, val, tb)
+                holder.exc_info = etype, val, tb
+            stop = __exit__
+
+        def with_custom_patch(target):
+            getter, attribute = _get_target(target)
+            return custom_patch(
+                getter, attribute, DEFAULT, None, False, False,
+                False, None, {}
+            )
+
+        @with_custom_patch('squizz.squozz')
+        def test(mock):
+            raise RuntimeError
+
+        self.assertRaises(RuntimeError, test)
+        self.assertIs(holder.exc_info[0], RuntimeError)
+        self.assertIsNotNone(holder.exc_info[1],
+                            'exception value not propgated')
+        self.assertIsNotNone(holder.exc_info[2],
+                            'exception traceback not propgated')
+
+
 
 if __name__ == '__main__':
     unittest2.main()