Allow raw parameter to accept callable like body elements

Author: ramast

doc/source/response.rst

@@ -70,6 +70,9 @@ Dynamic Response
 ================
 
 A callback can be provided in place of any of the body elements.
+raw attribute also accepts callable that returns HTTPResponse.
+The HTTPResponse should have `preload_content=False` or it may not work properly.
+
 Callbacks must be a function in the form of
 
 .. code:: python

releasenotes/notes/make-raw-accept-a-callback-973581da158430ab.yaml

@@ -0,0 +1,6 @@
+---
+features:
+  - |
+    When using dynamic responses, you can now use a callback function for the
+    raw HTTPResponse object. This aligns it with all the other parameters that
+    can be mocked.

requests_mock/adapter.pyi

@@ -66,7 +66,7 @@ class Adapter(BaseAdapter, _RequestHistoryTracker):
         text: Union[str, Callback[str]] = ...,
         content: Union[bytes, Callback[bytes]] = ...,
         body: Union[IOBase, Callback[IOBase]] = ...,
-        raw: HTTPResponse = ...,
+        raw: Union[HTTPResponse, Callback[HTTPResponse]] = ...,
         exc: Union[Exception, Type[Exception]] = ...,
         additional_matcher: AdditionalMatcher = ...,
         **kwargs: Any

requests_mock/mocker.pyi

@@ -54,7 +54,7 @@ class MockerCore:
       text: Union[str, Callback[str]] = ...,
       content: Union[bytes, Callback[bytes]] = ...,
       body: Union[IOBase, Callback[IOBase]] = ...,
-      raw: HTTPResponse = ...,
+      raw: Union[HTTPResponse, Callback[HTTPResponse]] = ...,
       exc: Union[Exception, Type[Exception]] = ...,
       additional_matcher: AdditionalMatcher = ...,
       json_encoder: Optional[Type[JSONEncoder]] = ...,
@@ -77,7 +77,7 @@ class MockerCore:
       text: Union[str, Callback[str]] = ...,
       content: Union[bytes, Callback[bytes]] = ...,
       body: Union[IOBase, Callback[IOBase]] = ...,
-      raw: HTTPResponse = ...,
+      raw: Union[HTTPResponse, Callback[HTTPResponse]] = ...,
       exc: Union[Exception, Type[Exception]] = ...,
       additional_matcher: AdditionalMatcher = ...,
       json_encoder: Optional[Type[JSONEncoder]] = ...,
@@ -99,7 +99,7 @@ class MockerCore:
       text: Union[str, Callback[str]] = ...,
       content: Union[bytes, Callback[bytes]] = ...,
       body: Union[IOBase, Callback[IOBase]] = ...,
-      raw: HTTPResponse = ...,
+      raw: Union[HTTPResponse, Callback[HTTPResponse]] = ...,
       exc: Union[Exception, Type[Exception]] = ...,
       additional_matcher: AdditionalMatcher = ...,
       json_encoder: Optional[Type[JSONEncoder]] = ...,
@@ -121,7 +121,7 @@ class MockerCore:
       text: Union[str, Callback[str]] = ...,
       content: Union[bytes, Callback[bytes]] = ...,
       body: Union[IOBase, Callback[IOBase]] = ...,
-      raw: HTTPResponse = ...,
+      raw: Union[HTTPResponse, Callback[HTTPResponse]] = ...,
       exc: Union[Exception, Type[Exception]] = ...,
       additional_matcher: AdditionalMatcher = ...,
       json_encoder: Optional[Type[JSONEncoder]] = ...,
@@ -143,7 +143,7 @@ class MockerCore:
       text: Union[str, Callback[str]] = ...,
       content: Union[bytes, Callback[bytes]] = ...,
       body: Union[IOBase, Callback[IOBase]] = ...,
-      raw: HTTPResponse = ...,
+      raw: Union[HTTPResponse, Callback[HTTPResponse]] = ...,
       exc: Union[Exception, Type[Exception]] = ...,
       additional_matcher: AdditionalMatcher = ...,
       json_encoder: Optional[Type[JSONEncoder]] = ...,
