Make safe_mode handle exceptions more consistently
authorOri Livneh <ori.livneh@gmail.com>
Fri, 23 Dec 2011 21:14:46 +0000 (16:14 -0500)
committerOri Livneh <ori.livneh@gmail.com>
Fri, 23 Dec 2011 21:14:46 +0000 (16:14 -0500)
requests/models.py
test_requests.py

index a51ee7001065ac367990ed98ad3a8acaa20e26a2..1611bd072d413caf5657970083352871eb2be6c3 100644 (file)
@@ -18,13 +18,15 @@ from .structures import CaseInsensitiveDict
 from .status_codes import codes
 from .packages import oreos
 from .auth import HTTPBasicAuth, HTTPProxyAuth
+from .packages.urllib3.response import HTTPResponse
 from .packages.urllib3.exceptions import MaxRetryError
 from .packages.urllib3.exceptions import SSLError as _SSLError
 from .packages.urllib3.exceptions import HTTPError as _HTTPError
 from .packages.urllib3 import connectionpool, poolmanager
 from .packages.urllib3.filepost import encode_multipart_formdata
 from .exceptions import (
-    Timeout, URLRequired, TooManyRedirects, HTTPError, ConnectionError)
+    ConnectionError, HTTPError, RequestException, Timeout, TooManyRedirects,
+    URLRequired)
 from .utils import (
     get_encoding_from_headers, stream_decode_response_unicode,
     decode_gzip, stream_decode_gzip, guess_filename, requote_path)
@@ -171,6 +173,9 @@ class Request(object):
                 # Save cookies in Response.
                 response.cookies = cookies
 
+                # No exceptions were harmed in the making of this request.
+                response.error = getattr(resp, 'error', None)
+
             # Save original response for later.
             response.raw = resp
 
@@ -438,35 +443,41 @@ class Request(object):
                     # Attach Cookie header to request.
                     self.headers['Cookie'] = cookie_header
 
-            # If the request fails but exceptions are suppressed,
-            # we'll build a blank response.
-            r = None
-
             try:
-                # Send the request.
-                r = conn.urlopen(
-                    method=self.method,
-                    url=self.path_url,
-                    body=body,
-                    headers=self.headers,
-                    redirect=False,
-                    assert_same_host=False,
-                    preload_content=False,
-                    decode_content=False,
-                    retries=self.config.get('max_retries', 0),
-                    timeout=self.timeout,
-                )
-                self.sent = True
-
-
-            except MaxRetryError, e:
-                if not self.config.get('safe_mode', False):
+                # The inner try .. except re-raises certain exceptions as
+                # internal exception types; the outer suppresses exceptions
+                # when safe mode is set.
+                try:
+                    # Send the request.
+                    r = conn.urlopen(
+                        method=self.method,
+                        url=self.path_url,
+                        body=body,
+                        headers=self.headers,
+                        redirect=False,
+                        assert_same_host=False,
+                        preload_content=False,
+                        decode_content=False,
+                        retries=self.config.get('max_retries', 0),
+                        timeout=self.timeout,
+                    )
+                    self.sent = True
+
+                except MaxRetryError, e:
                     raise ConnectionError(e)
 
-            except (_SSLError, _HTTPError), e:
-                if not self.config.get('safe_mode', False):
+                except (_SSLError, _HTTPError), e:
                     raise Timeout('Request timed out.')
 
+            except RequestException, e:
+                if self.config.get('safe_mode', False):
+                    # In safe mode, catch the exception and attach it to 
+                    # a blank urllib3.HTTPResponse object.
+                    r = HTTPResponse()
+                    r.error = e
+                else:
+                    raise
+
             self._build_response(r)
 
             # Response manipulation hook.
@@ -559,13 +570,11 @@ class Response(object):
             )
 
         def generate():
-            # self.raw can be None if we're in safe_mode and the request failed
-            if self.raw is not None:
-                while 1:
-                    chunk = self.raw.read(chunk_size)
-                    if not chunk:
-                        break
-                    yield chunk
+            while 1:
+                chunk = self.raw.read(chunk_size)
+                if not chunk:
+                    break
+                yield chunk
             self._content_consumed = True
 
         gen = generate()
index 2543103f2879f3cf81cffccfb4613c67510e0e58..b3e55c8cd981a76e8b662687e0216ba0b00134be 100755 (executable)
@@ -619,18 +619,19 @@ class RequestsTestSuite(unittest.TestCase):
         lines = '\n'.join(r.iter_lines())
         self.assertEqual(lines, quote)
 
-    def test_null_response(self):
+    def test_safe_mode(self):
 
-        # Safe mode creates empty responses for failed requests.
+        safe = requests.session(config=dict(safe_mode=True))
 
+        # Safe mode creates empty responses for failed requests.
         # Iterating on these responses should produce empty sequences
-        r = requests.get('http://_/', config=dict(safe_mode=True))
+        r = safe.get('http://_/')
         self.assertEquals(list(r.iter_lines()), [])
+        self.assertIsInstance(r.error, requests.exceptions.ConnectionError)
 
-        r = requests.get('http://_/', config=dict(safe_mode=True))
+        r = safe.get('http://_/')
         self.assertEquals(list(r.iter_content()), [])
-
-    def test_timeout(self):
+        self.assertIsInstance(r.error, requests.exceptions.ConnectionError)
 
         # When not in safe mode, should raise Timeout exception
         with self.assertRaises(requests.exceptions.Timeout):
@@ -640,6 +641,7 @@ class RequestsTestSuite(unittest.TestCase):
         r = requests.get(httpbin('stream', '1000'), timeout=0.0001,
                 config=dict(safe_mode=True))
         self.assertIsNone(r.content)
+        self.assertIsInstance(r.error, requests.exceptions.Timeout)
 
 
 if __name__ == '__main__':