Handle socket errors in iter_content

diff --git a/requests/models.py b/requests/models.py
index 0dc5556838c734e0a24d4795dcf16cd6003871c5..2ababaf8daa43cdd6881db5beb31b9a5759501e6 100644 (file)
--- a/requests/models.py
+++ b/requests/models.py
@@ -9,6 +9,7 @@ This module contains the primary objects that power Requests.
 
 import collections
 import datetime
+import socket
 
 from io import BytesIO, UnsupportedOperation
 from .hooks import default_hooks
@@ -22,7 +23,7 @@ from .packages.urllib3.util import parse_url
 from .packages.urllib3.exceptions import DecodeError
 from .exceptions import (
     HTTPError, RequestException, MissingSchema, InvalidURL,
-    ChunkedEncodingError, ContentDecodingError)
+    ChunkedEncodingError, ContentDecodingError, ConnectionError)
 from .utils import (
     guess_filename, get_auth_from_url, requote_uri,
     stream_decode_response_unicode, to_key_val_list, parse_header_links,
@@ -640,6 +641,8 @@ class Response(object):
                     raise ChunkedEncodingError(e)
                 except DecodeError as e:
                     raise ContentDecodingError(e)
+                except socket.error as e:
+                    raise ConnectionError(e)
             except AttributeError:
                 # Standard file-like object.
                 while True:

diff --git a/test_requests.py b/test_requests.py
index 283353b9196d965ae027ea0037ee89f6a25494b5..004d6e8f3a8bef2769ba3e1344c15b6a3a1b1b4f 100755 (executable)
--- a/test_requests.py
+++ b/test_requests.py
@@ -18,7 +18,7 @@ from requests.auth import HTTPDigestAuth, _basic_auth_str
 from requests.compat import (
     Morsel, cookielib, getproxies, str, urljoin, urlparse, is_py3, builtin_str)
 from requests.cookies import cookiejar_from_dict, morsel_to_cookie
-from requests.exceptions import InvalidURL, MissingSchema
+from requests.exceptions import InvalidURL, MissingSchema, ConnectionError
 from requests.models import PreparedRequest
 from requests.structures import CaseInsensitiveDict
 from requests.sessions import SessionRedirectMixin
@@ -720,6 +720,18 @@ class RequestsTestCase(unittest.TestCase):
         assert next(iter(r))
         io.close()
 
+    def test_iter_content_handles_socket_error(self):
+        r = requests.Response()
+        import socket
+
+        class RawMock(object):
+            def stream(self, chunk_size, decode_content=None):
+                raise socket.error()
+
+        setattr(r, 'raw', RawMock())
+        with pytest.raises(ConnectionError):
+            list(r.iter_content())
+
     def test_response_decode_unicode(self):
         """
         When called with decode_unicode, Response.iter_content should always