Handle socket errors in iter_content

This commit is contained in:
Roman Levin
2014-07-24 23:00:50 +02:00
parent fe693c4922
commit 9eab9f7415
2 changed files with 17 additions and 2 deletions
+4 -1
View File
@@ -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:
+13 -1
View File
@@ -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