Repository: encode/httpx

Issue:
Can't read response in Auth.auth_flow
I want to implement auto-refreshing of a token but in order to get the new token I need to be able to read the server response in the custom auth_flow method.

It seems like a trivial change to get the client to `read()` or `aread()` the response before sending it back (like it does with the request), and things actually seem to work fine, but some tests break which I don't understand so wanted to check that I'm not doing completely the wrong thing here.


Discussion:
Hi!

I assume you’d like to access the response body because the server returns an auth token as a JSON payload, eg `{"token": "123abc"}`, right?

Being able to access the request body is documented at the end of [Customizing Authentication](https://www.python-httpx.org/advanced/#customizing-authentication), but we don’t have such a flag for responses yet.

Sounds like a sensible symmetric addition to me, though! Want to throw a PR together to see what it would look like? :)
Indeed, adding a `requires_response_body` attribute would make sense.
OK, great! I'll have a go at it today.

Resolving patch:
diff --git a/httpx/_auth.py b/httpx/_auth.py
--- a/httpx/_auth.py
+++ b/httpx/_auth.py
@@ -23,6 +23,7 @@ class Auth:
     """
 
     requires_request_body = False
+    requires_response_body = False
 
     def auth_flow(self, request: Request) -> typing.Generator[Request, Response, None]:
         """
diff --git a/httpx/_client.py b/httpx/_client.py
--- a/httpx/_client.py
+++ b/httpx/_client.py
@@ -657,6 +657,8 @@ def send_handling_auth(
         request = next(auth_flow)
         while True:
             response = self.send_single_request(request, timeout)
+            if auth.requires_response_body:
+                response.read()
             try:
                 next_request = auth_flow.send(response)
             except StopIteration:
@@ -1182,6 +1184,8 @@ async def send_handling_auth(
         request = next(auth_flow)
         while True:
             response = await self.send_single_request(request, timeout)
+            if auth.requires_response_body:
+                await response.aread()
             try:
                 next_request = auth_flow.send(response)
             except StopIteration:


Test patch:
diff --git a/tests/client/test_auth.py b/tests/client/test_auth.py
--- a/tests/client/test_auth.py
+++ b/tests/client/test_auth.py
@@ -467,3 +467,40 @@ async def streaming_body():
 
     with pytest.raises(RequestBodyUnavailable):
         await client.post(url, data=streaming_body(), auth=auth)
+
+
+@pytest.mark.asyncio
+async def test_auth_reads_response_body() -> None:
+    """
+    Test that we can read the response body in an auth flow if `requires_response_body`
+    is set.
+    """
+
+    class ResponseBodyAuth(Auth):
+        """
+        A mock authentication scheme that requires clients to send an 'Authorization'
+        header, then send back the contents of the response in the 'Authorization'
+        header.
+        """
+
+        requires_response_body = True
+
+        def __init__(self, token):
+            self.token = token
+
+        def auth_flow(
+            self, request: Request
+        ) -> typing.Generator[Request, Response, None]:
+            request.headers["Authorization"] = self.token
+            response = yield request
+            data = response.text
+            request.headers["Authorization"] = data
+            yield request
+
+    url = "https://example.org/"
+    auth = ResponseBodyAuth("xyz")
+    client = AsyncClient(dispatch=MockDispatch())
+
+    response = await client.get(url, auth=auth)
+    assert response.status_code == 200
+    assert response.json() == {"auth": '{"auth": "xyz"}'}