Repository: encode/httpx
Pull request: Add Auth.requires_response_body attribute
If set then responses are read by the client before being sent back into the auth flow.

This is a work-in-progress. This branch has no test coverage right now because I don't fully understand the implications of this change. Currently there are five tests which fail if `requires_response_body` is set to `True`. These are all asserting that the response hasn't yet been read, or it isn't closed.

I can't figure out a way to ensure these tests all pass with `requires_response_body=True`. It's difficult because the response needs to be read before passing it back to `auth_flow`, but we need to pass it back in order to get the `StopIteration` to determine when to return.

Fixes #802 

Problem statement:
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.

Hints from issue 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.

Commit message: Update tests and docs

Patch to review:
diff --git a/docs/advanced.md b/docs/advanced.md
index 809d84383c..895ee66e8d 100644
--- a/docs/advanced.md
+++ b/docs/advanced.md
@@ -421,6 +421,30 @@ class MyCustomAuth(httpx.Auth):
         ...
 ```
 
+Similarly, if you are implementing a scheme that requires access to the response body, then use the `requires_response_body` property.
+
+```python
+class MyCustomAuth(httpx.Auth):
+    requires_response_body = True
+
+    def __init__(self, token, refresh_url):
+        self.token = token
+        self.refresh_url = refresh_url
+
+    def auth_flow(self, request):
+        request.headers["X-Auth"] = self.token["access"]
+        if response.status_code == 401:
+            # refresh the token
+            refresh_request = httpx.Request(
+                "POST", refresh_url, data={"refresh_token": self.token["refresh"]}
+            )
+            response = yield refresh_request
+            self.token = response.json()
+
+            request.headers["X-Auth"] = self.token["access"]
+            yield requestclass MyCustomAuth(httpx.Auth):
+```
+
 ## SSL certificates
 
 When making a request over HTTPS, HTTPX needs to verify the identity of the requested host. To do this, it uses a bundle of SSL certificates (a.k.a. CA bundle) delivered by a trusted certificate authority (CA).
diff --git a/httpx/_auth.py b/httpx/_auth.py
index c1fa892535..a8a9e62446 100644
--- 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
index 109c8351f8..6c87858947 100644
--- 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:
diff --git a/tests/client/test_auth.py b/tests/client/test_auth.py
index 10bdb87151..4c53fd9751 100644
--- a/tests/client/test_auth.py
+++ b/tests/client/test_auth.py
@@ -467,3 +467,37 @@ 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 auth_flow(
+            self, request: Request
+        ) -> typing.Generator[Request, Response, None]:
+            request.headers["Authorization"] = "xyz"
+            response = yield request
+            data = response.text
+            request.headers["Authorization"] = data
+            yield request
+
+    url = "https://example.org/"
+    auth = ResponseBodyAuth()
+    client = AsyncClient(dispatch=MockDispatch())
+
+    response = await client.get(url, auth=auth)
+    assert response.status_code == 200
+    assert response.json() == {"auth": '{"auth": "xyz"}'}


Code review comments:
diff_hunk: @@ -421,6 +421,30 @@ class MyCustomAuth(httpx.Auth):
         ...
 ```
 
+Similarly, if you are implementing a scheme that requires access to the response body, then use the `requires_response_body` property.
+
+```python
+class MyCustomAuth(httpx.Auth):
+    requires_response_body = True
+
+    def __init__(self, token, refresh_url):
+        self.token = token
+        self.refresh_url = refresh_url
+
+    def auth_flow(self, request):
+        request.headers["X-Auth"] = self.token["access"]
+        if response.status_code == 401:
+            # refresh the token
+            refresh_request = httpx.Request(
+                "POST", refresh_url, data={"refresh_token": self.token["refresh"]}
+            )
+            response = yield refresh_request
+            self.token = response.json()
+
+            request.headers["X-Auth"] = self.token["access"]
+            yield requestclass MyCustomAuth(httpx.Auth):
path: docs/advanced.md
text: @user1:
May I suggest this alternative stripped-down snippet?

```python
class MyCustomAuth(httpx.Auth):
    requires_response_body = True

    def __init__(self, access_token, refresh_token, refresh_url):
        self.access_token = access_token
        self.refresh_token = refresh_token
        self.refresh_url = refresh_url

    def auth_flow(self, request):
        request.headers["X-Authentication"] = self.access_token
        response = yield request

        if response.status_code == 401:
            # If the server issues a 401 response, then issue a request to
            # refresh tokens, and resend the request.
            refresh_response = yield self.build_refresh_request()
            self.update_tokens(refresh_response)

            request.headers["X-Authentication"] = self.access_token
            yield request

    def build_refresh_request(self):
        # Return an `httpx.Request` for refreshing tokens.
        ...

    def update_tokens(self, response):
        # Update the `.access_token` and `.refresh_token` tokens
        # based on a refresh response.
        data = response.json()
        ...
```

@author:
Great, I like it.

diff_hunk: @@ -421,6 +421,30 @@ class MyCustomAuth(httpx.Auth):
         ...
 ```
 
+Similarly, if you are implementing a scheme that requires access to the response body, then use the `requires_response_body` property.
+
+```python
+class MyCustomAuth(httpx.Auth):
+    requires_response_body = True
+
+    def __init__(self, token, refresh_url):
+        self.token = token
+        self.refresh_url = refresh_url
+
+    def auth_flow(self, request):
+        request.headers["X-Auth"] = self.token["access"]
+        if response.status_code == 401:
+            # refresh the token
+            refresh_request = httpx.Request(
+                "POST", refresh_url, data={"refresh_token": self.token["refresh"]}
+            )
+            response = yield refresh_request
+            self.token = response.json()
+
+            request.headers["X-Auth"] = self.token["access"]
+            yield requestclass MyCustomAuth(httpx.Auth):
path: docs/advanced.md
text: @user1:
Probably a typo?

```suggestion
            yield request
```

@author:
Oops, something went wrong with my markdown editing it seems.

diff_hunk: @@ -467,3 +467,37 @@ def auth_flow(
 
     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 auth_flow(
+            self, request: Request
+        ) -> typing.Generator[Request, Response, None]:
+            request.headers["Authorization"] = "xyz"
+            response = yield request
+            data = response.text
+            request.headers["Authorization"] = data
+            yield request
+
+    url = "https://example.org/"
+    auth = ResponseBodyAuth()
path: tests/client/test_auth.py
text: @user1:
For clarity, can we make it so that `"xyz"` is passed as a parameter here?

```suggestion
    auth = ResponseBodyAuth("xyz")
```