@@ -165,7 +165,7 @@ class MockerCore:
       text: Union[str, Callback[str]] = ...,
       content: Union[bytes, Callback[bytes]] = ...,
       body: Union[IOBase, Callback[IOBase]] = ...,
-      raw: HTTPResponse = ...,
+      raw: Union[HTTPResponse, Callback[HTTPResponse]] = ...,
       exc: Union[Exception, Type[Exception]] = ...,
       additional_matcher: AdditionalMatcher = ...,
       json_encoder: Optional[Type[JSONEncoder]] = ...,
@@ -187,7 +187,7 @@ class MockerCore:
       text: Union[str, Callback[str]] = ...,
       content: Union[bytes, Callback[bytes]] = ...,
       body: Union[IOBase, Callback[IOBase]] = ...,
-      raw: HTTPResponse = ...,
+      raw: Union[HTTPResponse, Callback[HTTPResponse]] = ...,
       exc: Union[Exception, Type[Exception]] = ...,
       additional_matcher: AdditionalMatcher = ...,
       json_encoder: Optional[Type[JSONEncoder]] = ...,
@@ -209,7 +209,7 @@ class MockerCore:
       text: Union[str, Callback[str]] = ...,
       content: Union[bytes, Callback[bytes]] = ...,
       body: Union[IOBase, Callback[IOBase]] = ...,
-      raw: HTTPResponse = ...,
+      raw: Union[HTTPResponse, Callback[HTTPResponse]] = ...,
       exc: Union[Exception, Type[Exception]] = ...,
       additional_matcher: AdditionalMatcher = ...,
       json_encoder: Optional[Type[JSONEncoder]] = ...,
@@ -231,7 +231,7 @@ class MockerCore:
       text: Union[str, Callback[str]] = ...,
       content: Union[bytes, Callback[bytes]] = ...,
       body: Union[IOBase, Callback[IOBase]] = ...,
-      raw: HTTPResponse = ...,
+      raw: Union[HTTPResponse, Callback[HTTPResponse]] = ...,
       exc: Union[Exception, Type[Exception]] = ...,
       additional_matcher: AdditionalMatcher = ...,
       json_encoder: Optional[Type[JSONEncoder]] = ...,

requests_mock/response.py

@@ -270,7 +270,7 @@ def _call(f, *args, **kwargs):
                                text=_call(self._params.get('text')),
                                content=_call(self._params.get('content')),
                                body=_call(self._params.get('body')),
-                               raw=self._params.get('raw'),
+                               raw=_call(self._params.get('raw')),
                                json_encoder=self._params.get('json_encoder'),
                                status_code=context.status_code,
                                reason=context.reason,

tests/test_adapter.py

@@ -10,12 +10,15 @@
 # License for the specific language governing permissions and limitations
 # under the License.
 
+import http.client
+import io
 import json
 import re
 import urllib.parse
 
 import purl
 import requests
+from urllib3 import HTTPResponse
 
 import requests_mock
 from . import base
@@ -119,6 +122,27 @@ def _text_cb(request, context):
         self.assertHeaders(resp)
         self.assertLastRequest()
 
+    def test_raw_callback(self):
+        status_code = 401
+        data = 'testdata'
+
+        def _raw_cb(request, context):
+            return HTTPResponse(
+                status=status_code,
+                headers=self.headers,
+                body=io.BytesIO(data.encode('utf-8')),
+                preload_content=False,
+                reason=http.client.responses.get(status_code),
+            )
+
+        self.adapter.register_uri('GET', self.url, raw=_raw_cb)
+        resp = self.session.get(self.url)
+        self.assertEqual(status_code, resp.status_code)
+        self.assertEqual(data, resp.text)
+        self.assertEqual(data.encode('utf-8'), resp.content)
+        self.assertHeaders(resp)
+        self.assertLastRequest()
+
     def test_json(self):
         json_data = {'hello': 'world'}
         self.adapter.register_uri('GET',