response callback work

CHANGES

@@ -2,6 +2,7 @@ Unreleased
 ----------
 
 - Allow empty list/dict as json object (GH-100)
+- Added `response_callback`
 
 0.5.1
 -----

README.rst

@@ -188,3 +188,31 @@ You can also use ``assert_all_requests_are_fired`` to add multiple responses for
             assert resp.status_code == 500
             resp = requests.get('http://twitter.com/api/1/foobar')
             assert resp.status_code == 200
+
+Using a callback to modify the response
+---------------------------------------
+
+If you use customized processing in `requests` via subclassing/mixins, or if you
+have library tools that interact with `requests` at a low level, you may need
+to add extended processing to the mocked Response object to fully simlulate the
+environment for your tests.  A `response_callback` can be used, which will be
+wrapped by the library before being returned to the caller.  The callback
+accepts a `response` as it's single argument, and is expected to return a 
+single `response` object.
+
+
+.. code-block:: python
+
+    import responses
+    import requests
+
+	def response_callback(resp):
+		resp.callback_processed = True
+		return resp
+
+	with responses.RequestsMock(response_callback=response_callback) as m:
+		m.add(responses.GET, 'http://example.com', body=b'test')
+		resp = requests.get('http://example.com')
+		assert resp.text == "test"
+		assert hasattr(resp, 'callback_processed')
+		assert resp.callback_processed is True

responses.py

@@ -118,11 +118,14 @@ class RequestsMock(object):
     PATCH = 'PATCH'
     POST = 'POST'
     PUT = 'PUT'
+    response_callback = None
 
-    def __init__(self, assert_all_requests_are_fired=True):
+    def __init__(self, assert_all_requests_are_fired=True,
+                 response_callback=None):
         self._calls = CallList()
         self.reset()
         self.assert_all_requests_are_fired = assert_all_requests_are_fired
+        self.response_callback = response_callback
 
     def reset(self):
         self._urls = []
@@ -230,6 +233,7 @@ def _has_strict_url_match(self, url, other):
 
     def _on_request(self, adapter, request, **kwargs):
         match = self._find_match(request)
+        resp_callback = self.response_callback
         # TODO(dcramer): find the correct class for this
         if match is None:
             error_msg = 'Connection refused: {0} {1}'.format(request.method,
@@ -238,6 +242,7 @@ def _on_request(self, adapter, request, **kwargs):
             response.request = request
 
             self._calls.add(request, response)
+            response = resp_callback(response) if resp_callback else response
             raise response
 
         if 'body' in match and isinstance(match['body'], Exception):
@@ -283,6 +288,7 @@ def _on_request(self, adapter, request, **kwargs):
         except (KeyError, TypeError):
             pass
 
+        response = resp_callback(response) if resp_callback else response
         self._calls.add(request, response)
 
         return response

test_responses.py

@@ -373,6 +373,22 @@ def run():
     assert_reset()
 
 
+def test_response_callback():
+    """adds a callback to decorate the response, then checks it"""
+    def run():
+        def response_callback(resp):
+            resp._is_mocked = True
+            return resp
+        with responses.RequestsMock(response_callback=response_callback) as m:
+            m.add(responses.GET, 'http://example.com', body=b'test')
+            resp = requests.get('http://example.com')
+            assert resp.text == "test"
+            assert hasattr(resp, '_is_mocked')
+            assert resp._is_mocked is True
+    run()
+    assert_reset()
+
+
 def test_assert_all_requests_are_fired():
     def run():
         with pytest.raises(AssertionError) as excinfo: