Neaten up our try...except structure for ensuring responses (#1525)

encode · Mar 23, 2021 · 41ee5e7 · 41ee5e7
1 parent 4b4769a
commit 41ee5e7
Show file tree

Hide file tree

Showing 2 changed files with 69 additions and 54 deletions.
diff --git a/httpx/_client.py b/httpx/_client.py
@@ -762,21 +762,18 @@ def send(
             allow_redirects=allow_redirects,
             history=[],
         )
-
-        if not stream:
-            try:
+        try:
+            if not stream:
                 response.read()
-            finally:
-                response.close()
 
-        try:
             for hook in self._event_hooks["response"]:
                 hook(response)
-        except Exception:
-            response.close()
-            raise
 
-        return response
+            return response
+
+        except Exception as exc:
+            response.close()
+            raise exc
 
     def _send_handling_auth(
         self,
@@ -800,18 +797,20 @@ def _send_handling_auth(
                 history=history,
             )
             try:
-                next_request = auth_flow.send(response)
-            except StopIteration:
-                return response
-            except BaseException as exc:
-                response.close()
-                raise exc from None
-            else:
+                try:
+                    next_request = auth_flow.send(response)
+                except StopIteration:
+                    return response
+
                 response.history = list(history)
                 response.read()
                 request = next_request
                 history.append(response)
 
+            except Exception as exc:
+                response.close()
+                raise exc
+
     def _send_handling_redirects(
         self,
         request: Request,
@@ -826,19 +825,24 @@ def _send_handling_redirects(
                 )
 
             response = self._send_single_request(request, timeout)
-            response.history = list(history)
+            try:
+                response.history = list(history)
 
-            if not response.is_redirect:
-                return response
+                if not response.is_redirect:
+                    return response
 
-            if allow_redirects:
-                response.read()
-            request = self._build_redirect_request(request, response)
-            history = history + [response]
+                request = self._build_redirect_request(request, response)
+                history = history + [response]
 
-            if not allow_redirects:
-                response.next_request = request
-                return response
+                if allow_redirects:
+                    response.read()
+                else:
+                    response.next_request = request
+                    return response
+
+            except Exception as exc:
+                response.close()
+                raise exc
 
     def _send_single_request(self, request: Request, timeout: Timeout) -> Response:
         """
@@ -1394,21 +1398,18 @@ async def send(
             allow_redirects=allow_redirects,
             history=[],
         )
-
-        if not stream:
-            try:
+        try:
+            if not stream:
                 await response.aread()
-            finally:
-                await response.aclose()
 
-        try:
             for hook in self._event_hooks["response"]:
                 await hook(response)
-        except Exception:
-            await response.aclose()
-            raise
 
-        return response
+            return response
+
+        except Exception as exc:
+            await response.aclose()
+            raise exc
 
     async def _send_handling_auth(
         self,
@@ -1432,18 +1433,20 @@ async def _send_handling_auth(
                 history=history,
             )
             try:
-                next_request = await auth_flow.asend(response)
-            except StopAsyncIteration:
-                return response
-            except BaseException as exc:
-                await response.aclose()
-                raise exc from None
-            else:
+                try:
+                    next_request = await auth_flow.asend(response)
+                except StopAsyncIteration:
+                    return response
+
                 response.history = list(history)
                 await response.aread()
                 request = next_request
                 history.append(response)
 
+            except Exception as exc:
+                await response.aclose()
+                raise exc
+
     async def _send_handling_redirects(
         self,
         request: Request,
@@ -1458,19 +1461,24 @@ async def _send_handling_redirects(
                 )
 
             response = await self._send_single_request(request, timeout)
-            response.history = list(history)
+            try:
+                response.history = list(history)
 
-            if not response.is_redirect:
-                return response
+                if not response.is_redirect:
+                    return response
 
-            if allow_redirects:
-                await response.aread()
-            request = self._build_redirect_request(request, response)
-            history = history + [response]
+                request = self._build_redirect_request(request, response)
+                history = history + [response]
 
-            if not allow_redirects:
-                response.next_request = request
-                return response
+                if allow_redirects:
+                    await response.aread()
+                else:
+                    response.next_request = request
+                    return response
+
+            except Exception as exc:
+                await response.aclose()
+                raise exc
 
     async def _send_single_request(
         self, request: Request, timeout: Timeout

diff --git a/tests/client/test_redirects.py b/tests/client/test_redirects.py
@@ -393,3 +393,10 @@ def test_redirect_custom_scheme():
     with pytest.raises(httpx.UnsupportedProtocol) as e:
         client.post("https://example.org/redirect_custom_scheme")
     assert str(e.value) == "Scheme 'market' not supported."
+
+
+@pytest.mark.usefixtures("async_environment")
+async def test_async_invalid_redirect():
+    async with httpx.AsyncClient(transport=httpx.MockTransport(redirects)) as client:
+        with pytest.raises(httpx.RemoteProtocolError):
+            await client.get("http://example.org/invalid_redirect")