From d20faf36d991ccc0ec432bf3c8049b0f136a4405 Mon Sep 17 00:00:00 2001
From: Stan Seibert <stan.seibert@continuum.io>
Date: Fri, 27 Jun 2014 17:01:21 -0500
Subject: [PATCH] Add support for response.raw attribute.

---
 httmock.py | 19 ++++++++++++++-----
 tests.py   | 14 ++++++++++++++
 2 files changed, 28 insertions(+), 5 deletions(-)

diff --git a/httmock.py b/httmock.py
index b78d8b3..3091927 100644
--- a/httmock.py
+++ b/httmock.py
@@ -10,6 +10,10 @@
     import urlparse
 except ImportError:
     import urllib.parse as urlparse
+try:
+    from cStringIO import StringIO
+except ImportError:
+    from io import StringIO
 
 if sys.version_info >= (3, 0, 0):
     basestring = str
@@ -27,7 +31,7 @@ def getheaders(self, name):
 
 
 def response(status_code=200, content='', headers=None, reason=None, elapsed=0,
-             request=None):
+             request=None, stream=False):
     res = requests.Response()
     res.status_code = status_code
     if isinstance(content, dict):
@@ -41,6 +45,10 @@ def response(status_code=200, content='', headers=None, reason=None, elapsed=0,
     res.reason = reason
     res.elapsed = datetime.timedelta(elapsed)
     res.request = request
+    if stream:
+        res.raw = StringIO(content if content is not None else '')
+    else:
+        res.raw = StringIO('')
     if hasattr(request, 'url'):
         res.url = request.url
         if isinstance(request.url, bytes):
@@ -106,7 +114,7 @@ def __enter__(self):
         self._real_session_send = requests.Session.send
 
         def _fake_send(session, request, **kwargs):
-            response = self.intercept(request)
+            response = self.intercept(request, stream=kwargs['stream'])
             if isinstance(response, requests.Response):
                 # this is pasted from requests to handle redirects properly:
                 kwargs.setdefault('stream', session.stream)
@@ -145,7 +153,7 @@ def _fake_send(session, request, **kwargs):
     def __exit__(self, exc_type, exc_val, exc_tb):
         requests.Session.send = self._real_session_send
 
-    def intercept(self, request):
+    def intercept(self, request, stream):
         url = urlparse.urlsplit(request.url)
         res = first_of(self.handlers, url, request)
         if isinstance(res, requests.Response):
@@ -156,9 +164,10 @@ def intercept(self, request):
                             res.get('headers'),
                             res.get('reason'),
                             res.get('elapsed', 0),
-                            request)
+                            request,
+                            stream=stream)
         elif isinstance(res, basestring):
-            return response(content=res)
+            return response(content=res, stream=stream)
         elif res is None:
             return None
         else:
diff --git a/tests.py b/tests.py
index c8ee0d9..4ef1f4d 100644
--- a/tests.py
+++ b/tests.py
@@ -63,6 +63,11 @@ def test_method_fallback(self):
             r = requests.get('http://example.com/')
         self.assertEqual(r.content, 'Hello from example.com')
 
+    def test_stream(self):
+        with HTTMock(any_mock):
+            r = requests.get('http://example.com/', stream=True)
+        self.assertEqual(r.raw.read(), 'Hello from example.com')
+
     def test_netloc_fallback(self):
         with HTTMock(google_mock, facebook_mock):
             r = requests.get('http://google.com/')
@@ -187,6 +192,14 @@ def test_response_headers(self):
         r = response(200, None, {'Content-Type': 'application/json'})
         self.assertEqual(r.headers['content-type'], 'application/json')
 
+    def test_response_stream_raw(self):
+        r = response(200, 'content', {'Content-Type': 'application/json'}, stream=True)
+        self.assertEqual(r.raw.read(), 'content')
+
+    def test_response_nostream_raw(self):
+        r = response(200, 'content', {'Content-Type': 'application/json'})
+        self.assertEqual(r.raw.read(), '')
+
     def test_response_cookies(self):
         @all_requests
         def response_content(url, request):
@@ -224,3 +237,4 @@ def get_mock(url, request):
             response = requests.get('http://example.com/')
             self.assertEqual(len(response.history), 1)
             self.assertEqual(response.content, 'Hello from Google')
+