Map httpcore exceptions for Response read methods (#1190)

* Map httpcore exceptions for `Response.iter_*` methods * Tweak * Change wording Co-authored-by: Florimond Manca <[email protected]> Co-authored-by: Florimond Manca <[email protected]> Co-authored-by: Tom Christie <[email protected]>
encode · Aug 19, 2020 · 03cd88c · 03cd88c
1 parent d10b7cd
commit 03cd88c
Show file tree

Hide file tree

Showing 3 changed files with 31 additions and 5 deletions.
diff --git a/httpx/_models.py b/httpx/_models.py
@@ -24,6 +24,7 @@
     TextDecoder,
 )
 from ._exceptions import (
+    HTTPCORE_EXC_MAP,
     CookieConflict,
     HTTPStatusError,
     InvalidURL,
@@ -32,6 +33,7 @@
     ResponseClosed,
     ResponseNotRead,
     StreamConsumed,
+    map_exceptions,
 )
 from ._status_codes import codes
 from ._types import (
@@ -931,8 +933,9 @@ def iter_raw(self) -> typing.Iterator[bytes]:
             raise ResponseClosed()
 
         self.is_stream_consumed = True
-        for part in self._raw_stream:
-            yield part
+        with map_exceptions(HTTPCORE_EXC_MAP, request=self.request):
+            for part in self._raw_stream:
+                yield part
         self.close()
 
     def next(self) -> "Response":
@@ -1007,8 +1010,9 @@ async def aiter_raw(self) -> typing.AsyncIterator[bytes]:
             raise ResponseClosed()
 
         self.is_stream_consumed = True
-        async for part in self._raw_stream:
-            yield part
+        with map_exceptions(HTTPCORE_EXC_MAP, request=self.request):
+            async for part in self._raw_stream:
+                yield part
         await self.aclose()
 
     async def anext(self) -> "Response":

diff --git a/tests/conftest.py b/tests/conftest.py
@@ -76,6 +76,8 @@ async def app(scope, receive, send):
     assert scope["type"] == "http"
     if scope["path"].startswith("/slow_response"):
         await slow_response(scope, receive, send)
+    elif scope["path"].startswith("/slow_stream_response"):
+        await slow_stream_response(scope, receive, send)
     elif scope["path"].startswith("/status"):
         await status_code(scope, receive, send)
     elif scope["path"].startswith("/echo_body"):
@@ -111,6 +113,19 @@ async def slow_response(scope, receive, send):
     await send({"type": "http.response.body", "body": b"Hello, world!"})
 
 
+async def slow_stream_response(scope, receive, send):
+    await send(
+        {
+            "type": "http.response.start",
+            "status": 200,
+            "headers": [[b"content-type", b"text/plain"]],
+        }
+    )
+
+    await sleep(1)
+    await send({"type": "http.response.body", "body": b"", "more_body": False})
+
+
 async def status_code(scope, receive, send):
     status_code = int(scope["path"].replace("/status/", ""))
     await send(

diff --git a/tests/test_exceptions.py b/tests/test_exceptions.py
@@ -24,7 +24,7 @@ def test_httpcore_all_exceptions_mapped() -> None:
         pytest.fail(f"Unmapped httpcore exceptions: {not_mapped}")
 
 
-def test_httpcore_exception_mapping() -> None:
+def test_httpcore_exception_mapping(server) -> None:
     """
     HTTPCore exception mapping works as expected.
     """
@@ -33,6 +33,13 @@ def test_httpcore_exception_mapping() -> None:
     with pytest.raises(httpx.ConnectError):
         httpx.get("http://doesnotexist")
 
+    # Make sure streaming methods also map exceptions.
+    url = server.url.copy_with(path="/slow_stream_response")
+    timeout = httpx.Timeout(None, read=0.1)
+    with httpx.stream("GET", url, timeout=timeout) as stream:
+        with pytest.raises(httpx.ReadTimeout):
+            stream.read()
+
     # Make sure it also works with custom transports.
     class MockTransport(httpcore.SyncHTTPTransport):
         def request(self, *args: Any, **kwargs: Any) -> Any: