22from contextlib import contextmanager
33
44import requests
5+ import six
56from requests .cookies import MockRequest
7+ from urllib3 import HTTPResponse
68
79from . import dispatch
810from ._compat import httplib , iteritems , gzip_decompress
@@ -23,6 +25,29 @@ def __init__(self, request, code):
2325 self .response .request = request
2426
2527
28+ class _IOReader (six .BytesIO ):
29+ """A reader that makes a BytesIO look like a HTTPResponse.
30+ A HTTPResponse will return an empty string when you read from it after
31+ the socket has been closed. A BytesIO will raise a ValueError. For
32+ compatibility we want to do the same thing a HTTPResponse does.
33+ """
34+
35+ def read (self , * args , ** kwargs ): # pylint: disable=arguments-differ
36+ if self .closed :
37+ return six .b ('' )
38+
39+ # not a new style object in python 2
40+ result = six .BytesIO .read (self , * args , ** kwargs )
41+
42+ # when using resp.iter_content(None) it'll go through a different
43+ # request path in urllib3. This path checks whether the object is
44+ # marked closed instead of the return value. see gh124.
45+ if result == six .b ('' ):
46+ self .close ()
47+
48+ return result
49+
50+
2651class FlaskLoopback (object ):
2752
2853 def __init__ (self , flask_app ):
@@ -72,12 +97,12 @@ def handle_request(self, session, url, request):
7297 with ExitStack () as stack :
7398 for handler in self ._request_context_handlers :
7499 try :
75- stack .enter_context (handler (request ))
100+ stack .enter_context (handler (request )) # pylint: disable=no-member
76101 except CustomHTTPResponse as e :
77102 return e .response
78103
79104 self ._test_client .cookie_jar .clear ()
80- for cookie in request ._cookies :
105+ for cookie in request ._cookies : # pylint: disable=protected-access
81106 self ._test_client .cookie_jar .set_cookie (cookie )
82107 resp = self ._test_client .open (path , ** open_kwargs )
83108 returned = requests .Response ()
@@ -89,8 +114,16 @@ def handle_request(self, session, url, request):
89114 resp_data = resp .get_data ()
90115 if resp .headers .get ('content-encoding' ) == 'gzip' :
91116 resp_data = gzip_decompress (resp_data )
92- returned ._content = resp_data # pylint: disable=protected-access
117+ returned ._content = resp_data # pylint: disable=protected-access
93118 returned .headers .update (resp .headers .items ())
119+ returned .raw = HTTPResponse (
120+ status = returned .status_code ,
121+ reason = returned .reason ,
122+ headers = returned .headers ,
123+ body = _IOReader (resp_data ) or _IOReader (six .b ('' )),
124+ decode_content = False ,
125+ preload_content = False ,
126+ )
94127 self ._extract_cookies (session , request , resp , returned )
95128 return returned
96129
@@ -122,8 +155,9 @@ def get_all(self, name, default=None):
122155
123156_hostname = None
124157
158+
125159def _get_hostname ():
126- global _hostname # pylint: disable=global-statement
160+ global _hostname # pylint: disable=global-statement
127161 if _hostname is None :
128162 _hostname = socket .getfqdn ()
129163 return _hostname
0 commit comments