Merge pull request #16 from d1618033/issue-15

Add raw field

Author: vmalloc

flask_loopback/flask_loopback.py

@@ -2,7 +2,9 @@
 from contextlib import contextmanager
 
 import requests
+import six
 from requests.cookies import MockRequest
+from urllib3 import HTTPResponse
 
 from . import dispatch
 from ._compat import httplib, iteritems, gzip_decompress
@@ -23,6 +25,29 @@ def __init__(self, request, code):
         self.response.request = request
 
 
+class _IOReader(six.BytesIO):
+    """A reader that makes a BytesIO look like a HTTPResponse.
+    A HTTPResponse will return an empty string when you read from it after
+    the socket has been closed. A BytesIO will raise a ValueError. For
+    compatibility we want to do the same thing a HTTPResponse does.
+    """
+
+    def read(self, *args, **kwargs):  # pylint: disable=arguments-differ
+        if self.closed:
+            return six.b('')
+
+        # not a new style object in python 2
+        result = six.BytesIO.read(self, *args, **kwargs)
+
+        # when using resp.iter_content(None) it'll go through a different
+        # request path in urllib3. This path checks whether the object is
+        # marked closed instead of the return value. see gh124.
+        if result == six.b(''):
+            self.close()
+
+        return result
+
+
 class FlaskLoopback(object):
 
     def __init__(self, flask_app):
@@ -72,7 +97,7 @@ def handle_request(self, session, url, request):
         with ExitStack() as stack:
             for handler in self._request_context_handlers:
                 try:
-                    stack.enter_context(handler(request))
+                    stack.enter_context(handler(request))  # pylint: disable=no-member
                 except CustomHTTPResponse as e:
                     return e.response
 
@@ -89,8 +114,16 @@ def handle_request(self, session, url, request):
             resp_data = resp.get_data()
             if resp.headers.get('content-encoding') == 'gzip':
                 resp_data = gzip_decompress(resp_data)
-            returned._content = resp_data # pylint: disable=protected-access
+            returned._content = resp_data  # pylint: disable=protected-access
             returned.headers.update(resp.headers.items())
+            returned.raw = HTTPResponse(
+                status=returned.status_code,
+                reason=returned.reason,
+                headers=returned.headers,
+                body=_IOReader(resp_data) or _IOReader(six.b('')),
+                decode_content=False,
+                preload_content=False,
+            )
             self._extract_cookies(session, request, resp, returned)
             return returned
 
@@ -122,8 +155,9 @@ def get_all(self, name, default=None):
 
 _hostname = None
 
+
 def _get_hostname():
-    global _hostname            # pylint: disable=global-statement
+    global _hostname  # pylint: disable=global-statement
     if _hostname is None:
         _hostname = socket.getfqdn()
     return _hostname

tests/conftest.py

@@ -1,7 +1,8 @@
 import hashlib
+import os
 
 import pytest
-from flask import Flask, g, jsonify, request, Response
+from flask import Flask, g, jsonify, request, Response, send_file
 from flask_loopback import FlaskLoopback
 from flask_loopback._compat import gzip_compress
 from urlobject import URLObject as URL
@@ -65,6 +66,14 @@ def compressed():  # pylint: disable=unused-variable
         orig = "uncompressed!".encode("utf-8")
         return Response(gzip_compress(orig), headers={"Content-Encoding": "gzip"})
 
+    @returned.route("/files/binary")
+    def binary_file():  # pylint: disable=unused-variable
+        return send_file(os.path.join(os.path.dirname(__file__), "files", "binary"))
+
+    @returned.route("/files/text")
+    def text_file():  # pylint: disable=unused-variable
+        return send_file(os.path.join(os.path.dirname(__file__), "files", "text"))
+
     @returned.before_request
     def before():  # pylint: disable=unused-variable
         g.cookies = []

tests/files/binary

@@ -0,0 +1 @@
+hello world

tests/files/text

@@ -0,0 +1,22 @@
+import os
+
+import pytest
+import requests
+
+
+def read_response_using_iter_content(resp):
+    return b''.join([chunk for chunk in resp.iter_content(chunk_size=1024)])
+
+
+def read_response_using_content(resp):
+    return resp.content
+
+
+@pytest.mark.parametrize("read_response", [read_response_using_content, read_response_using_iter_content])
+@pytest.mark.parametrize("file_type", ["binary", "text"])
+def test_download_file(active_app, url, read_response, file_type):  # pylint: disable=unused-argument
+    resp = requests.get(url.add_path("files").add_path(file_type))
+    resp.raise_for_status()
+    result = read_response(resp)
+    with open(os.path.join(os.path.dirname(__file__), 'files', file_type), 'rb') as f:
+        assert result == f.read()