diff --git a/packages/google-auth/google/oauth2/_client_async.py b/packages/google-auth/google/oauth2/_client_async.py index 1cb3ee0243f5..3ed5470ad8ea 100644 --- a/packages/google-auth/google/oauth2/_client_async.py +++ b/packages/google-auth/google/oauth2/_client_async.py @@ -39,8 +39,9 @@ async def _token_endpoint_request_no_throw( This function doesn't throw on response errors. Args: - request (google.auth.transport.Request): A callable used to make - HTTP requests. + request (google.auth.aio.transport.Request): A callable used to make + HTTP requests. Responses may expose ``read()`` and ``status_code`` + or the legacy ``content()`` and ``status`` interface. token_uri (str): The OAuth 2.0 authorizations server's token endpoint URI. body (Mapping[str, str]): The parameters to send in the request body. @@ -74,8 +75,11 @@ async def _token_endpoint_request_no_throw( method="POST", url=token_uri, headers=headers, body=body ) - # Using data.read() resulted in zlib decompression errors. This may require future investigation. - response_body1 = await response.content() + if hasattr(response, "read"): + response_body1 = await response.read() + else: + # Legacy content() decodes compressed bodies, unlike data.read(). + response_body1 = await response.content() response_body = ( response_body1.decode("utf-8") @@ -88,11 +92,17 @@ async def _token_endpoint_request_no_throw( except ValueError: response_data = response_body - if response.status == http_client.OK: + status_code = ( + response.status_code + if hasattr(response, "status_code") + else response.status + ) + + if status_code == http_client.OK: return True, response_data, None retryable_error = client._can_retry( - status_code=response.status, response_data=response_data + status_code=status_code, response_data=response_data ) if not can_retry or not retryable_error: @@ -107,8 +117,9 @@ async def _token_endpoint_request( """Makes a request to the OAuth 2.0 authorization server's token endpoint. Args: - request (google.auth.transport.Request): A callable used to make - HTTP requests. + request (google.auth.aio.transport.Request): A callable used to make + HTTP requests. Responses may expose ``read()`` and ``status_code`` + or the legacy ``content()`` and ``status`` interface. token_uri (str): The OAuth 2.0 authorizations server's token endpoint URI. body (Mapping[str, str]): The parameters to send in the request body. diff --git a/packages/google-auth/tests_async/oauth2/test__client_async.py b/packages/google-auth/tests_async/oauth2/test__client_async.py index c572c3e791ad..3a1602036d7e 100644 --- a/packages/google-auth/tests_async/oauth2/test__client_async.py +++ b/packages/google-auth/tests_async/oauth2/test__client_async.py @@ -14,6 +14,7 @@ import asyncio import datetime +import gzip import http.client as http_client import json import urllib @@ -24,6 +25,8 @@ from google.auth import _helpers, exceptions from google.auth import _jwt_async as jwt from google.auth.aio import transport as aio_transport +from google.auth.aio.transport import aiohttp as aiohttp_transport +from google.auth.transport import _aiohttp_requests from google.oauth2 import _client as sync_client from google.oauth2 import _client_async as _client from tests.oauth2 import test__client as test_client @@ -53,8 +56,9 @@ def make_aio_request(response_data, status_code=http_client.OK, text=False): @pytest.mark.asyncio -async def test__token_endpoint_request(): - request = make_request({"test": "response"}) +@pytest.mark.parametrize("request_factory", [make_request, make_aio_request]) +async def test__token_endpoint_request(request_factory): + request = request_factory({"test": "response"}) result = await _client._token_endpoint_request( request, "http://example.com", {"test": "params"} @@ -73,8 +77,9 @@ async def test__token_endpoint_request(): @pytest.mark.asyncio -async def test__token_endpoint_request_text(): - request = make_request("response", text=True) +@pytest.mark.parametrize("request_factory", [make_request, make_aio_request]) +async def test__token_endpoint_request_text(request_factory): + request = request_factory("response", text=True) result = await _client._token_endpoint_request( request, "http://example.com", {"test": "params"} @@ -93,8 +98,9 @@ async def test__token_endpoint_request_text(): @pytest.mark.asyncio -async def test__token_endpoint_request_json(): - request = make_request({"test": "response"}) +@pytest.mark.parametrize("request_factory", [make_request, make_aio_request]) +async def test__token_endpoint_request_json(request_factory): + request = request_factory({"test": "response"}) access_token = "access_token" result = await _client._token_endpoint_request( @@ -121,14 +127,42 @@ async def test__token_endpoint_request_json(): @pytest.mark.asyncio -async def test__token_endpoint_request_error(): - request = make_request({}, status=http_client.BAD_REQUEST) +@pytest.mark.parametrize("request_factory", [make_request, make_aio_request]) +async def test__token_endpoint_request_error(request_factory): + request = request_factory({}, http_client.BAD_REQUEST) with pytest.raises(exceptions.RefreshError) as excinfo: await _client._token_endpoint_request(request, "http://example.com", {}) assert not excinfo.value.retryable +@pytest.mark.asyncio +async def test__token_endpoint_request_aiohttp_response(): + response = mock.Mock() + response.status = http_client.OK + response.read = mock.AsyncMock(return_value=b'{"access_token": "token"}') + request = mock.AsyncMock(return_value=aiohttp_transport.Response(response)) + + result = await _client._token_endpoint_request(request, "http://example.com", {}) + + assert result == {"access_token": "token"} + + +@pytest.mark.asyncio +async def test__token_endpoint_request_legacy_compressed_response(): + response = mock.Mock() + response.status = http_client.OK + response.headers = {"Content-Encoding": "gzip"} + response.content.read = mock.AsyncMock( + return_value=gzip.compress(b'{"access_token": "token"}') + ) + request = mock.AsyncMock(return_value=_aiohttp_requests._CombinedResponse(response)) + + result = await _client._token_endpoint_request(request, "http://example.com", {}) + + assert result == {"access_token": "token"} + + @pytest.mark.asyncio async def test__token_endpoint_request_internal_failure_error(): request = make_request( @@ -485,9 +519,12 @@ async def test_refresh_grant_retry_with_retry( @pytest.mark.asyncio @pytest.mark.parametrize("can_retry", [True, False]) +@pytest.mark.parametrize("request_factory", [make_request, make_aio_request]) @mock.patch("time.sleep", return_value=None) -async def test__token_endpoint_request_no_throw_with_retry(mock_sleep, can_retry): - mock_request = make_request( +async def test__token_endpoint_request_no_throw_with_retry( + mock_sleep, request_factory, can_retry +): + mock_request = request_factory( {"error": "help", "error_description": "I'm alive"}, http_client.INTERNAL_SERVER_ERROR, )