From 94a1d956a78acbc9e3e2391b8335a0907191a435 Mon Sep 17 00:00:00 2001 From: agrawalradhika-cell Date: Tue, 25 Aug 2026 20:56:24 -0700 Subject: [PATCH 01/79] feat: Add retry for cert rotation handling feat: Add retry for cert rotation handling --- .../google/auth/aio/transport/sessions.py | 30 +++++++++++++++++++ 1 file changed, 30 insertions(+) diff --git a/packages/google-auth/google/auth/aio/transport/sessions.py b/packages/google-auth/google/auth/aio/transport/sessions.py index d88162667bda..c55be4670549 100644 --- a/packages/google-auth/google/auth/aio/transport/sessions.py +++ b/packages/google-auth/google/auth/aio/transport/sessions.py @@ -15,6 +15,8 @@ import asyncio from contextlib import asynccontextmanager import functools +import http.client as http_client +import logging import time from typing import Mapping, Optional, TYPE_CHECKING, Union import warnings @@ -26,6 +28,8 @@ from google.auth.exceptions import TimeoutError import google.auth.transport._mtls_helper +_LOGGER = logging.getLogger(__name__) + if TYPE_CHECKING: # pragma: NO COVER import aiohttp from aiohttp import ClientTimeout # type: ignore @@ -310,6 +314,32 @@ async def request( url, method, data, headers, actual_timeout, **kwargs ) ) + if response.status_code == http_client.UNAUTHORIZED: + if self.is_mtls: + call_cert_bytes, call_key_bytes, cached_fingerprint, current_cert_fingerprint = google.auth.transport._mtls_helper.check_parameters_for_unauthorized_response( + self._cached_cert + ) + if cached_fingerprint != current_cert_fingerprint: + try: + _LOGGER.info( + "Client certificate has changed, reconfiguring mTLS " + "channel." + ) + await self.configure_mtls_channel( + lambda: (call_cert_bytes, call_key_bytes) + ) + continue + except Exception as e: + _LOGGER.error("Failed to reconfigure mTLS channel: %s", e) + raise exceptions.MutualTLSChannelError( + "Failed to reconfigure mTLS channel" + ) from e + else: + _LOGGER.info( + "Skipping reconfiguration of mTLS channel because the client" + " certificate has not changed." + ) + if response.status_code not in transport.DEFAULT_RETRYABLE_STATUS_CODES: break return response From 420447c1d1bc71f29dc765c391f4bfb5ce6008bf Mon Sep 17 00:00:00 2001 From: agrawalradhika-cell Date: Tue, 25 Aug 2026 20:58:25 -0700 Subject: [PATCH 02/79] chore: Add tests for MTLS certificate rotation behavior --- .../tests/transport/aio/test_sessions_mtls.py | 76 +++++++++++++++++++ 1 file changed, 76 insertions(+) diff --git a/packages/google-auth/tests/transport/aio/test_sessions_mtls.py b/packages/google-auth/tests/transport/aio/test_sessions_mtls.py index b68766ca5b5d..740c6ac84ba2 100644 --- a/packages/google-auth/tests/transport/aio/test_sessions_mtls.py +++ b/packages/google-auth/tests/transport/aio/test_sessions_mtls.py @@ -344,3 +344,79 @@ async def test_configure_mtls_channel_close_exception_does_not_abort(self): assert session._is_mtls is True assert session._cached_cert == b"fake_cert_data" await session.close() + + @pytest.mark.asyncio + async def test_cert_rotation_failure_raises_error(self): + mock_creds = mock.AsyncMock(spec=credentials.Credentials) + mock_creds.before_request = mock.AsyncMock(return_value=None) + + mock_auth_req = mock.AsyncMock() + mock_resp = mock.Mock() + import http.client as http_client + mock_resp.status_code = http_client.UNAUTHORIZED + mock_auth_req.return_value = mock_resp + + session = sessions.AsyncAuthorizedSession(mock_creds, auth_request=mock_auth_req) + session._is_mtls = True + session._cached_cert = b"old_cert" + + new_cert = b"new_cert" + new_key = b"new_key" + + with mock.patch("google.auth.transport._mtls_helper.check_parameters_for_unauthorized_response") as mock_check, \ + mock.patch.object(session, "configure_mtls_channel", new_callable=mock.AsyncMock) as mock_conf: + + mock_check.return_value = (new_cert, new_key, b"old_fp", b"new_fp") + mock_conf.side_effect = Exception("Failed to reconfigure") + + with pytest.raises(exceptions.MutualTLSChannelError): + await session.request("GET", "http://example.com") + + mock_check.assert_called_once() + mock_conf.assert_called_once() + + @pytest.mark.asyncio + async def test_cert_rotation_check_params_fails(self): + mock_creds = mock.AsyncMock(spec=credentials.Credentials) + mock_auth_req = mock.AsyncMock() + mock_resp = mock.Mock() + import http.client as http_client + mock_resp.status_code = http_client.UNAUTHORIZED + mock_auth_req.return_value = mock_resp + + session = sessions.AsyncAuthorizedSession(mock_creds, auth_request=mock_auth_req) + session._is_mtls = True + session._cached_cert = b"cached_cert" + + with mock.patch( + "google.auth.transport._mtls_helper.check_parameters_for_unauthorized_response", + side_effect=Exception("check_params failed"), + ) as mock_check_params: + with pytest.raises(Exception, match="check_params failed"): + await session.request("GET", "http://example.com") + + mock_check_params.assert_called_once() + + @pytest.mark.asyncio + async def test_no_cert_rotation_when_cert_match_and_mTLS_enabled(self): + mock_creds = mock.AsyncMock(spec=credentials.Credentials) + mock_auth_req = mock.AsyncMock() + mock_resp = mock.Mock() + import http.client as http_client + mock_resp.status_code = http_client.UNAUTHORIZED + mock_auth_req.return_value = mock_resp + + session = sessions.AsyncAuthorizedSession(mock_creds, auth_request=mock_auth_req) + session._is_mtls = True + session._cached_cert = b"old_cert" + + with mock.patch( + "google.auth.transport._mtls_helper.check_parameters_for_unauthorized_response", + ) as mock_check, mock.patch.object(session, "configure_mtls_channel", new_callable=mock.AsyncMock) as mock_conf: + # same fingerprint, so no call to configure_mtls_channel + mock_check.return_value = (b"new_cert", b"new_key", b"same_fp", b"same_fp") + + await session.request("GET", "http://example.com") + + mock_check.assert_called_once() + mock_conf.assert_not_called() From 907cf0083a8120259ab10b2b7a503c8e8ea139b1 Mon Sep 17 00:00:00 2001 From: agrawalradhika-cell Date: Wed, 26 Aug 2026 10:34:10 -0700 Subject: [PATCH 03/79] Update packages/google-auth/tests/transport/aio/test_sessions_mtls.py Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- .../google-auth/tests/transport/aio/test_sessions_mtls.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/packages/google-auth/tests/transport/aio/test_sessions_mtls.py b/packages/google-auth/tests/transport/aio/test_sessions_mtls.py index 740c6ac84ba2..f448c57073b2 100644 --- a/packages/google-auth/tests/transport/aio/test_sessions_mtls.py +++ b/packages/google-auth/tests/transport/aio/test_sessions_mtls.py @@ -392,9 +392,8 @@ async def test_cert_rotation_check_params_fails(self): "google.auth.transport._mtls_helper.check_parameters_for_unauthorized_response", side_effect=Exception("check_params failed"), ) as mock_check_params: - with pytest.raises(Exception, match="check_params failed"): - await session.request("GET", "http://example.com") - + resp = await session.request("GET", "http://example.com") + assert resp == mock_resp mock_check_params.assert_called_once() @pytest.mark.asyncio From cc850b13e16e2ed1b297a00d4e2ecd958eaa5b6c Mon Sep 17 00:00:00 2001 From: agrawalradhika-cell Date: Wed, 26 Aug 2026 10:38:02 -0700 Subject: [PATCH 04/79] Improve error handling for mTLS reconfiguration Handle exceptions during mTLS reconfiguration with warnings instead of errors. --- .../google/auth/aio/transport/sessions.py | 17 +++++++++++------ 1 file changed, 11 insertions(+), 6 deletions(-) diff --git a/packages/google-auth/google/auth/aio/transport/sessions.py b/packages/google-auth/google/auth/aio/transport/sessions.py index c55be4670549..6ace722c2814 100644 --- a/packages/google-auth/google/auth/aio/transport/sessions.py +++ b/packages/google-auth/google/auth/aio/transport/sessions.py @@ -315,7 +315,7 @@ async def request( ) ) if response.status_code == http_client.UNAUTHORIZED: - if self.is_mtls: + try: call_cert_bytes, call_key_bytes, cached_fingerprint, current_cert_fingerprint = google.auth.transport._mtls_helper.check_parameters_for_unauthorized_response( self._cached_cert ) @@ -330,15 +330,20 @@ async def request( ) continue except Exception as e: - _LOGGER.error("Failed to reconfigure mTLS channel: %s", e) - raise exceptions.MutualTLSChannelError( - "Failed to reconfigure mTLS channel" - ) from e + _LOGGER.warning( + "Failed to reconfigure mTLS channel: %s. Proceeding with original response.", + e + ) else: - _LOGGER.info( + _LOGGER.info( "Skipping reconfiguration of mTLS channel because the client" " certificate has not changed." ) + except Exception as e: + _LOGGER.warning( + "Failed to check client certificate parameters: %s. Proceeding with original response.", + e, + ) if response.status_code not in transport.DEFAULT_RETRYABLE_STATUS_CODES: break From a44acb0654b0ec21cbe49e98aa8ad1c638001d38 Mon Sep 17 00:00:00 2001 From: agrawalradhika-cell Date: Wed, 26 Aug 2026 11:19:57 -0700 Subject: [PATCH 05/79] fix: Rename test_cert_rotation_failure to test_cert_rotation_failure_logs Updated test logic to assert response instead of expecting an error. --- .../google-auth/tests/transport/aio/test_sessions_mtls.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/packages/google-auth/tests/transport/aio/test_sessions_mtls.py b/packages/google-auth/tests/transport/aio/test_sessions_mtls.py index f448c57073b2..0b0a0ec87229 100644 --- a/packages/google-auth/tests/transport/aio/test_sessions_mtls.py +++ b/packages/google-auth/tests/transport/aio/test_sessions_mtls.py @@ -346,7 +346,7 @@ async def test_configure_mtls_channel_close_exception_does_not_abort(self): await session.close() @pytest.mark.asyncio - async def test_cert_rotation_failure_raises_error(self): + async def test_cert_rotation_failure_logs(self): mock_creds = mock.AsyncMock(spec=credentials.Credentials) mock_creds.before_request = mock.AsyncMock(return_value=None) @@ -369,8 +369,8 @@ async def test_cert_rotation_failure_raises_error(self): mock_check.return_value = (new_cert, new_key, b"old_fp", b"new_fp") mock_conf.side_effect = Exception("Failed to reconfigure") - with pytest.raises(exceptions.MutualTLSChannelError): - await session.request("GET", "http://example.com") + resp = await session.request("GET", "http://example.com") + assert resp == mock_resp mock_check.assert_called_once() mock_conf.assert_called_once() From 984e47c3388abd43271ac9e0d7c6f198534db176 Mon Sep 17 00:00:00 2001 From: agrawalradhika-cell Date: Wed, 26 Aug 2026 14:54:48 -0700 Subject: [PATCH 06/79] chore: Refactor MTLS parameter check on unauthorized response o use async executor Refactor unauthorized response handling to use async executor for MTLS parameter checks. --- .../google-auth/google/auth/aio/transport/sessions.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/packages/google-auth/google/auth/aio/transport/sessions.py b/packages/google-auth/google/auth/aio/transport/sessions.py index 6ace722c2814..1a4b9c468f9a 100644 --- a/packages/google-auth/google/auth/aio/transport/sessions.py +++ b/packages/google-auth/google/auth/aio/transport/sessions.py @@ -316,7 +316,13 @@ async def request( ) if response.status_code == http_client.UNAUTHORIZED: try: - call_cert_bytes, call_key_bytes, cached_fingerprint, current_cert_fingerprint = google.auth.transport._mtls_helper.check_parameters_for_unauthorized_response( + ( + call_cert_bytes, + call_key_bytes, + cached_fingerprint, + current_cert_fingerprint + ) = await mtls._run_in_executor( + google.auth.transport._mtls_helper.check_parameters_for_unauthorized_response, self._cached_cert ) if cached_fingerprint != current_cert_fingerprint: From 30341bc5acfbc6155165c8062852e670acd7dedb Mon Sep 17 00:00:00 2001 From: agrawalradhika-cell Date: Wed, 26 Aug 2026 22:03:28 -0700 Subject: [PATCH 07/79] chore: Reset mTLS init task upon client certificate change chore: Reset mTLS init task upon client certificate change --- packages/google-auth/google/auth/aio/transport/sessions.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/packages/google-auth/google/auth/aio/transport/sessions.py b/packages/google-auth/google/auth/aio/transport/sessions.py index 1a4b9c468f9a..c4ef00528d0a 100644 --- a/packages/google-auth/google/auth/aio/transport/sessions.py +++ b/packages/google-auth/google/auth/aio/transport/sessions.py @@ -331,6 +331,8 @@ async def request( "Client certificate has changed, reconfiguring mTLS " "channel." ) + if self._mtls_init_task and self._mtls_init_task.done(): + self._mtls_init_task = None await self.configure_mtls_channel( lambda: (call_cert_bytes, call_key_bytes) ) From 1c068dce457d288fa3a98a3534158cc3a5b53656 Mon Sep 17 00:00:00 2001 From: Radhika Agrawal Date: Thu, 27 Aug 2026 17:48:12 +0000 Subject: [PATCH 08/79] fix: fix the lint errors Signed-off-by: Radhika Agrawal --- .../google/auth/aio/transport/sessions.py | 18 ++++++------- .../tests/transport/aio/test_sessions_mtls.py | 27 ++++++++++++++----- 2 files changed, 29 insertions(+), 16 deletions(-) diff --git a/packages/google-auth/google/auth/aio/transport/sessions.py b/packages/google-auth/google/auth/aio/transport/sessions.py index c4ef00528d0a..22e11c74f577 100644 --- a/packages/google-auth/google/auth/aio/transport/sessions.py +++ b/packages/google-auth/google/auth/aio/transport/sessions.py @@ -28,8 +28,6 @@ from google.auth.exceptions import TimeoutError import google.auth.transport._mtls_helper -_LOGGER = logging.getLogger(__name__) - if TYPE_CHECKING: # pragma: NO COVER import aiohttp from aiohttp import ClientTimeout # type: ignore @@ -41,6 +39,8 @@ except (ImportError, AttributeError): ClientTimeout = None +_LOGGER = logging.getLogger(__name__) + # Tracks the internal aiohttp installation and usage try: @@ -317,13 +317,13 @@ async def request( if response.status_code == http_client.UNAUTHORIZED: try: ( - call_cert_bytes, - call_key_bytes, - cached_fingerprint, - current_cert_fingerprint + call_cert_bytes, + call_key_bytes, + cached_fingerprint, + current_cert_fingerprint, ) = await mtls._run_in_executor( google.auth.transport._mtls_helper.check_parameters_for_unauthorized_response, - self._cached_cert + self._cached_cert, ) if cached_fingerprint != current_cert_fingerprint: try: @@ -340,10 +340,10 @@ async def request( except Exception as e: _LOGGER.warning( "Failed to reconfigure mTLS channel: %s. Proceeding with original response.", - e + e, ) else: - _LOGGER.info( + _LOGGER.info( "Skipping reconfiguration of mTLS channel because the client" " certificate has not changed." ) diff --git a/packages/google-auth/tests/transport/aio/test_sessions_mtls.py b/packages/google-auth/tests/transport/aio/test_sessions_mtls.py index 0b0a0ec87229..764bfc8c40e3 100644 --- a/packages/google-auth/tests/transport/aio/test_sessions_mtls.py +++ b/packages/google-auth/tests/transport/aio/test_sessions_mtls.py @@ -353,19 +353,24 @@ async def test_cert_rotation_failure_logs(self): mock_auth_req = mock.AsyncMock() mock_resp = mock.Mock() import http.client as http_client + mock_resp.status_code = http_client.UNAUTHORIZED mock_auth_req.return_value = mock_resp - session = sessions.AsyncAuthorizedSession(mock_creds, auth_request=mock_auth_req) + session = sessions.AsyncAuthorizedSession( + mock_creds, auth_request=mock_auth_req + ) session._is_mtls = True session._cached_cert = b"old_cert" new_cert = b"new_cert" new_key = b"new_key" - with mock.patch("google.auth.transport._mtls_helper.check_parameters_for_unauthorized_response") as mock_check, \ - mock.patch.object(session, "configure_mtls_channel", new_callable=mock.AsyncMock) as mock_conf: - + with mock.patch( + "google.auth.transport._mtls_helper.check_parameters_for_unauthorized_response" + ) as mock_check, mock.patch.object( + session, "configure_mtls_channel", new_callable=mock.AsyncMock + ) as mock_conf: mock_check.return_value = (new_cert, new_key, b"old_fp", b"new_fp") mock_conf.side_effect = Exception("Failed to reconfigure") @@ -381,10 +386,13 @@ async def test_cert_rotation_check_params_fails(self): mock_auth_req = mock.AsyncMock() mock_resp = mock.Mock() import http.client as http_client + mock_resp.status_code = http_client.UNAUTHORIZED mock_auth_req.return_value = mock_resp - session = sessions.AsyncAuthorizedSession(mock_creds, auth_request=mock_auth_req) + session = sessions.AsyncAuthorizedSession( + mock_creds, auth_request=mock_auth_req + ) session._is_mtls = True session._cached_cert = b"cached_cert" @@ -402,16 +410,21 @@ async def test_no_cert_rotation_when_cert_match_and_mTLS_enabled(self): mock_auth_req = mock.AsyncMock() mock_resp = mock.Mock() import http.client as http_client + mock_resp.status_code = http_client.UNAUTHORIZED mock_auth_req.return_value = mock_resp - session = sessions.AsyncAuthorizedSession(mock_creds, auth_request=mock_auth_req) + session = sessions.AsyncAuthorizedSession( + mock_creds, auth_request=mock_auth_req + ) session._is_mtls = True session._cached_cert = b"old_cert" with mock.patch( "google.auth.transport._mtls_helper.check_parameters_for_unauthorized_response", - ) as mock_check, mock.patch.object(session, "configure_mtls_channel", new_callable=mock.AsyncMock) as mock_conf: + ) as mock_check, mock.patch.object( + session, "configure_mtls_channel", new_callable=mock.AsyncMock + ) as mock_conf: # same fingerprint, so no call to configure_mtls_channel mock_check.return_value = (b"new_cert", b"new_key", b"same_fp", b"same_fp") From 6fb1e863826bfc34c76f0d3bfbbf5610766fd07f Mon Sep 17 00:00:00 2001 From: agrawalradhika-cell Date: Thu, 27 Aug 2026 13:00:32 -0700 Subject: [PATCH 09/79] chore: Refactor mTLS channel reconfiguration logic for adding mTLS check after 401 check chore: Refactor mTLS channel reconfiguration logic for adding mTLS check after 401 check --- .../google/auth/aio/transport/sessions.py | 72 ++++++++++--------- 1 file changed, 38 insertions(+), 34 deletions(-) diff --git a/packages/google-auth/google/auth/aio/transport/sessions.py b/packages/google-auth/google/auth/aio/transport/sessions.py index 22e11c74f577..29d22eaa0e1d 100644 --- a/packages/google-auth/google/auth/aio/transport/sessions.py +++ b/packages/google-auth/google/auth/aio/transport/sessions.py @@ -40,6 +40,7 @@ ClientTimeout = None _LOGGER = logging.getLogger(__name__) +MTLS_URL_PREFIXES = ["mtls.googleapis.com", "mtls.sandbox.googleapis.com"] # Tracks the internal aiohttp installation and usage @@ -315,43 +316,46 @@ async def request( ) ) if response.status_code == http_client.UNAUTHORIZED: - try: - ( - call_cert_bytes, - call_key_bytes, - cached_fingerprint, - current_cert_fingerprint, - ) = await mtls._run_in_executor( - google.auth.transport._mtls_helper.check_parameters_for_unauthorized_response, - self._cached_cert, - ) - if cached_fingerprint != current_cert_fingerprint: - try: + if getattr(self, "is_mtls", False) and any( + prefix in url for prefix in MTLS_URL_PREFIXES + ): + try: + ( + call_cert_bytes, + call_key_bytes, + cached_fingerprint, + current_cert_fingerprint, + ) = await mtls._run_in_executor( + google.auth.transport._mtls_helper.check_parameters_for_unauthorized_response, + self._cached_cert, + ) + if cached_fingerprint != current_cert_fingerprint: + try: + _LOGGER.info( + "Client certificate has changed, reconfiguring mTLS " + "channel." + ) + if self._mtls_init_task and self._mtls_init_task.done(): + self._mtls_init_task = None + await self.configure_mtls_channel( + lambda: (call_cert_bytes, call_key_bytes) + ) + continue + except Exception as e: + _LOGGER.warning( + "Failed to reconfigure mTLS channel: %s. Proceeding with original response.", + e, + ) + else: _LOGGER.info( - "Client certificate has changed, reconfiguring mTLS " - "channel." - ) - if self._mtls_init_task and self._mtls_init_task.done(): - self._mtls_init_task = None - await self.configure_mtls_channel( - lambda: (call_cert_bytes, call_key_bytes) + "Skipping reconfiguration of mTLS channel because the client" + " certificate has not changed." ) - continue - except Exception as e: - _LOGGER.warning( - "Failed to reconfigure mTLS channel: %s. Proceeding with original response.", - e, - ) - else: - _LOGGER.info( - "Skipping reconfiguration of mTLS channel because the client" - " certificate has not changed." + except Exception as e: + _LOGGER.warning( + "Failed to check client certificate parameters: %s. Proceeding with original response.", + e, ) - except Exception as e: - _LOGGER.warning( - "Failed to check client certificate parameters: %s. Proceeding with original response.", - e, - ) if response.status_code not in transport.DEFAULT_RETRYABLE_STATUS_CODES: break From 2cdfe2d5197dcf29a1ffa613ff993169107a8a70 Mon Sep 17 00:00:00 2001 From: agrawalradhika-cell Date: Thu, 27 Aug 2026 16:05:18 -0700 Subject: [PATCH 10/79] chore: Add mTLS rotation lock for certificate management Implement mTLS rotation lock to prevent race conditions during certificate reconfiguration. --- .../google/auth/aio/transport/sessions.py | 74 +++++++++++-------- 1 file changed, 43 insertions(+), 31 deletions(-) diff --git a/packages/google-auth/google/auth/aio/transport/sessions.py b/packages/google-auth/google/auth/aio/transport/sessions.py index 29d22eaa0e1d..b9bc251053b4 100644 --- a/packages/google-auth/google/auth/aio/transport/sessions.py +++ b/packages/google-auth/google/auth/aio/transport/sessions.py @@ -153,6 +153,7 @@ def __init__( "`auth_request` must either be configured or the external package `aiohttp` must be installed to use the default value." ) self._auth_request = _auth_request + self._mtls_rotation_lock = asyncio.Lock() async def configure_mtls_channel(self, client_cert_callback=None): """Configure the client certificate and key for SSL connection. @@ -319,43 +320,54 @@ async def request( if getattr(self, "is_mtls", False) and any( prefix in url for prefix in MTLS_URL_PREFIXES ): - try: - ( - call_cert_bytes, - call_key_bytes, - cached_fingerprint, - current_cert_fingerprint, - ) = await mtls._run_in_executor( - google.auth.transport._mtls_helper.check_parameters_for_unauthorized_response, - self._cached_cert, - ) - if cached_fingerprint != current_cert_fingerprint: + # Snapshot the stale certificate state BEFORE acquiring the lock. + # This represents the cert that caused the 401 rejection. + stale_cert = self._cached_cert + + # Wait in line to acquire the lock + async with self._mtls_rotation_lock: + # Check Did another coroutine already reconfigure mTLS + if self._cached_cert != stale_cert: + # Yes! Another request already updated the channel + pass + else: try: - _LOGGER.info( - "Client certificate has changed, reconfiguring mTLS " - "channel." + ( + call_cert_bytes, + call_key_bytes, + cached_fingerprint, + current_cert_fingerprint, + ) = await mtls._run_in_executor( + google.auth.transport._mtls_helper.check_parameters_for_unauthorized_response, + self._cached_cert, ) - if self._mtls_init_task and self._mtls_init_task.done(): - self._mtls_init_task = None - await self.configure_mtls_channel( - lambda: (call_cert_bytes, call_key_bytes) - ) - continue + if cached_fingerprint != current_cert_fingerprint: + try: + _LOGGER.info( + "Client certificate has changed, reconfiguring mTLS " + "channel." + ) + if self._mtls_init_task and self._mtls_init_task.done(): + self._mtls_init_task = None + await self.configure_mtls_channel( + lambda: (call_cert_bytes, call_key_bytes) + ) + continue + except Exception as e: + _LOGGER.warning( + "Failed to reconfigure mTLS channel: %s. Proceeding with original response.", + e, + ) + else: + _LOGGER.info( + "Skipping reconfiguration of mTLS channel because the client" + " certificate has not changed." + ) except Exception as e: _LOGGER.warning( - "Failed to reconfigure mTLS channel: %s. Proceeding with original response.", + "Failed to check client certificate parameters: %s. Proceeding with original response.", e, ) - else: - _LOGGER.info( - "Skipping reconfiguration of mTLS channel because the client" - " certificate has not changed." - ) - except Exception as e: - _LOGGER.warning( - "Failed to check client certificate parameters: %s. Proceeding with original response.", - e, - ) if response.status_code not in transport.DEFAULT_RETRYABLE_STATUS_CODES: break From d734731ca21a37b137ff2c0c65c9fd54d4e55458 Mon Sep 17 00:00:00 2001 From: agrawalradhika-cell Date: Thu, 27 Aug 2026 16:21:53 -0700 Subject: [PATCH 11/79] chore: Log mTLS channel reconfiguration failure as error chore: Change warning to error log for mTLS channel reconfiguration failure. --- .../google-auth/google/auth/aio/transport/sessions.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/packages/google-auth/google/auth/aio/transport/sessions.py b/packages/google-auth/google/auth/aio/transport/sessions.py index b9bc251053b4..6d7c8d5199ea 100644 --- a/packages/google-auth/google/auth/aio/transport/sessions.py +++ b/packages/google-auth/google/auth/aio/transport/sessions.py @@ -283,6 +283,8 @@ async def request( google.auth.exceptions.TimeoutError: If the method does not complete within the configured `max_allowed_time` or the request exceeds the configured `timeout`. + google.auth.exceptions.MutualTLSChannelError: If mutual TLS + channel reconfiguration fails for any reason during certificate rotation. """ if self._mtls_init_task: try: @@ -354,10 +356,10 @@ async def request( ) continue except Exception as e: - _LOGGER.warning( - "Failed to reconfigure mTLS channel: %s. Proceeding with original response.", - e, - ) + _LOGGER.error("Failed to reconfigure mTLS channel: %s", e) + raise exceptions.MutualTLSChannelError( + "Failed to reconfigure mTLS channel" + ) from e else: _LOGGER.info( "Skipping reconfiguration of mTLS channel because the client" From 97e91d0e8a7fc35cfe6f05b9252cf50ec77a8aeb Mon Sep 17 00:00:00 2001 From: agrawalradhika-cell Date: Fri, 28 Aug 2026 11:16:03 -0700 Subject: [PATCH 12/79] chore: Refactor mTLS handling for unauthorized responses chore: Refactor mTLS handling for unauthorized responses --- .../google/auth/aio/transport/sessions.py | 127 +++++++++++------- 1 file changed, 75 insertions(+), 52 deletions(-) diff --git a/packages/google-auth/google/auth/aio/transport/sessions.py b/packages/google-auth/google/auth/aio/transport/sessions.py index 6d7c8d5199ea..df73569c7465 100644 --- a/packages/google-auth/google/auth/aio/transport/sessions.py +++ b/packages/google-auth/google/auth/aio/transport/sessions.py @@ -318,61 +318,84 @@ async def request( url, method, data, headers, actual_timeout, **kwargs ) ) - if response.status_code == http_client.UNAUTHORIZED: - if getattr(self, "is_mtls", False) and any( - prefix in url for prefix in MTLS_URL_PREFIXES - ): - # Snapshot the stale certificate state BEFORE acquiring the lock. - # This represents the cert that caused the 401 rejection. - stale_cert = self._cached_cert - - # Wait in line to acquire the lock - async with self._mtls_rotation_lock: - # Check Did another coroutine already reconfigure mTLS - if self._cached_cert != stale_cert: - # Yes! Another request already updated the channel - pass - else: - try: - ( - call_cert_bytes, - call_key_bytes, - cached_fingerprint, - current_cert_fingerprint, - ) = await mtls._run_in_executor( - google.auth.transport._mtls_helper.check_parameters_for_unauthorized_response, - self._cached_cert, - ) - if cached_fingerprint != current_cert_fingerprint: - try: - _LOGGER.info( - "Client certificate has changed, reconfiguring mTLS " - "channel." - ) - if self._mtls_init_task and self._mtls_init_task.done(): - self._mtls_init_task = None - await self.configure_mtls_channel( - lambda: (call_cert_bytes, call_key_bytes) - ) - continue - except Exception as e: - _LOGGER.error("Failed to reconfigure mTLS channel: %s", e) - raise exceptions.MutualTLSChannelError( - "Failed to reconfigure mTLS channel" - ) from e - else: - _LOGGER.info( - "Skipping reconfiguration of mTLS channel because the client" - " certificate has not changed." - ) - except Exception as e: - _LOGGER.warning( - "Failed to check client certificate parameters: %s. Proceeding with original response.", - e, - ) if response.status_code not in transport.DEFAULT_RETRYABLE_STATUS_CODES: break + + if response.status_code == http_client.UNAUTHORIZED: + _auth_retry_count = kwargs.pop("_auth_retry_count", 0) + if _auth_retry_count < 2: + is_streaming = data is not None and isinstance(data, (collections.abc.Iterator, collections.abc.AsyncIterable)) or hasattr(data, "read") + if getattr(self, "is_mtls", False) and any( + prefix in url for prefix in MTLS_URL_PREFIXES + ): + # Snapshot the stale certificate state BEFORE acquiring the lock. + # This represents the cert that caused the 401 rejection. + stale_cert = self._cached_cert + + # Wait in line to acquire the lock + async with self._mtls_rotation_lock: + # Check Did another coroutine already reconfigure mTLS + if self._cached_cert != stale_cert: + # Yes! Another request already updated the channel + pass + else: + try: + ( + call_cert_bytes, + call_key_bytes, + cached_fingerprint, + current_cert_fingerprint, + ) = await mtls._run_in_executor( + google.auth.transport._mtls_helper.check_parameters_for_unauthorized_response, + self._cached_cert, + ) + if cached_fingerprint != current_cert_fingerprint: + try: + _LOGGER.info( + "Client certificate has changed, reconfiguring mTLS " + "channel." + ) + if self._mtls_init_task and self._mtls_init_task.done(): + self._mtls_init_task = None + await self.configure_mtls_channel( + lambda: (call_cert_bytes, call_key_bytes) + ) + continue + except Exception as e: + _LOGGER.error("Failed to reconfigure mTLS channel: %s", e) + raise exceptions.MutualTLSChannelError( + "Failed to reconfigure mTLS channel" + ) from e + else: + _LOGGER.info( + "Skipping reconfiguration of mTLS channel because the client" + " certificate has not changed." + ) + except Exception as e: + _LOGGER.warning( + "Failed to check client certificate parameters: %s. Proceeding with original response.", + e, + ) + if is_streaming: + return response + if hasattr(response, "close"): + if asyncio.iscoroutinefunction(response.close): + await response.close() + else: + response.close() + await self._credentials.refresh(self._auth_request) + kwargs["_auth_retry_count"] = _auth_retry_count + 1 + return await self.request( + method, + url, + data=data, + headers=headers, + max_allowed_time=max_allowed_time, + timeout=timeout, + total_attempts=total_attempts, + **kwargs + ) return response @functools.wraps(request) From d0da58b56a40e78fc05961607ba5a71fc318548f Mon Sep 17 00:00:00 2001 From: agrawalradhika-cell Date: Fri, 28 Aug 2026 11:19:45 -0700 Subject: [PATCH 13/79] fix: Remove unnecessary continue statement after mTLS configuration. Remove unnecessary continue statement after mTLS configuration. --- packages/google-auth/google/auth/aio/transport/sessions.py | 1 - 1 file changed, 1 deletion(-) diff --git a/packages/google-auth/google/auth/aio/transport/sessions.py b/packages/google-auth/google/auth/aio/transport/sessions.py index df73569c7465..2be1dbe00147 100644 --- a/packages/google-auth/google/auth/aio/transport/sessions.py +++ b/packages/google-auth/google/auth/aio/transport/sessions.py @@ -361,7 +361,6 @@ async def request( await self.configure_mtls_channel( lambda: (call_cert_bytes, call_key_bytes) ) - continue except Exception as e: _LOGGER.error("Failed to reconfigure mTLS channel: %s", e) raise exceptions.MutualTLSChannelError( From 825426d04ad33c0a26aae7de280f78191bbc4be8 Mon Sep 17 00:00:00 2001 From: agrawalradhika-cell Date: Fri, 28 Aug 2026 11:47:13 -0700 Subject: [PATCH 14/79] fix: Fix cert rotation tests and improve error handling Refactor tests for certificate rotation and error handling in AsyncAuthorizedSession. Update test names for clarity and ensure proper logging of errors. --- .../tests/transport/aio/test_sessions_mtls.py | 166 +++++++++++++----- 1 file changed, 120 insertions(+), 46 deletions(-) diff --git a/packages/google-auth/tests/transport/aio/test_sessions_mtls.py b/packages/google-auth/tests/transport/aio/test_sessions_mtls.py index 764bfc8c40e3..c3110d513818 100644 --- a/packages/google-auth/tests/transport/aio/test_sessions_mtls.py +++ b/packages/google-auth/tests/transport/aio/test_sessions_mtls.py @@ -15,6 +15,7 @@ import json import os import ssl +import http.client as http_client from unittest import mock import pytest @@ -346,89 +347,162 @@ async def test_configure_mtls_channel_close_exception_does_not_abort(self): await session.close() @pytest.mark.asyncio - async def test_cert_rotation_failure_logs(self): + async def test_cert_rotation_failure_raises_error(self, caplog): mock_creds = mock.AsyncMock(spec=credentials.Credentials) mock_creds.before_request = mock.AsyncMock(return_value=None) - - mock_auth_req = mock.AsyncMock() + mock_resp = mock.Mock() - import http.client as http_client - mock_resp.status_code = http_client.UNAUTHORIZED - mock_auth_req.return_value = mock_resp + mock_auth_req = mock.AsyncMock(return_value=mock_resp) - session = sessions.AsyncAuthorizedSession( - mock_creds, auth_request=mock_auth_req - ) + session = sessions.AsyncAuthorizedSession(mock_creds, auth_request=mock_auth_req) session._is_mtls = True session._cached_cert = b"old_cert" new_cert = b"new_cert" new_key = b"new_key" - with mock.patch( - "google.auth.transport._mtls_helper.check_parameters_for_unauthorized_response" - ) as mock_check, mock.patch.object( - session, "configure_mtls_channel", new_callable=mock.AsyncMock - ) as mock_conf: + with mock.patch("google.auth.transport._mtls_helper.check_parameters_for_unauthorized_response") as mock_check, \ + mock.patch.object(session, "configure_mtls_channel", new_callable=mock.AsyncMock) as mock_conf: + mock_check.return_value = (new_cert, new_key, b"old_fp", b"new_fp") mock_conf.side_effect = Exception("Failed to reconfigure") - resp = await session.request("GET", "http://example.com") - assert resp == mock_resp + with pytest.raises(exceptions.MutualTLSChannelError): + await session.request("GET", "https://pubsub.mtls.googleapis.com/test") mock_check.assert_called_once() mock_conf.assert_called_once() + assert "Failed to reconfigure mTLS channel" in caplog.text + + await session.close() + @pytest.mark.asyncio - async def test_cert_rotation_check_params_fails(self): + async def test_cert_rotation_check_params_fails(self, caplog): mock_creds = mock.AsyncMock(spec=credentials.Credentials) - mock_auth_req = mock.AsyncMock() + mock_creds.before_request = mock.AsyncMock(return_value=None) + mock_resp = mock.Mock() - import http.client as http_client - mock_resp.status_code = http_client.UNAUTHORIZED - mock_auth_req.return_value = mock_resp + mock_auth_req = mock.AsyncMock(return_value=mock_resp) - session = sessions.AsyncAuthorizedSession( - mock_creds, auth_request=mock_auth_req - ) + session = sessions.AsyncAuthorizedSession(mock_creds, auth_request=mock_auth_req) session._is_mtls = True - session._cached_cert = b"cached_cert" + session._cached_cert = b"old_cert" + + with mock.patch("google.auth.transport._mtls_helper.check_parameters_for_unauthorized_response") as mock_check, \ + mock.patch.object(session, "configure_mtls_channel", new_callable=mock.AsyncMock) as mock_conf: + + mock_check.side_effect = Exception("Failed to check params") + + resp = await session.request("GET", "https://pubsub.mtls.googleapis.com/test") - with mock.patch( - "google.auth.transport._mtls_helper.check_parameters_for_unauthorized_response", - side_effect=Exception("check_params failed"), - ) as mock_check_params: - resp = await session.request("GET", "http://example.com") assert resp == mock_resp - mock_check_params.assert_called_once() + mock_check.assert_called_once() + mock_conf.assert_not_called() + assert "Failed to check client certificate parameters" in caplog.text + + await session.close() + @pytest.mark.asyncio async def test_no_cert_rotation_when_cert_match_and_mTLS_enabled(self): mock_creds = mock.AsyncMock(spec=credentials.Credentials) - mock_auth_req = mock.AsyncMock() + mock_creds.before_request = mock.AsyncMock(return_value=None) + mock_resp = mock.Mock() - import http.client as http_client - mock_resp.status_code = http_client.UNAUTHORIZED - mock_auth_req.return_value = mock_resp + mock_auth_req = mock.AsyncMock(return_value=mock_resp) - session = sessions.AsyncAuthorizedSession( - mock_creds, auth_request=mock_auth_req - ) + session = sessions.AsyncAuthorizedSession(mock_creds, auth_request=mock_auth_req) session._is_mtls = True session._cached_cert = b"old_cert" - with mock.patch( - "google.auth.transport._mtls_helper.check_parameters_for_unauthorized_response", - ) as mock_check, mock.patch.object( - session, "configure_mtls_channel", new_callable=mock.AsyncMock - ) as mock_conf: - # same fingerprint, so no call to configure_mtls_channel - mock_check.return_value = (b"new_cert", b"new_key", b"same_fp", b"same_fp") + new_cert = b"new_cert" + new_key = b"new_key" + + with mock.patch("google.auth.transport._mtls_helper.check_parameters_for_unauthorized_response") as mock_check, \ + mock.patch.object(session, "configure_mtls_channel", new_callable=mock.AsyncMock) as mock_conf: - await session.request("GET", "http://example.com") + # Matching fingerprints mean no layout rotation is needed + mock_check.return_value = (new_cert, new_key, b"old_fp", b"old_fp") + resp = await session.request("GET", "https://pubsub.mtls.googleapis.com/test") + + assert resp == mock_resp mock_check.assert_called_once() mock_conf.assert_not_called() + + await session.close() + + + @pytest.mark.asyncio + async def test_cert_rotation_success_and_retry(self): + mock_creds = mock.AsyncMock(spec=credentials.Credentials) + mock_creds.before_request = mock.AsyncMock(return_value=None) + mock_creds.refresh = mock.AsyncMock(return_value=None) + + # Initial request fails natively with 401. Retry succeeds with 200. + mock_resp_401 = mock.Mock() + mock_resp_401.status_code = http_client.UNAUTHORIZED + mock_resp_200 = mock.Mock() + mock_resp_200.status_code = http_client.OK + + # Use side_effect to dynamically yield responses + mock_auth_req = mock.AsyncMock(side_effect=[mock_resp_401, mock_resp_200]) + + session = sessions.AsyncAuthorizedSession(mock_creds, auth_request=mock_auth_req) + session._is_mtls = True + session._cached_cert = b"old_cert" + + new_cert = b"new_cert" + new_key = b"new_key" + + with mock.patch("google.auth.transport._mtls_helper.check_parameters_for_unauthorized_response") as mock_check, \ + mock.patch.object(session, "configure_mtls_channel", new_callable=mock.AsyncMock) as mock_conf: + + mock_check.return_value = (new_cert, new_key, b"old_fp", b"new_fp") + + resp = await session.request("GET", "https://pubsub.mtls.googleapis.com/test") + + # 1. Assert the retried 200 response is successfully returned to the user + assert resp == mock_resp_200 + + # 2. Assert rotation logic correctly executed + mock_check.assert_called_once() + mock_conf.assert_called_once_with(mock.ANY) + + # 3. Assert credentials were explicitly refreshed + mock_creds.refresh.assert_called_once() + + # 4. Assert headers were explicitly rebound on the recursive retry (2 invocations) + assert mock_creds.before_request.call_count == 2 + + await session.close() + + + @pytest.mark.asyncio + async def test_non_mtls_url_bypasses_rotation(self): + mock_creds = mock.AsyncMock(spec=credentials.Credentials) + mock_resp_401 = mock.Mock() + mock_resp_401.status_code = http_client.UNAUTHORIZED + mock_auth_req = mock.AsyncMock(return_value=mock_resp_401) + + session = sessions.AsyncAuthorizedSession(mock_creds, auth_request=mock_auth_req) + + # Even if mTLS is enabled globally... + session._is_mtls = True + session._cached_cert = b"old_cert" + + with mock.patch("google.auth.transport._mtls_helper.check_parameters_for_unauthorized_response") as mock_check, \ + mock.patch.object(session, "configure_mtls_channel", new_callable=mock.AsyncMock) as mock_conf: + + # ...a 401 on a regular domain bypasses checks and just returns the 401 locally + resp = await session.request("GET", "https://pubsub.googleapis.com/test") + + assert resp == mock_resp_401 + mock_check.assert_not_called() + mock_conf.assert_not_called() + + await session.close() From 63e587cc13fa74d5bfbdebec0df196eea23f57df Mon Sep 17 00:00:00 2001 From: Radhika Agrawal Date: Fri, 28 Aug 2026 20:58:50 +0000 Subject: [PATCH 15/79] fix: fix unit tests for the checks Signed-off-by: Radhika Agrawal --- .../google/auth/aio/transport/sessions.py | 35 +++++++++++++------ .../tests/transport/aio/test_sessions_mtls.py | 6 ++-- 2 files changed, 28 insertions(+), 13 deletions(-) diff --git a/packages/google-auth/google/auth/aio/transport/sessions.py b/packages/google-auth/google/auth/aio/transport/sessions.py index 2be1dbe00147..8fa68dd73683 100644 --- a/packages/google-auth/google/auth/aio/transport/sessions.py +++ b/packages/google-auth/google/auth/aio/transport/sessions.py @@ -13,6 +13,7 @@ # limitations under the License. import asyncio +import collections.abc from contextlib import asynccontextmanager import functools import http.client as http_client @@ -321,11 +322,17 @@ async def request( if response.status_code not in transport.DEFAULT_RETRYABLE_STATUS_CODES: break - + if response.status_code == http_client.UNAUTHORIZED: _auth_retry_count = kwargs.pop("_auth_retry_count", 0) if _auth_retry_count < 2: - is_streaming = data is not None and isinstance(data, (collections.abc.Iterator, collections.abc.AsyncIterable)) or hasattr(data, "read") + is_streaming = ( + data is not None + and isinstance( + data, (collections.abc.Iterator, collections.abc.AsyncIterable) + ) + or hasattr(data, "read") + ) if getattr(self, "is_mtls", False) and any( prefix in url for prefix in MTLS_URL_PREFIXES ): @@ -335,7 +342,7 @@ async def request( # Wait in line to acquire the lock async with self._mtls_rotation_lock: - # Check Did another coroutine already reconfigure mTLS + # Check Did another coroutine already reconfigure mTLS if self._cached_cert != stale_cert: # Yes! Another request already updated the channel pass @@ -350,19 +357,30 @@ async def request( google.auth.transport._mtls_helper.check_parameters_for_unauthorized_response, self._cached_cert, ) + except Exception as e: + _LOGGER.warning( + "Failed to check client certificate parameters: %s. Proceeding with original response.", + e, + ) + else: if cached_fingerprint != current_cert_fingerprint: try: _LOGGER.info( "Client certificate has changed, reconfiguring mTLS " "channel." ) - if self._mtls_init_task and self._mtls_init_task.done(): + if ( + self._mtls_init_task + and self._mtls_init_task.done() + ): self._mtls_init_task = None await self.configure_mtls_channel( lambda: (call_cert_bytes, call_key_bytes) ) except Exception as e: - _LOGGER.error("Failed to reconfigure mTLS channel: %s", e) + _LOGGER.error( + "Failed to reconfigure mTLS channel: %s", e + ) raise exceptions.MutualTLSChannelError( "Failed to reconfigure mTLS channel" ) from e @@ -371,11 +389,6 @@ async def request( "Skipping reconfiguration of mTLS channel because the client" " certificate has not changed." ) - except Exception as e: - _LOGGER.warning( - "Failed to check client certificate parameters: %s. Proceeding with original response.", - e, - ) if is_streaming: return response if hasattr(response, "close"): @@ -393,7 +406,7 @@ async def request( max_allowed_time=max_allowed_time, timeout=timeout, total_attempts=total_attempts, - **kwargs + **kwargs, ) return response diff --git a/packages/google-auth/tests/transport/aio/test_sessions_mtls.py b/packages/google-auth/tests/transport/aio/test_sessions_mtls.py index c3110d513818..15eefccabc44 100644 --- a/packages/google-auth/tests/transport/aio/test_sessions_mtls.py +++ b/packages/google-auth/tests/transport/aio/test_sessions_mtls.py @@ -348,6 +348,8 @@ async def test_configure_mtls_channel_close_exception_does_not_abort(self): @pytest.mark.asyncio async def test_cert_rotation_failure_raises_error(self, caplog): + import logging + caplog.set_level(logging.ERROR) mock_creds = mock.AsyncMock(spec=credentials.Credentials) mock_creds.before_request = mock.AsyncMock(return_value=None) @@ -399,7 +401,7 @@ async def test_cert_rotation_check_params_fails(self, caplog): resp = await session.request("GET", "https://pubsub.mtls.googleapis.com/test") assert resp == mock_resp - mock_check.assert_called_once() + assert mock_check.call_count >= 1 mock_conf.assert_not_called() assert "Failed to check client certificate parameters" in caplog.text @@ -431,7 +433,7 @@ async def test_no_cert_rotation_when_cert_match_and_mTLS_enabled(self): resp = await session.request("GET", "https://pubsub.mtls.googleapis.com/test") assert resp == mock_resp - mock_check.assert_called_once() + assert mock_check.call_count >= 1 mock_conf.assert_not_called() await session.close() From 71b3bf545c6b926c68202c45345f2bf62062504e Mon Sep 17 00:00:00 2001 From: Radhika Agrawal Date: Fri, 28 Aug 2026 22:09:00 +0000 Subject: [PATCH 16/79] fix: Fix unit tests for the change Signed-off-by: Radhika Agrawal --- .../google/auth/aio/transport/sessions.py | 2 +- .../tests/transport/aio/test_sessions_mtls.py | 120 +++++++++++------- 2 files changed, 74 insertions(+), 48 deletions(-) diff --git a/packages/google-auth/google/auth/aio/transport/sessions.py b/packages/google-auth/google/auth/aio/transport/sessions.py index 8fa68dd73683..0a425aca20e8 100644 --- a/packages/google-auth/google/auth/aio/transport/sessions.py +++ b/packages/google-auth/google/auth/aio/transport/sessions.py @@ -287,6 +287,7 @@ async def request( google.auth.exceptions.MutualTLSChannelError: If mutual TLS channel reconfiguration fails for any reason during certificate rotation. """ + _auth_retry_count = kwargs.pop("_auth_retry_count", 0) if self._mtls_init_task: try: await self._mtls_init_task @@ -324,7 +325,6 @@ async def request( break if response.status_code == http_client.UNAUTHORIZED: - _auth_retry_count = kwargs.pop("_auth_retry_count", 0) if _auth_retry_count < 2: is_streaming = ( data is not None diff --git a/packages/google-auth/tests/transport/aio/test_sessions_mtls.py b/packages/google-auth/tests/transport/aio/test_sessions_mtls.py index 15eefccabc44..810883d7df44 100644 --- a/packages/google-auth/tests/transport/aio/test_sessions_mtls.py +++ b/packages/google-auth/tests/transport/aio/test_sessions_mtls.py @@ -12,10 +12,10 @@ # See the License for the specific language governing permissions and # limitations under the License. +import http.client as http_client import json import os import ssl -import http.client as http_client from unittest import mock import pytest @@ -349,24 +349,29 @@ async def test_configure_mtls_channel_close_exception_does_not_abort(self): @pytest.mark.asyncio async def test_cert_rotation_failure_raises_error(self, caplog): import logging - caplog.set_level(logging.ERROR) + + caplog.set_level(logging.ERROR, logger="google.auth.aio.transport.sessions") mock_creds = mock.AsyncMock(spec=credentials.Credentials) mock_creds.before_request = mock.AsyncMock(return_value=None) - + mock_resp = mock.Mock() mock_resp.status_code = http_client.UNAUTHORIZED mock_auth_req = mock.AsyncMock(return_value=mock_resp) - session = sessions.AsyncAuthorizedSession(mock_creds, auth_request=mock_auth_req) + session = sessions.AsyncAuthorizedSession( + mock_creds, auth_request=mock_auth_req + ) session._is_mtls = True session._cached_cert = b"old_cert" new_cert = b"new_cert" new_key = b"new_key" - with mock.patch("google.auth.transport._mtls_helper.check_parameters_for_unauthorized_response") as mock_check, \ - mock.patch.object(session, "configure_mtls_channel", new_callable=mock.AsyncMock) as mock_conf: - + with mock.patch( + "google.auth.transport._mtls_helper.check_parameters_for_unauthorized_response" + ) as mock_check, mock.patch.object( + session, "configure_mtls_channel", new_callable=mock.AsyncMock + ) as mock_conf: mock_check.return_value = (new_cert, new_key, b"old_fp", b"new_fp") mock_conf.side_effect = Exception("Failed to reconfigure") @@ -379,26 +384,34 @@ async def test_cert_rotation_failure_raises_error(self, caplog): await session.close() - @pytest.mark.asyncio async def test_cert_rotation_check_params_fails(self, caplog): + import logging + + caplog.set_level(logging.WARNING, logger="google.auth.aio.transport.sessions") mock_creds = mock.AsyncMock(spec=credentials.Credentials) mock_creds.before_request = mock.AsyncMock(return_value=None) - + mock_resp = mock.Mock() mock_resp.status_code = http_client.UNAUTHORIZED mock_auth_req = mock.AsyncMock(return_value=mock_resp) - session = sessions.AsyncAuthorizedSession(mock_creds, auth_request=mock_auth_req) + session = sessions.AsyncAuthorizedSession( + mock_creds, auth_request=mock_auth_req + ) session._is_mtls = True session._cached_cert = b"old_cert" - with mock.patch("google.auth.transport._mtls_helper.check_parameters_for_unauthorized_response") as mock_check, \ - mock.patch.object(session, "configure_mtls_channel", new_callable=mock.AsyncMock) as mock_conf: - + with mock.patch( + "google.auth.transport._mtls_helper.check_parameters_for_unauthorized_response" + ) as mock_check, mock.patch.object( + session, "configure_mtls_channel", new_callable=mock.AsyncMock + ) as mock_conf: mock_check.side_effect = Exception("Failed to check params") - resp = await session.request("GET", "https://pubsub.mtls.googleapis.com/test") + resp = await session.request( + "GET", "https://pubsub.mtls.googleapis.com/test" + ) assert resp == mock_resp assert mock_check.call_count >= 1 @@ -407,82 +420,91 @@ async def test_cert_rotation_check_params_fails(self, caplog): await session.close() - @pytest.mark.asyncio async def test_no_cert_rotation_when_cert_match_and_mTLS_enabled(self): mock_creds = mock.AsyncMock(spec=credentials.Credentials) mock_creds.before_request = mock.AsyncMock(return_value=None) - + mock_resp = mock.Mock() mock_resp.status_code = http_client.UNAUTHORIZED mock_auth_req = mock.AsyncMock(return_value=mock_resp) - session = sessions.AsyncAuthorizedSession(mock_creds, auth_request=mock_auth_req) + session = sessions.AsyncAuthorizedSession( + mock_creds, auth_request=mock_auth_req + ) session._is_mtls = True session._cached_cert = b"old_cert" new_cert = b"new_cert" new_key = b"new_key" - with mock.patch("google.auth.transport._mtls_helper.check_parameters_for_unauthorized_response") as mock_check, \ - mock.patch.object(session, "configure_mtls_channel", new_callable=mock.AsyncMock) as mock_conf: - + with mock.patch( + "google.auth.transport._mtls_helper.check_parameters_for_unauthorized_response" + ) as mock_check, mock.patch.object( + session, "configure_mtls_channel", new_callable=mock.AsyncMock + ) as mock_conf: # Matching fingerprints mean no layout rotation is needed mock_check.return_value = (new_cert, new_key, b"old_fp", b"old_fp") - resp = await session.request("GET", "https://pubsub.mtls.googleapis.com/test") - + resp = await session.request( + "GET", "https://pubsub.mtls.googleapis.com/test" + ) + assert resp == mock_resp assert mock_check.call_count >= 1 mock_conf.assert_not_called() await session.close() - @pytest.mark.asyncio async def test_cert_rotation_success_and_retry(self): mock_creds = mock.AsyncMock(spec=credentials.Credentials) mock_creds.before_request = mock.AsyncMock(return_value=None) mock_creds.refresh = mock.AsyncMock(return_value=None) - + # Initial request fails natively with 401. Retry succeeds with 200. mock_resp_401 = mock.Mock() mock_resp_401.status_code = http_client.UNAUTHORIZED mock_resp_200 = mock.Mock() mock_resp_200.status_code = http_client.OK - + # Use side_effect to dynamically yield responses mock_auth_req = mock.AsyncMock(side_effect=[mock_resp_401, mock_resp_200]) - session = sessions.AsyncAuthorizedSession(mock_creds, auth_request=mock_auth_req) + session = sessions.AsyncAuthorizedSession( + mock_creds, auth_request=mock_auth_req + ) session._is_mtls = True session._cached_cert = b"old_cert" new_cert = b"new_cert" new_key = b"new_key" - with mock.patch("google.auth.transport._mtls_helper.check_parameters_for_unauthorized_response") as mock_check, \ - mock.patch.object(session, "configure_mtls_channel", new_callable=mock.AsyncMock) as mock_conf: - + with mock.patch( + "google.auth.transport._mtls_helper.check_parameters_for_unauthorized_response" + ) as mock_check, mock.patch.object( + session, "configure_mtls_channel", new_callable=mock.AsyncMock + ) as mock_conf: mock_check.return_value = (new_cert, new_key, b"old_fp", b"new_fp") - - resp = await session.request("GET", "https://pubsub.mtls.googleapis.com/test") - + + resp = await session.request( + "GET", "https://pubsub.mtls.googleapis.com/test" + ) + # 1. Assert the retried 200 response is successfully returned to the user assert resp == mock_resp_200 - + # 2. Assert rotation logic correctly executed mock_check.assert_called_once() mock_conf.assert_called_once_with(mock.ANY) - + # 3. Assert credentials were explicitly refreshed mock_creds.refresh.assert_called_once() - + # 4. Assert headers were explicitly rebound on the recursive retry (2 invocations) assert mock_creds.before_request.call_count == 2 - - await session.close() + await session.close() @pytest.mark.asyncio async def test_non_mtls_url_bypasses_rotation(self): @@ -490,21 +512,25 @@ async def test_non_mtls_url_bypasses_rotation(self): mock_resp_401 = mock.Mock() mock_resp_401.status_code = http_client.UNAUTHORIZED mock_auth_req = mock.AsyncMock(return_value=mock_resp_401) - - session = sessions.AsyncAuthorizedSession(mock_creds, auth_request=mock_auth_req) - + + session = sessions.AsyncAuthorizedSession( + mock_creds, auth_request=mock_auth_req + ) + # Even if mTLS is enabled globally... - session._is_mtls = True + session._is_mtls = True session._cached_cert = b"old_cert" - - with mock.patch("google.auth.transport._mtls_helper.check_parameters_for_unauthorized_response") as mock_check, \ - mock.patch.object(session, "configure_mtls_channel", new_callable=mock.AsyncMock) as mock_conf: - + + with mock.patch( + "google.auth.transport._mtls_helper.check_parameters_for_unauthorized_response" + ) as mock_check, mock.patch.object( + session, "configure_mtls_channel", new_callable=mock.AsyncMock + ) as mock_conf: # ...a 401 on a regular domain bypasses checks and just returns the 401 locally resp = await session.request("GET", "https://pubsub.googleapis.com/test") - + assert resp == mock_resp_401 mock_check.assert_not_called() mock_conf.assert_not_called() - + await session.close() From 7d92d30c8d1f0f224fa2703721f1c5b5f7e8b559 Mon Sep 17 00:00:00 2001 From: Radhika Agrawal Date: Fri, 28 Aug 2026 22:19:37 +0000 Subject: [PATCH 17/79] test: remove fragile async caplog assertions --- .../tests/transport/aio/test_sessions_mtls.py | 12 ++---------- 1 file changed, 2 insertions(+), 10 deletions(-) diff --git a/packages/google-auth/tests/transport/aio/test_sessions_mtls.py b/packages/google-auth/tests/transport/aio/test_sessions_mtls.py index 810883d7df44..3e5ff3fc3f86 100644 --- a/packages/google-auth/tests/transport/aio/test_sessions_mtls.py +++ b/packages/google-auth/tests/transport/aio/test_sessions_mtls.py @@ -347,10 +347,7 @@ async def test_configure_mtls_channel_close_exception_does_not_abort(self): await session.close() @pytest.mark.asyncio - async def test_cert_rotation_failure_raises_error(self, caplog): - import logging - - caplog.set_level(logging.ERROR, logger="google.auth.aio.transport.sessions") + async def test_cert_rotation_failure_raises_error(self): mock_creds = mock.AsyncMock(spec=credentials.Credentials) mock_creds.before_request = mock.AsyncMock(return_value=None) @@ -380,15 +377,11 @@ async def test_cert_rotation_failure_raises_error(self, caplog): mock_check.assert_called_once() mock_conf.assert_called_once() - assert "Failed to reconfigure mTLS channel" in caplog.text await session.close() @pytest.mark.asyncio - async def test_cert_rotation_check_params_fails(self, caplog): - import logging - - caplog.set_level(logging.WARNING, logger="google.auth.aio.transport.sessions") + async def test_cert_rotation_check_params_fails(self): mock_creds = mock.AsyncMock(spec=credentials.Credentials) mock_creds.before_request = mock.AsyncMock(return_value=None) @@ -416,7 +409,6 @@ async def test_cert_rotation_check_params_fails(self, caplog): assert resp == mock_resp assert mock_check.call_count >= 1 mock_conf.assert_not_called() - assert "Failed to check client certificate parameters" in caplog.text await session.close() From 8b2efcf80c2a9c4c84a943973e84885f09cb8600 Mon Sep 17 00:00:00 2001 From: agrawalradhika-cell Date: Fri, 28 Aug 2026 15:22:38 -0700 Subject: [PATCH 18/79] fix: Add error handling for credential refresh failures Handle RefreshError during credential refresh to prevent unhandled exceptions. --- .../google-auth/google/auth/aio/transport/sessions.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/packages/google-auth/google/auth/aio/transport/sessions.py b/packages/google-auth/google/auth/aio/transport/sessions.py index 0a425aca20e8..5084f7d5fa90 100644 --- a/packages/google-auth/google/auth/aio/transport/sessions.py +++ b/packages/google-auth/google/auth/aio/transport/sessions.py @@ -396,7 +396,13 @@ async def request( await response.close() else: response.close() - await self._credentials.refresh(self._auth_request) + try: + await self._credentials.refresh(self._auth_request) + except exceptions.RefreshError as e: + _LOGGER.debug( + "Credential refresh failed, returning 401 response. Error: %s", e + ) + return response kwargs["_auth_retry_count"] = _auth_retry_count + 1 return await self.request( method, From a4d0405bcb6137221121d2654063a350cf8e136f Mon Sep 17 00:00:00 2001 From: Radhika Agrawal Date: Fri, 28 Aug 2026 22:37:59 +0000 Subject: [PATCH 19/79] fix: Fix lint errors Signed-off-by: Radhika Agrawal --- packages/google-auth/google/auth/aio/transport/sessions.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/packages/google-auth/google/auth/aio/transport/sessions.py b/packages/google-auth/google/auth/aio/transport/sessions.py index 5084f7d5fa90..28b331e3723e 100644 --- a/packages/google-auth/google/auth/aio/transport/sessions.py +++ b/packages/google-auth/google/auth/aio/transport/sessions.py @@ -400,7 +400,8 @@ async def request( await self._credentials.refresh(self._auth_request) except exceptions.RefreshError as e: _LOGGER.debug( - "Credential refresh failed, returning 401 response. Error: %s", e + "Credential refresh failed, returning 401 response. Error: %s", + e, ) return response kwargs["_auth_retry_count"] = _auth_retry_count + 1 From 2d52a21c6ad7faea220fbef88915a63bb714bce0 Mon Sep 17 00:00:00 2001 From: agrawalradhika-cell Date: Sat, 29 Aug 2026 20:51:52 -0700 Subject: [PATCH 20/79] chore: Refactor mTLS endpoint handling in sessions.py --- .../google/auth/aio/transport/sessions.py | 108 +++++++++--------- 1 file changed, 57 insertions(+), 51 deletions(-) diff --git a/packages/google-auth/google/auth/aio/transport/sessions.py b/packages/google-auth/google/auth/aio/transport/sessions.py index 28b331e3723e..de0c2778f7ec 100644 --- a/packages/google-auth/google/auth/aio/transport/sessions.py +++ b/packages/google-auth/google/auth/aio/transport/sessions.py @@ -333,62 +333,68 @@ async def request( ) or hasattr(data, "read") ) - if getattr(self, "is_mtls", False) and any( - prefix in url for prefix in MTLS_URL_PREFIXES - ): + is_mtls_endpoint = False + if getattr(self, "is_mtls", False): + hostname = urllib.parse.urlsplit(url).hostname + if hostname: + is_mtls_endpoint = any( + hostname == prefix or hostname.endswith("." + prefix) + for prefix in MTLS_URL_PREFIXES + ) # Snapshot the stale certificate state BEFORE acquiring the lock. # This represents the cert that caused the 401 rejection. - stale_cert = self._cached_cert - - # Wait in line to acquire the lock - async with self._mtls_rotation_lock: - # Check Did another coroutine already reconfigure mTLS - if self._cached_cert != stale_cert: - # Yes! Another request already updated the channel - pass - else: - try: - ( - call_cert_bytes, - call_key_bytes, - cached_fingerprint, - current_cert_fingerprint, - ) = await mtls._run_in_executor( - google.auth.transport._mtls_helper.check_parameters_for_unauthorized_response, - self._cached_cert, - ) - except Exception as e: - _LOGGER.warning( - "Failed to check client certificate parameters: %s. Proceeding with original response.", - e, - ) + if is_mtls_endpoint: + stale_cert = self._cached_cert + + # Wait in line to acquire the lock + async with self._mtls_rotation_lock: + # Check Did another coroutine already reconfigure mTLS + if self._cached_cert != stale_cert: + # Yes! Another request already updated the channel + pass else: - if cached_fingerprint != current_cert_fingerprint: - try: + try: + ( + call_cert_bytes, + call_key_bytes, + cached_fingerprint, + current_cert_fingerprint, + ) = await mtls._run_in_executor( + google.auth.transport._mtls_helper.check_parameters_for_unauthorized_response, + self._cached_cert, + ) + except Exception as e: + _LOGGER.warning( + "Failed to check client certificate parameters: %s. Proceeding with original response.", + e, + ) + else: + if cached_fingerprint != current_cert_fingerprint: + try: + _LOGGER.info( + "Client certificate has changed, reconfiguring mTLS " + "channel." + ) + if ( + self._mtls_init_task + and self._mtls_init_task.done() + ): + self._mtls_init_task = None + await self.configure_mtls_channel( + lambda: (call_cert_bytes, call_key_bytes) + ) + except Exception as e: + _LOGGER.error( + "Failed to reconfigure mTLS channel: %s", e + ) + raise exceptions.MutualTLSChannelError( + "Failed to reconfigure mTLS channel" + ) from e + else: _LOGGER.info( - "Client certificate has changed, reconfiguring mTLS " - "channel." - ) - if ( - self._mtls_init_task - and self._mtls_init_task.done() - ): - self._mtls_init_task = None - await self.configure_mtls_channel( - lambda: (call_cert_bytes, call_key_bytes) - ) - except Exception as e: - _LOGGER.error( - "Failed to reconfigure mTLS channel: %s", e + "Skipping reconfiguration of mTLS channel because the client" + " certificate has not changed." ) - raise exceptions.MutualTLSChannelError( - "Failed to reconfigure mTLS channel" - ) from e - else: - _LOGGER.info( - "Skipping reconfiguration of mTLS channel because the client" - " certificate has not changed." - ) if is_streaming: return response if hasattr(response, "close"): From 2806f4f52b57c93c5988c00c0f99f9e30ea9548b Mon Sep 17 00:00:00 2001 From: agrawalradhika-cell Date: Sat, 29 Aug 2026 22:28:43 -0700 Subject: [PATCH 21/79] chore: Reorder response closing logic for clarity chore: Reorder response closing logic for clarity --- .../google/auth/aio/transport/sessions.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/packages/google-auth/google/auth/aio/transport/sessions.py b/packages/google-auth/google/auth/aio/transport/sessions.py index de0c2778f7ec..c5cfec1d6021 100644 --- a/packages/google-auth/google/auth/aio/transport/sessions.py +++ b/packages/google-auth/google/auth/aio/transport/sessions.py @@ -397,11 +397,6 @@ async def request( ) if is_streaming: return response - if hasattr(response, "close"): - if asyncio.iscoroutinefunction(response.close): - await response.close() - else: - response.close() try: await self._credentials.refresh(self._auth_request) except exceptions.RefreshError as e: @@ -410,6 +405,13 @@ async def request( e, ) return response + + if hasattr(response, "close"): + if asyncio.iscoroutinefunction(response.close): + await response.close() + else: + response.close() + kwargs["_auth_retry_count"] = _auth_retry_count + 1 return await self.request( method, From d5426f2cf6943255ffe2621166fd2fa9b44e4d76 Mon Sep 17 00:00:00 2001 From: agrawalradhika-cell Date: Sat, 29 Aug 2026 22:31:32 -0700 Subject: [PATCH 22/79] chore: Handle additional exception during credential refresh chore: Handle additional exception during credential refresh --- packages/google-auth/google/auth/aio/transport/sessions.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/google-auth/google/auth/aio/transport/sessions.py b/packages/google-auth/google/auth/aio/transport/sessions.py index c5cfec1d6021..c521371ad51d 100644 --- a/packages/google-auth/google/auth/aio/transport/sessions.py +++ b/packages/google-auth/google/auth/aio/transport/sessions.py @@ -399,7 +399,7 @@ async def request( return response try: await self._credentials.refresh(self._auth_request) - except exceptions.RefreshError as e: + except (exceptions.RefreshError, getattr(exceptions, "InvalidOperation", Exception)) as e: _LOGGER.debug( "Credential refresh failed, returning 401 response. Error: %s", e, From 968a9fd54b588d56c2b029460980380ee29afe2c Mon Sep 17 00:00:00 2001 From: agrawalradhika-cell Date: Sat, 29 Aug 2026 23:40:14 -0700 Subject: [PATCH 23/79] fix: Modify mTLS rotation lock initialization Change _mtls_rotation_lock initialization to None and update its usage. --- .../google-auth/google/auth/aio/transport/sessions.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/packages/google-auth/google/auth/aio/transport/sessions.py b/packages/google-auth/google/auth/aio/transport/sessions.py index c521371ad51d..b7123689298f 100644 --- a/packages/google-auth/google/auth/aio/transport/sessions.py +++ b/packages/google-auth/google/auth/aio/transport/sessions.py @@ -154,7 +154,7 @@ def __init__( "`auth_request` must either be configured or the external package `aiohttp` must be installed to use the default value." ) self._auth_request = _auth_request - self._mtls_rotation_lock = asyncio.Lock() + self._mtls_rotation_lock = None async def configure_mtls_channel(self, client_cert_callback=None): """Configure the client certificate and key for SSL connection. @@ -345,8 +345,10 @@ async def request( # This represents the cert that caused the 401 rejection. if is_mtls_endpoint: stale_cert = self._cached_cert - - # Wait in line to acquire the lock + + if self._mtls_rotation_lock is None: + self._mtls_rotation_lock = asyncio.Lock() + async with self._mtls_rotation_lock: # Check Did another coroutine already reconfigure mTLS if self._cached_cert != stale_cert: From 9d1a690f22e5f6ccab469c9c33e1df41b5033c5e Mon Sep 17 00:00:00 2001 From: agrawalradhika-cell Date: Sat, 29 Aug 2026 23:59:47 -0700 Subject: [PATCH 24/79] fix: Handle response closure in mTLS error handling fix: Handle response closure in mTLS error handling --- packages/google-auth/google/auth/aio/transport/sessions.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/packages/google-auth/google/auth/aio/transport/sessions.py b/packages/google-auth/google/auth/aio/transport/sessions.py index b7123689298f..340cdf6f2cad 100644 --- a/packages/google-auth/google/auth/aio/transport/sessions.py +++ b/packages/google-auth/google/auth/aio/transport/sessions.py @@ -389,6 +389,11 @@ async def request( _LOGGER.error( "Failed to reconfigure mTLS channel: %s", e ) + if hasattr(response, "close"): + if asyncio.iscoroutinefunction(response.close): + await response.close() + else: + response.close() raise exceptions.MutualTLSChannelError( "Failed to reconfigure mTLS channel" ) from e From 85d4a766e2d669f13852dac07f311b1f272ecc9f Mon Sep 17 00:00:00 2001 From: agrawalradhika-cell Date: Sun, 30 Aug 2026 00:11:19 -0700 Subject: [PATCH 25/79] Fix: Fix improperly falling through to the credential refresh logic. Fix: Fix improperly falling through to the credential refresh logic. --- packages/google-auth/google/auth/aio/transport/sessions.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/google-auth/google/auth/aio/transport/sessions.py b/packages/google-auth/google/auth/aio/transport/sessions.py index 340cdf6f2cad..7394e63f45cc 100644 --- a/packages/google-auth/google/auth/aio/transport/sessions.py +++ b/packages/google-auth/google/auth/aio/transport/sessions.py @@ -352,7 +352,6 @@ async def request( async with self._mtls_rotation_lock: # Check Did another coroutine already reconfigure mTLS if self._cached_cert != stale_cert: - # Yes! Another request already updated the channel pass else: try: @@ -370,6 +369,7 @@ async def request( "Failed to check client certificate parameters: %s. Proceeding with original response.", e, ) + return response else: if cached_fingerprint != current_cert_fingerprint: try: From d1c6512497b5eedaa38b575885c7e8979b93692c Mon Sep 17 00:00:00 2001 From: agrawalradhika-cell Date: Sun, 30 Aug 2026 11:10:20 -0700 Subject: [PATCH 26/79] chore: Track and close old auth requests in sessions.py Add support for tracking and closing old authentication requests. --- .../google/auth/aio/transport/sessions.py | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/packages/google-auth/google/auth/aio/transport/sessions.py b/packages/google-auth/google/auth/aio/transport/sessions.py index 7394e63f45cc..266a3db88b44 100644 --- a/packages/google-auth/google/auth/aio/transport/sessions.py +++ b/packages/google-auth/google/auth/aio/transport/sessions.py @@ -149,6 +149,7 @@ def __init__( self._is_mtls = False self._mtls_init_task = None self._cached_cert = None + self._old_auth_requests = [] if _auth_request is None: raise exceptions.TransportError( "`auth_request` must either be configured or the external package `aiohttp` must be installed to use the default value." @@ -211,12 +212,8 @@ async def _do_configure(): old_auth_request = self._auth_request self._auth_request = AiohttpRequest(session=new_session) + self._old_auth_requests.append(old_auth_request) - try: - await old_auth_request.close() - except Exception: - # Suppress so it doesn't abort the mTLS configuration - pass else: is_mtls = False warnings.warn( @@ -713,3 +710,10 @@ async def close(self) -> None: except asyncio.CancelledError: pass await self._auth_request.close() + + for old_request in self._old_auth_requests: + try: + await old_request.close() + except Exception: + pass + self._old_auth_requests.clear() From b28caea4d9387c426b1fa7573bd5ff5132638034 Mon Sep 17 00:00:00 2001 From: agrawalradhika-cell Date: Sun, 30 Aug 2026 11:25:40 -0700 Subject: [PATCH 27/79] fix: Adjust max_allowed_time based on elapsed time fix: Adjust max_allowed_time based on elapsed time --- packages/google-auth/google/auth/aio/transport/sessions.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/packages/google-auth/google/auth/aio/transport/sessions.py b/packages/google-auth/google/auth/aio/transport/sessions.py index 266a3db88b44..3d01c930bc74 100644 --- a/packages/google-auth/google/auth/aio/transport/sessions.py +++ b/packages/google-auth/google/auth/aio/transport/sessions.py @@ -297,6 +297,7 @@ async def request( ) if headers is None: headers = {} + start_time = time.monotonic() async with timeout_guard(max_allowed_time) as with_timeout: await with_timeout( # Note: before_request will attempt to refresh credentials if expired. @@ -417,12 +418,14 @@ async def request( response.close() kwargs["_auth_retry_count"] = _auth_retry_count + 1 + elapsed_time = time.monotonic() - start_time + remaining_time = max(0.0, max_allowed_time - elapsed_time) return await self.request( method, url, data=data, headers=headers, - max_allowed_time=max_allowed_time, + max_allowed_time=remaining_time, timeout=timeout, total_attempts=total_attempts, **kwargs, From f176eed4d90e244e082d62dbf97992b0a0e1c70d Mon Sep 17 00:00:00 2001 From: agrawalradhika-cell Date: Sun, 30 Aug 2026 20:39:13 -0700 Subject: [PATCH 28/79] chore: Add client_cert_callback to transport session Added support for client certificate callback in the transport session. --- packages/google-auth/google/auth/aio/transport/sessions.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/packages/google-auth/google/auth/aio/transport/sessions.py b/packages/google-auth/google/auth/aio/transport/sessions.py index 3d01c930bc74..33c1751f7828 100644 --- a/packages/google-auth/google/auth/aio/transport/sessions.py +++ b/packages/google-auth/google/auth/aio/transport/sessions.py @@ -149,6 +149,7 @@ def __init__( self._is_mtls = False self._mtls_init_task = None self._cached_cert = None + self._client_cert_callback = None self._old_auth_requests = [] if _auth_request is None: raise exceptions.TransportError( @@ -183,6 +184,7 @@ async def configure_mtls_channel(self, client_cert_callback=None): creation failed for any reason. """ if self._mtls_init_task is None: + self._client_cert_callback = client_cert_callback async def _do_configure(): # Run the blocking check in an executor @@ -361,6 +363,7 @@ async def request( ) = await mtls._run_in_executor( google.auth.transport._mtls_helper.check_parameters_for_unauthorized_response, self._cached_cert, + self._client_cert_callback, ) except Exception as e: _LOGGER.warning( From a436abe9321a3d598034e8bdc7af1a257be192cf Mon Sep 17 00:00:00 2001 From: agrawalradhika-cell Date: Sun, 30 Aug 2026 20:41:14 -0700 Subject: [PATCH 29/79] chore: Enhance check_parameters_for_unauthorized_response with callback Added an optional client_cert_callback parameter to check_parameters_for_unauthorized_response to allow custom client certificate retrieval. --- .../google-auth/google/auth/transport/_mtls_helper.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/packages/google-auth/google/auth/transport/_mtls_helper.py b/packages/google-auth/google/auth/transport/_mtls_helper.py index 7779c484c713..7695a23d10ce 100644 --- a/packages/google-auth/google/auth/transport/_mtls_helper.py +++ b/packages/google-auth/google/auth/transport/_mtls_helper.py @@ -808,11 +808,13 @@ def check_use_client_cert(): return False -def check_parameters_for_unauthorized_response(cached_cert): +def check_parameters_for_unauthorized_response(cached_cert, client_cert_callback=None): """Returns the cached and current cert fingerprint for reconfiguring mTLS. Args: cached_cert(bytes): The cached client certificate. + client_cert_callback(Optional[Callable[[], (bytes, bytes)]]): + The optional callback that returns the client certificate and private key bytes. Returns: bytes: The client callback cert bytes. @@ -820,7 +822,10 @@ def check_parameters_for_unauthorized_response(cached_cert): str: The base64-encoded SHA256 cached fingerprint. str: The base64-encoded SHA256 current cert fingerprint. """ - call_cert_bytes, call_key_bytes = call_client_cert_callback() + if client_cert_callback: + call_cert_bytes, call_key_bytes = client_cert_callback() + else: + call_cert_bytes, call_key_bytes = call_client_cert_callback() cert_obj = _agent_identity_utils.parse_certificate(call_cert_bytes) current_cert_fingerprint = _agent_identity_utils.calculate_certificate_fingerprint( cert_obj From fbee9903d0a65eefddc53e4cef12c642bc76f31b Mon Sep 17 00:00:00 2001 From: agrawalradhika-cell Date: Sun, 30 Aug 2026 20:58:06 -0700 Subject: [PATCH 30/79] fix: Add test for certificate rotation lock contention Added a test for certificate rotation lock contention to ensure that multiple concurrent requests handle unauthorized responses correctly and only reconfigure the MTLS channel once. --- .../tests/transport/aio/test_sessions_mtls.py | 59 ++++++++++++++++++- 1 file changed, 57 insertions(+), 2 deletions(-) diff --git a/packages/google-auth/tests/transport/aio/test_sessions_mtls.py b/packages/google-auth/tests/transport/aio/test_sessions_mtls.py index 3e5ff3fc3f86..492917eb803f 100644 --- a/packages/google-auth/tests/transport/aio/test_sessions_mtls.py +++ b/packages/google-auth/tests/transport/aio/test_sessions_mtls.py @@ -384,6 +384,7 @@ async def test_cert_rotation_failure_raises_error(self): async def test_cert_rotation_check_params_fails(self): mock_creds = mock.AsyncMock(spec=credentials.Credentials) mock_creds.before_request = mock.AsyncMock(return_value=None) + mock_creds.refresh = mock.AsyncMock() mock_resp = mock.Mock() mock_resp.status_code = http_client.UNAUTHORIZED @@ -407,7 +408,8 @@ async def test_cert_rotation_check_params_fails(self): ) assert resp == mock_resp - assert mock_check.call_count >= 1 + mock_check.assert_called_once() + mock_creds.refresh.assert_not_called() mock_conf.assert_not_called() await session.close() @@ -498,6 +500,60 @@ async def test_cert_rotation_success_and_retry(self): await session.close() + @pytest.mark.asyncio + async def test_cert_rotation_lock_contention(self): + mock_creds = mock.AsyncMock(spec=credentials.Credentials) + mock_creds.before_request = mock.AsyncMock(return_value=None) + mock_creds.refresh = mock.AsyncMock(return_value=None) + + mock_resp_401 = mock.Mock() + mock_resp_401.status_code = http_client.UNAUTHORIZED + mock_auth_req = mock.AsyncMock(return_value=mock_resp_401) + + session = sessions.AsyncAuthorizedSession( + mock_creds, auth_request=mock_auth_req + ) + session._is_mtls = True + session._cached_cert = b"old_cert" + + new_cert = b"new_cert" + new_key = b"new_key" + + async def mock_configure_mtls_channel(*args, **kwargs): + # Introduce a small delay to guarantee lock contention from asyncio.gather tasks + await asyncio.sleep(0.01) + # Simulate the channel update to ensure following tasks take the skip branch + session._cached_cert = new_cert + + def mock_check_side_effect(cached_cert, callback): + # Return new fingerprints on first check, but matching fingerprints on retries + if cached_cert == b"old_cert": + return (new_cert, new_key, b"old_fp", b"new_fp") + return (new_cert, new_key, b"new_fp", b"new_fp") + + with mock.patch( + "google.auth.transport._mtls_helper.check_parameters_for_unauthorized_response" + ) as mock_check, mock.patch.object( + session, "configure_mtls_channel", new_callable=mock.AsyncMock + ) as mock_conf: + mock_check.side_effect = mock_check_side_effect + mock_conf.side_effect = mock_configure_mtls_channel + + # Launch multiple concurrent requests triggering 401s + tasks = [ + session.request("GET", "https://pubsub.mtls.googleapis.com/test") + for _ in range(3) + ] + responses = await asyncio.gather(*tasks) + + for resp in responses: + assert resp == mock_resp_401 + + # Confirm the channel is only reconfigured once across all concurrent requests + mock_conf.assert_called_once() + + await session.close() + @pytest.mark.asyncio async def test_non_mtls_url_bypasses_rotation(self): mock_creds = mock.AsyncMock(spec=credentials.Credentials) @@ -518,7 +574,6 @@ async def test_non_mtls_url_bypasses_rotation(self): ) as mock_check, mock.patch.object( session, "configure_mtls_channel", new_callable=mock.AsyncMock ) as mock_conf: - # ...a 401 on a regular domain bypasses checks and just returns the 401 locally resp = await session.request("GET", "https://pubsub.googleapis.com/test") assert resp == mock_resp_401 From b310e274e968700d4a01913cd9d4c9bf1b229deb Mon Sep 17 00:00:00 2001 From: agrawalradhika-cell Date: Mon, 31 Aug 2026 09:08:17 -0700 Subject: [PATCH 31/79] fix: Enhance MTLS session tests with various scenarios Add tests for MTLS session handling, including cert rotation, credential refresh failures, and session closure. --- .../tests/transport/aio/test_sessions_mtls.py | 152 ++++++++++++++++++ 1 file changed, 152 insertions(+) diff --git a/packages/google-auth/tests/transport/aio/test_sessions_mtls.py b/packages/google-auth/tests/transport/aio/test_sessions_mtls.py index 492917eb803f..6d0c0fea2cd5 100644 --- a/packages/google-auth/tests/transport/aio/test_sessions_mtls.py +++ b/packages/google-auth/tests/transport/aio/test_sessions_mtls.py @@ -574,6 +574,7 @@ async def test_non_mtls_url_bypasses_rotation(self): ) as mock_check, mock.patch.object( session, "configure_mtls_channel", new_callable=mock.AsyncMock ) as mock_conf: + # ...a 401 on a regular domain bypasses checks and just returns the 401 locally resp = await session.request("GET", "https://pubsub.googleapis.com/test") assert resp == mock_resp_401 @@ -581,3 +582,154 @@ async def test_non_mtls_url_bypasses_rotation(self): mock_conf.assert_not_called() await session.close() + + + @pytest.mark.asyncio + async def test_cert_rotation_skips_retry_for_streaming(self): + """Covers the `if is_streaming:` branch which bypasses retry for streaming data.""" + mock_creds = mock.AsyncMock(spec=credentials.Credentials) + mock_creds.before_request = mock.AsyncMock(return_value=None) + mock_creds.refresh = mock.AsyncMock() + + mock_resp = mock.Mock() + mock_resp.status_code = http_client.UNAUTHORIZED + mock_auth_req = mock.AsyncMock(return_value=mock_resp) + + session = sessions.AsyncAuthorizedSession( + mock_creds, auth_request=mock_auth_req + ) + session._is_mtls = True + session._cached_cert = b"old_cert" + + # Simulate a streaming object (hasattr(data, 'read')) + class MockStream: + def read(self): + pass + + with mock.patch( + "google.auth.transport._mtls_helper.check_parameters_for_unauthorized_response" + ) as mock_check, mock.patch.object( + session, "configure_mtls_channel", new_callable=mock.AsyncMock + ) as mock_conf: + mock_check.return_value = (b"new", b"new", b"old_fp", b"new_fp") + + resp = await session.request( + "GET", "https://pubsub.mtls.googleapis.com/test", data=MockStream() + ) + + assert resp == mock_resp + mock_conf.assert_called_once() + # Because it is streaming, it skips credentials refresh and retry + mock_creds.refresh.assert_not_called() + + await session.close() + + @pytest.mark.asyncio + async def test_cert_rotation_credential_refresh_fails(self): + """Covers the except block for RefreshError when credentials fail to refresh.""" + mock_creds = mock.AsyncMock(spec=credentials.Credentials) + mock_creds.before_request = mock.AsyncMock(return_value=None) + # Simulate a credentials refresh failure + mock_creds.refresh = mock.AsyncMock( + side_effect=exceptions.RefreshError("Refresh failed") + ) + + mock_resp = mock.Mock() + mock_resp.status_code = http_client.UNAUTHORIZED + mock_auth_req = mock.AsyncMock(return_value=mock_resp) + + session = sessions.AsyncAuthorizedSession( + mock_creds, auth_request=mock_auth_req + ) + session._is_mtls = True + session._cached_cert = b"old_cert" + + with mock.patch( + "google.auth.transport._mtls_helper.check_parameters_for_unauthorized_response" + ) as mock_check, mock.patch.object( + session, "configure_mtls_channel", new_callable=mock.AsyncMock + ) as mock_conf: + mock_check.return_value = (b"new", b"new", b"old_fp", b"new_fp") + + resp = await session.request( + "GET", "https://pubsub.mtls.googleapis.com/test" + ) + + # It should catch the RefreshError and return the 401 response directly + assert resp == mock_resp + mock_creds.refresh.assert_called_once() + # Ensure it didn't retry the request by checking auth_request was only called once + mock_auth_req.assert_called_once() + + await session.close() + + @pytest.mark.asyncio + async def test_cert_rotation_max_retries_exceeded(self): + """Covers the `if _auth_retry_count < 2:` max retry limit.""" + mock_creds = mock.AsyncMock(spec=credentials.Credentials) + mock_creds.before_request = mock.AsyncMock(return_value=None) + mock_creds.refresh = mock.AsyncMock(return_value=None) + + mock_resp = mock.Mock() + mock_resp.status_code = http_client.UNAUTHORIZED + mock_resp.close = mock.AsyncMock() + # Always return 401 + mock_auth_req = mock.AsyncMock(return_value=mock_resp) + + session = sessions.AsyncAuthorizedSession( + mock_creds, auth_request=mock_auth_req + ) + session._is_mtls = True + session._cached_cert = b"old_cert" + + with mock.patch( + "google.auth.transport._mtls_helper.check_parameters_for_unauthorized_response" + ) as mock_check, mock.patch.object( + session, "configure_mtls_channel", new_callable=mock.AsyncMock + ) as mock_conf: + # Yield new fingerprints to trigger reconfiguration branches + mock_check.return_value = (b"new", b"new", b"old_fp", b"new_fp") + + resp = await session.request( + "GET", "https://pubsub.mtls.googleapis.com/test" + ) + + assert resp == mock_resp + # Expecting exactly 3 requests (Initial try, Retry 1, Retry 2). + # When _auth_retry_count reaches 2, it drops out of the retry condition. + assert mock_auth_req.call_count == 3 + # It should have checked the cert twice before hitting the retry cap + assert mock_check.call_count == 2 + # Verify the stale response is closed before the retry + assert mock_resp.close.call_count == 2 + + await session.close() + + @pytest.mark.asyncio + async def test_session_close_cleans_old_auth_requests(self): + """Covers the loop in the `close()` method that drains `_old_auth_requests`.""" + mock_creds = mock.AsyncMock(spec=credentials.Credentials) + session = sessions.AsyncAuthorizedSession( + mock_creds, auth_request=mock.AsyncMock() + ) + + # Manually inject mocked old requests simulating past channel reconfigurations + mock_old_req_1 = mock.AsyncMock() + mock_old_req_2 = mock.AsyncMock() + mock_old_req_3_fails = mock.AsyncMock() + mock_old_req_3_fails.close.side_effect = Exception("Close error") + + session._old_auth_requests.extend( + [mock_old_req_1, mock_old_req_2, mock_old_req_3_fails] + ) + + # Triggers `await old_request.close()` for each item + await session.close() + + # Ensure all were called + mock_old_req_1.close.assert_called_once() + mock_old_req_2.close.assert_called_once() + mock_old_req_3_fails.close.assert_called_once() + + # Ensure the list was cleared + assert len(session._old_auth_requests) == 0 From 55f1ad437951957178fc5e38278ceaa97367f897 Mon Sep 17 00:00:00 2001 From: Radhika Agrawal Date: Mon, 31 Aug 2026 17:12:45 +0000 Subject: [PATCH 32/79] Fix: Fix lint and unit tetsts Signed-off-by: Radhika Agrawal --- .../google/auth/aio/transport/sessions.py | 24 +++++++++++++------ .../google/auth/transport/_mtls_helper.py | 2 +- .../tests/transport/aio/test_sessions.py | 2 +- .../tests/transport/aio/test_sessions_mtls.py | 10 ++++---- 4 files changed, 23 insertions(+), 15 deletions(-) diff --git a/packages/google-auth/google/auth/aio/transport/sessions.py b/packages/google-auth/google/auth/aio/transport/sessions.py index 33c1751f7828..ceb2440cde43 100644 --- a/packages/google-auth/google/auth/aio/transport/sessions.py +++ b/packages/google-auth/google/auth/aio/transport/sessions.py @@ -20,6 +20,7 @@ import logging import time from typing import Mapping, Optional, TYPE_CHECKING, Union +import urllib import warnings from google.auth import _exponential_backoff, exceptions @@ -149,7 +150,7 @@ def __init__( self._is_mtls = False self._mtls_init_task = None self._cached_cert = None - self._client_cert_callback = None + self._client_cert_callback = None self._old_auth_requests = [] if _auth_request is None: raise exceptions.TransportError( @@ -345,7 +346,7 @@ async def request( # This represents the cert that caused the 401 rejection. if is_mtls_endpoint: stale_cert = self._cached_cert - + if self._mtls_rotation_lock is None: self._mtls_rotation_lock = asyncio.Lock() @@ -384,14 +385,20 @@ async def request( ): self._mtls_init_task = None await self.configure_mtls_channel( - lambda: (call_cert_bytes, call_key_bytes) + lambda: ( + call_cert_bytes, + call_key_bytes, + ) ) except Exception as e: _LOGGER.error( - "Failed to reconfigure mTLS channel: %s", e + "Failed to reconfigure mTLS channel: %s", + e, ) if hasattr(response, "close"): - if asyncio.iscoroutinefunction(response.close): + if asyncio.iscoroutinefunction( + response.close + ): await response.close() else: response.close() @@ -407,7 +414,10 @@ async def request( return response try: await self._credentials.refresh(self._auth_request) - except (exceptions.RefreshError, getattr(exceptions, "InvalidOperation", Exception)) as e: + except ( + exceptions.RefreshError, + getattr(exceptions, "InvalidOperation", Exception), + ) as e: _LOGGER.debug( "Credential refresh failed, returning 401 response. Error: %s", e, @@ -419,7 +429,7 @@ async def request( await response.close() else: response.close() - + kwargs["_auth_retry_count"] = _auth_retry_count + 1 elapsed_time = time.monotonic() - start_time remaining_time = max(0.0, max_allowed_time - elapsed_time) diff --git a/packages/google-auth/google/auth/transport/_mtls_helper.py b/packages/google-auth/google/auth/transport/_mtls_helper.py index 7695a23d10ce..b36742f15455 100644 --- a/packages/google-auth/google/auth/transport/_mtls_helper.py +++ b/packages/google-auth/google/auth/transport/_mtls_helper.py @@ -813,7 +813,7 @@ def check_parameters_for_unauthorized_response(cached_cert, client_cert_callback Args: cached_cert(bytes): The cached client certificate. - client_cert_callback(Optional[Callable[[], (bytes, bytes)]]): + client_cert_callback(Optional[Callable[[], (bytes, bytes)]]): The optional callback that returns the client certificate and private key bytes. Returns: diff --git a/packages/google-auth/tests/transport/aio/test_sessions.py b/packages/google-auth/tests/transport/aio/test_sessions.py index de283b7b2e7f..b2eb2a12b8ac 100644 --- a/packages/google-auth/tests/transport/aio/test_sessions.py +++ b/packages/google-auth/tests/transport/aio/test_sessions.py @@ -255,7 +255,7 @@ async def test_request_raises_transport_error(self): async def test_request_max_allowed_time_exceeded_error(self): auth_request = MockRequest(side_effect=TransportError) authed_session = sessions.AsyncAuthorizedSession(self.credentials, auth_request) - with patch("time.monotonic", side_effect=[0, 1, 1]): + with patch("time.monotonic", side_effect=[0, 1, 1, 1, 1, 1, 1, 1]): with pytest.raises(TimeoutError): await authed_session.request("GET", self.TEST_URL, max_allowed_time=1) diff --git a/packages/google-auth/tests/transport/aio/test_sessions_mtls.py b/packages/google-auth/tests/transport/aio/test_sessions_mtls.py index 6d0c0fea2cd5..467ebc01b562 100644 --- a/packages/google-auth/tests/transport/aio/test_sessions_mtls.py +++ b/packages/google-auth/tests/transport/aio/test_sessions_mtls.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import asyncio import http.client as http_client import json import os @@ -583,10 +584,8 @@ async def test_non_mtls_url_bypasses_rotation(self): await session.close() - - @pytest.mark.asyncio + @pytest.mark.asyncio async def test_cert_rotation_skips_retry_for_streaming(self): - """Covers the `if is_streaming:` branch which bypasses retry for streaming data.""" mock_creds = mock.AsyncMock(spec=credentials.Credentials) mock_creds.before_request = mock.AsyncMock(return_value=None) mock_creds.refresh = mock.AsyncMock() @@ -648,7 +647,7 @@ async def test_cert_rotation_credential_refresh_fails(self): "google.auth.transport._mtls_helper.check_parameters_for_unauthorized_response" ) as mock_check, mock.patch.object( session, "configure_mtls_channel", new_callable=mock.AsyncMock - ) as mock_conf: + ): mock_check.return_value = (b"new", b"new", b"old_fp", b"new_fp") resp = await session.request( @@ -686,7 +685,7 @@ async def test_cert_rotation_max_retries_exceeded(self): "google.auth.transport._mtls_helper.check_parameters_for_unauthorized_response" ) as mock_check, mock.patch.object( session, "configure_mtls_channel", new_callable=mock.AsyncMock - ) as mock_conf: + ): # Yield new fingerprints to trigger reconfiguration branches mock_check.return_value = (b"new", b"new", b"old_fp", b"new_fp") @@ -730,6 +729,5 @@ async def test_session_close_cleans_old_auth_requests(self): mock_old_req_1.close.assert_called_once() mock_old_req_2.close.assert_called_once() mock_old_req_3_fails.close.assert_called_once() - # Ensure the list was cleared assert len(session._old_auth_requests) == 0 From 554a571741d8662e7e08799272d685a535a3d642 Mon Sep 17 00:00:00 2001 From: Radhika Agrawal Date: Mon, 31 Aug 2026 18:15:28 +0000 Subject: [PATCH 33/79] fix: fix unit tests for tests_sessions Signed-off-by: Radhika Agrawal --- packages/google-auth/google/auth/aio/transport/sessions.py | 4 ++-- packages/google-auth/tests/transport/aio/test_sessions.py | 6 ++++-- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/packages/google-auth/google/auth/aio/transport/sessions.py b/packages/google-auth/google/auth/aio/transport/sessions.py index ceb2440cde43..7584a4e178eb 100644 --- a/packages/google-auth/google/auth/aio/transport/sessions.py +++ b/packages/google-auth/google/auth/aio/transport/sessions.py @@ -151,13 +151,13 @@ def __init__( self._mtls_init_task = None self._cached_cert = None self._client_cert_callback = None - self._old_auth_requests = [] + self._old_auth_requests = [] # type: list if _auth_request is None: raise exceptions.TransportError( "`auth_request` must either be configured or the external package `aiohttp` must be installed to use the default value." ) self._auth_request = _auth_request - self._mtls_rotation_lock = None + self._mtls_rotation_lock = None # type: Optional[asyncio.Lock] async def configure_mtls_channel(self, client_cert_callback=None): """Configure the client certificate and key for SSL connection. diff --git a/packages/google-auth/tests/transport/aio/test_sessions.py b/packages/google-auth/tests/transport/aio/test_sessions.py index b2eb2a12b8ac..ecdb60e4af5a 100644 --- a/packages/google-auth/tests/transport/aio/test_sessions.py +++ b/packages/google-auth/tests/transport/aio/test_sessions.py @@ -105,7 +105,9 @@ async def test_timeout_with_simple_async_task_within_bounds( self, simple_async_task ): task = False - with patch("time.monotonic", side_effect=[0, 0.25, 0.75]): + with patch( + "time.monotonic", side_effect=lambda it=iter([0, 0.25]): next(it, 0.75) + ): with patch("asyncio.wait_for", lambda coro, _: coro): async with self.make_timeout_guard( timeout=self.default_timeout @@ -255,7 +257,7 @@ async def test_request_raises_transport_error(self): async def test_request_max_allowed_time_exceeded_error(self): auth_request = MockRequest(side_effect=TransportError) authed_session = sessions.AsyncAuthorizedSession(self.credentials, auth_request) - with patch("time.monotonic", side_effect=[0, 1, 1, 1, 1, 1, 1, 1]): + with patch("time.monotonic", side_effect=[0, 0] + [2] * 10): with pytest.raises(TimeoutError): await authed_session.request("GET", self.TEST_URL, max_allowed_time=1) From eb28f819df681bcbb785a615987707c3c01cf3ab Mon Sep 17 00:00:00 2001 From: agrawalradhika-cell Date: Mon, 31 Aug 2026 12:23:54 -0700 Subject: [PATCH 34/79] chore: Refactor mTLS channel configuration callback chore: Refactor mTLS channel configuration callback --- packages/google-auth/google/auth/aio/transport/sessions.py | 7 +------ 1 file changed, 1 insertion(+), 6 deletions(-) diff --git a/packages/google-auth/google/auth/aio/transport/sessions.py b/packages/google-auth/google/auth/aio/transport/sessions.py index 7584a4e178eb..1d8e97415181 100644 --- a/packages/google-auth/google/auth/aio/transport/sessions.py +++ b/packages/google-auth/google/auth/aio/transport/sessions.py @@ -384,12 +384,7 @@ async def request( and self._mtls_init_task.done() ): self._mtls_init_task = None - await self.configure_mtls_channel( - lambda: ( - call_cert_bytes, - call_key_bytes, - ) - ) + await self.configure_mtls_channel(self._client_cert_callback) except Exception as e: _LOGGER.error( "Failed to reconfigure mTLS channel: %s", From 221810e6aa76574ecf4672f0a29d0f0ab0469cf4 Mon Sep 17 00:00:00 2001 From: agrawalradhika-cell Date: Mon, 31 Aug 2026 12:37:07 -0700 Subject: [PATCH 35/79] fix: Import urllib.parse instead of urllib --- packages/google-auth/google/auth/aio/transport/sessions.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/google-auth/google/auth/aio/transport/sessions.py b/packages/google-auth/google/auth/aio/transport/sessions.py index 1d8e97415181..e47dbd821728 100644 --- a/packages/google-auth/google/auth/aio/transport/sessions.py +++ b/packages/google-auth/google/auth/aio/transport/sessions.py @@ -20,7 +20,7 @@ import logging import time from typing import Mapping, Optional, TYPE_CHECKING, Union -import urllib +import urllib.parse import warnings from google.auth import _exponential_backoff, exceptions From 445c57640031e71ff1e8cce7d95db1de4fb65774 Mon Sep 17 00:00:00 2001 From: agrawalradhika-cell Date: Mon, 31 Aug 2026 12:47:56 -0700 Subject: [PATCH 36/79] fix: Format mTLS channel configuration for readability --- packages/google-auth/google/auth/aio/transport/sessions.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/packages/google-auth/google/auth/aio/transport/sessions.py b/packages/google-auth/google/auth/aio/transport/sessions.py index e47dbd821728..fd9c146c0cf9 100644 --- a/packages/google-auth/google/auth/aio/transport/sessions.py +++ b/packages/google-auth/google/auth/aio/transport/sessions.py @@ -384,7 +384,9 @@ async def request( and self._mtls_init_task.done() ): self._mtls_init_task = None - await self.configure_mtls_channel(self._client_cert_callback) + await self.configure_mtls_channel( + self._client_cert_callback + ) except Exception as e: _LOGGER.error( "Failed to reconfigure mTLS channel: %s", From 6b0edd3f927e7c4d7b61583027ce1ad51c1f401c Mon Sep 17 00:00:00 2001 From: agrawalradhika-cell Date: Mon, 31 Aug 2026 14:46:51 -0700 Subject: [PATCH 37/79] fix: Fix test name for mTLS certificate matching --- packages/google-auth/tests/transport/aio/test_sessions_mtls.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/google-auth/tests/transport/aio/test_sessions_mtls.py b/packages/google-auth/tests/transport/aio/test_sessions_mtls.py index 467ebc01b562..3840008cebbd 100644 --- a/packages/google-auth/tests/transport/aio/test_sessions_mtls.py +++ b/packages/google-auth/tests/transport/aio/test_sessions_mtls.py @@ -416,7 +416,7 @@ async def test_cert_rotation_check_params_fails(self): await session.close() @pytest.mark.asyncio - async def test_no_cert_rotation_when_cert_match_and_mTLS_enabled(self): + async def test_no_cert_rotation_when_cert_matches_and_mtls_enabled(self): mock_creds = mock.AsyncMock(spec=credentials.Credentials) mock_creds.before_request = mock.AsyncMock(return_value=None) From b6e30b090a7edb2680f09da96cbf01193089e652 Mon Sep 17 00:00:00 2001 From: agrawalradhika-cell Date: Mon, 31 Aug 2026 15:09:18 -0700 Subject: [PATCH 38/79] fix: Refactor type annotations and error handling fix: Refactor type annotations and error handling --- .../google/auth/aio/transport/sessions.py | 17 ++++++++++------- 1 file changed, 10 insertions(+), 7 deletions(-) diff --git a/packages/google-auth/google/auth/aio/transport/sessions.py b/packages/google-auth/google/auth/aio/transport/sessions.py index fd9c146c0cf9..fe1f96ff3589 100644 --- a/packages/google-auth/google/auth/aio/transport/sessions.py +++ b/packages/google-auth/google/auth/aio/transport/sessions.py @@ -151,7 +151,7 @@ def __init__( self._mtls_init_task = None self._cached_cert = None self._client_cert_callback = None - self._old_auth_requests = [] # type: list + self._old_auth_requests: list[transport.Request] = [] if _auth_request is None: raise exceptions.TransportError( "`auth_request` must either be configured or the external package `aiohttp` must be installed to use the default value." @@ -327,11 +327,8 @@ async def request( if response.status_code == http_client.UNAUTHORIZED: if _auth_retry_count < 2: - is_streaming = ( - data is not None - and isinstance( - data, (collections.abc.Iterator, collections.abc.AsyncIterable) - ) + is_streaming = data is not None and ( + isinstance(data, (collections.abc.Iterator, collections.abc.AsyncIterable)) or hasattr(data, "read") ) is_mtls_endpoint = False @@ -366,7 +363,13 @@ async def request( self._cached_cert, self._client_cert_callback, ) - except Exception as e: + except ( + exceptions.ClientCertError, + exceptions.MutualTLSChannelError, + OSError, + ValueError, + ImportError, + ) as e: _LOGGER.warning( "Failed to check client certificate parameters: %s. Proceeding with original response.", e, From 7c32ec9376509ba76004de868d9bf8bce4ec78d0 Mon Sep 17 00:00:00 2001 From: agrawalradhika-cell Date: Mon, 31 Aug 2026 15:42:05 -0700 Subject: [PATCH 39/79] chore: Change exception type in test for MTLS session Change exception type in test for MTLS session --- packages/google-auth/tests/transport/aio/test_sessions_mtls.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/google-auth/tests/transport/aio/test_sessions_mtls.py b/packages/google-auth/tests/transport/aio/test_sessions_mtls.py index 3840008cebbd..48c4d72cbe10 100644 --- a/packages/google-auth/tests/transport/aio/test_sessions_mtls.py +++ b/packages/google-auth/tests/transport/aio/test_sessions_mtls.py @@ -402,7 +402,7 @@ async def test_cert_rotation_check_params_fails(self): ) as mock_check, mock.patch.object( session, "configure_mtls_channel", new_callable=mock.AsyncMock ) as mock_conf: - mock_check.side_effect = Exception("Failed to check params") + mock_check.side_effect = exceptions.MutualTLSChannelError("Failed to check params") resp = await session.request( "GET", "https://pubsub.mtls.googleapis.com/test" From f98c347b893677c555aa9cdc88f4d11d562ada58 Mon Sep 17 00:00:00 2001 From: agrawalradhika-cell Date: Mon, 31 Aug 2026 17:05:59 -0700 Subject: [PATCH 40/79] Fix duplicate isinstance check for data type --- packages/google-auth/google/auth/aio/transport/sessions.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/packages/google-auth/google/auth/aio/transport/sessions.py b/packages/google-auth/google/auth/aio/transport/sessions.py index fe1f96ff3589..651630da47a1 100644 --- a/packages/google-auth/google/auth/aio/transport/sessions.py +++ b/packages/google-auth/google/auth/aio/transport/sessions.py @@ -328,7 +328,9 @@ async def request( if response.status_code == http_client.UNAUTHORIZED: if _auth_retry_count < 2: is_streaming = data is not None and ( - isinstance(data, (collections.abc.Iterator, collections.abc.AsyncIterable)) + isinstance( + data, (collections.abc.Iterator, collections.abc.AsyncIterable) + ) or hasattr(data, "read") ) is_mtls_endpoint = False From c97bfdf3169aa4c6f75a5514de5d2d430aff5722 Mon Sep 17 00:00:00 2001 From: agrawalradhika-cell Date: Mon, 31 Aug 2026 17:06:49 -0700 Subject: [PATCH 41/79] fix: Format error message for MutualTLSChannelError --- .../google-auth/tests/transport/aio/test_sessions_mtls.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/packages/google-auth/tests/transport/aio/test_sessions_mtls.py b/packages/google-auth/tests/transport/aio/test_sessions_mtls.py index 48c4d72cbe10..6f1025489d74 100644 --- a/packages/google-auth/tests/transport/aio/test_sessions_mtls.py +++ b/packages/google-auth/tests/transport/aio/test_sessions_mtls.py @@ -402,7 +402,9 @@ async def test_cert_rotation_check_params_fails(self): ) as mock_check, mock.patch.object( session, "configure_mtls_channel", new_callable=mock.AsyncMock ) as mock_conf: - mock_check.side_effect = exceptions.MutualTLSChannelError("Failed to check params") + mock_check.side_effect = exceptions.MutualTLSChannelError( + "Failed to check params" + ) resp = await session.request( "GET", "https://pubsub.mtls.googleapis.com/test" From 2790119f0a0e3afc9f4536073ad339a8c4d165f1 Mon Sep 17 00:00:00 2001 From: agrawalradhika-cell Date: Mon, 31 Aug 2026 19:01:08 -0700 Subject: [PATCH 42/79] chore: Refactor mTLS handling and improve timeout logic Refactor mTLS handling and improve timeout logic --- .../google/auth/aio/transport/sessions.py | 225 ++++++++++-------- 1 file changed, 126 insertions(+), 99 deletions(-) diff --git a/packages/google-auth/google/auth/aio/transport/sessions.py b/packages/google-auth/google/auth/aio/transport/sessions.py index 651630da47a1..c158f8086e67 100644 --- a/packages/google-auth/google/auth/aio/transport/sessions.py +++ b/packages/google-auth/google/auth/aio/transport/sessions.py @@ -73,6 +73,8 @@ async def timeout_guard(timeout): total_timeout = timeout def _remaining_time(): + if total_timeout is None: + return None elapsed = time.monotonic() - start remaining = total_timeout - elapsed if remaining <= 0: @@ -327,104 +329,123 @@ async def request( if response.status_code == http_client.UNAUTHORIZED: if _auth_retry_count < 2: + if max_allowed_time is not None: + elapsed = time.monotonic() - start_time + remaining_time = max(0.0, max_allowed_time - elapsed) + if remaining_time == 0.0: + raise google.auth.exceptions.TimeoutError( + "Timeout exceeded before credential refresh could begin" + ) + else: + remaining_time = None is_streaming = data is not None and ( isinstance( data, (collections.abc.Iterator, collections.abc.AsyncIterable) ) or hasattr(data, "read") ) - is_mtls_endpoint = False - if getattr(self, "is_mtls", False): - hostname = urllib.parse.urlsplit(url).hostname - if hostname: - is_mtls_endpoint = any( - hostname == prefix or hostname.endswith("." + prefix) - for prefix in MTLS_URL_PREFIXES - ) - # Snapshot the stale certificate state BEFORE acquiring the lock. - # This represents the cert that caused the 401 rejection. - if is_mtls_endpoint: - stale_cert = self._cached_cert - - if self._mtls_rotation_lock is None: - self._mtls_rotation_lock = asyncio.Lock() - - async with self._mtls_rotation_lock: - # Check Did another coroutine already reconfigure mTLS - if self._cached_cert != stale_cert: - pass - else: - try: - ( - call_cert_bytes, - call_key_bytes, - cached_fingerprint, - current_cert_fingerprint, - ) = await mtls._run_in_executor( - google.auth.transport._mtls_helper.check_parameters_for_unauthorized_response, - self._cached_cert, - self._client_cert_callback, - ) - except ( - exceptions.ClientCertError, - exceptions.MutualTLSChannelError, - OSError, - ValueError, - ImportError, - ) as e: - _LOGGER.warning( - "Failed to check client certificate parameters: %s. Proceeding with original response.", - e, - ) - return response + async def _recover_auth_state(): + is_mtls_endpoint = False + if getattr(self, "is_mtls", False): + hostname = urllib.parse.urlsplit(url).hostname + if hostname: + is_mtls_endpoint = any( + hostname == prefix or hostname.endswith("." + prefix) + for prefix in MTLS_URL_PREFIXES + ) + # Snapshot the stale certificate state BEFORE acquiring the lock. + # This represents the cert that caused the 401 rejection. + if is_mtls_endpoint: + stale_cert = self._cached_cert + + if self._mtls_rotation_lock is None: + self._mtls_rotation_lock = asyncio.Lock() + + async with self._mtls_rotation_lock: + # Check Did another coroutine already reconfigure mTLS + if self._cached_cert != stale_cert: + pass else: - if cached_fingerprint != current_cert_fingerprint: - try: + try: + ( + call_cert_bytes, + call_key_bytes, + cached_fingerprint, + current_cert_fingerprint, + ) = await mtls._run_in_executor( + google.auth.transport._mtls_helper.check_parameters_for_unauthorized_response, + self._cached_cert, + self._client_cert_callback, + ) + except ( + exceptions.ClientCertError, + exceptions.MutualTLSChannelError, + OSError, + ValueError, + ImportError, + ) as e: + _LOGGER.warning( + "Failed to check client certificate parameters: %s. Proceeding with original response.", + e, + ) + return response + else: + if cached_fingerprint != current_cert_fingerprint: + try: + _LOGGER.info( + "Client certificate has changed, reconfiguring mTLS " + "channel." + ) + if ( + self._mtls_init_task + and self._mtls_init_task.done() + ): + self._mtls_init_task = None + await self.configure_mtls_channel( + self._client_cert_callback + ) + except Exception as e: + _LOGGER.error( + "Failed to reconfigure mTLS channel: %s", + e, + ) + if hasattr(response, "close"): + if asyncio.iscoroutinefunction( + response.close + ): + await response.close() + else: + response.close() + raise exceptions.MutualTLSChannelError( + "Failed to reconfigure mTLS channel" + ) from e + else: _LOGGER.info( - "Client certificate has changed, reconfiguring mTLS " - "channel." + "Skipping reconfiguration of mTLS channel because the client" + " certificate has not changed." ) - if ( - self._mtls_init_task - and self._mtls_init_task.done() - ): - self._mtls_init_task = None - await self.configure_mtls_channel( - self._client_cert_callback - ) - except Exception as e: - _LOGGER.error( - "Failed to reconfigure mTLS channel: %s", - e, - ) - if hasattr(response, "close"): - if asyncio.iscoroutinefunction( - response.close - ): - await response.close() - else: - response.close() - raise exceptions.MutualTLSChannelError( - "Failed to reconfigure mTLS channel" - ) from e - else: - _LOGGER.info( - "Skipping reconfiguration of mTLS channel because the client" - " certificate has not changed." - ) - if is_streaming: - return response - try: - await self._credentials.refresh(self._auth_request) - except ( - exceptions.RefreshError, - getattr(exceptions, "InvalidOperation", Exception), - ) as e: - _LOGGER.debug( - "Credential refresh failed, returning 401 response. Error: %s", - e, - ) - return response + if is_streaming: + return response + try: + await self._credentials.refresh(self._auth_request) + except ( + exceptions.RefreshError, + getattr(exceptions, "InvalidOperation", Exception), + ) as e: + _LOGGER.debug( + "Credential refresh failed, returning 401 response. Error: %s", + e, + ) + return response + + # Return None to explicitly signal successful recovery + return None + + async with timeout_guard(remaining_time) as auth_with_timeout: + early_return_response = await auth_with_timeout(_recover_auth_state()) + # If it returned a response (meaning streaming or error), bail out + if early_return_response is not None: + return early_return_response if hasattr(response, "close"): if asyncio.iscoroutinefunction(response.close): @@ -432,9 +453,14 @@ async def request( else: response.close() + if max_allowed_time is not None: + remaining_time = max(0.0, max_allowed_time - (time.monotonic() - start_time)) + if remaining_time == 0.0: + raise google.auth.exceptions.TimeoutError( + "Timeout exceeded before retrying the request" + ) + kwargs["_auth_retry_count"] = _auth_retry_count + 1 - elapsed_time = time.monotonic() - start_time - remaining_time = max(0.0, max_allowed_time - elapsed_time) return await self.request( method, url, @@ -727,11 +753,12 @@ async def close(self) -> None: await self._mtls_init_task except asyncio.CancelledError: pass - await self._auth_request.close() - - for old_request in self._old_auth_requests: - try: - await old_request.close() - except Exception: - pass - self._old_auth_requests.clear() + try: + await self._auth_request.close() + finally: + for old_request in self._old_auth_requests: + try: + await old_request.close() + except Exception: + pass + self._old_auth_requests.clear() From 498bcd6339122cfa8a274812d829cf92abb5f49b Mon Sep 17 00:00:00 2001 From: agrawalradhika-cell Date: Tue, 1 Sep 2026 07:45:20 -0700 Subject: [PATCH 43/79] chore: Implement mTLS check counter for configuration management Add a counter to track mTLS configuration checks and prevent redundant operations. --- .../google/auth/aio/transport/sessions.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/packages/google-auth/google/auth/aio/transport/sessions.py b/packages/google-auth/google/auth/aio/transport/sessions.py index c158f8086e67..882c0bbc9fde 100644 --- a/packages/google-auth/google/auth/aio/transport/sessions.py +++ b/packages/google-auth/google/auth/aio/transport/sessions.py @@ -160,6 +160,8 @@ def __init__( ) self._auth_request = _auth_request self._mtls_rotation_lock = None # type: Optional[asyncio.Lock] + self._mtls_check_counter = 0 + async def configure_mtls_channel(self, client_cert_callback=None): """Configure the client certificate and key for SSL connection. @@ -360,10 +362,13 @@ async def _recover_auth_state(): if self._mtls_rotation_lock is None: self._mtls_rotation_lock = asyncio.Lock() + # Snapshot the counter state BEFORE acquiring the lock. + check_counter_at_error = self._mtls_check_counter async with self._mtls_rotation_lock: - # Check Did another coroutine already reconfigure mTLS - if self._cached_cert != stale_cert: + # Check if another coroutine already reconfigured mTLS or + # ran the validation check. + if self._mtls_check_counter > check_counter_at_error: pass else: try: @@ -424,6 +429,9 @@ async def _recover_auth_state(): "Skipping reconfiguration of mTLS channel because the client" " certificate has not changed." ) + finally: + # Always increment so waiting tasks skip the check block + self._mtls_check_counter += 1 if is_streaming: return response try: From 242359aecdec34de7a1b97101e9a14c29c2e0e40 Mon Sep 17 00:00:00 2001 From: agrawalradhika-cell Date: Tue, 1 Sep 2026 08:07:23 -0700 Subject: [PATCH 44/79] fix: Implement test for cert rotation lock contention Add test for certificate rotation lock contention without cert change. --- .../tests/transport/aio/test_sessions_mtls.py | 65 ++++++++++++++++++- 1 file changed, 64 insertions(+), 1 deletion(-) diff --git a/packages/google-auth/tests/transport/aio/test_sessions_mtls.py b/packages/google-auth/tests/transport/aio/test_sessions_mtls.py index 6f1025489d74..4e0078df4028 100644 --- a/packages/google-auth/tests/transport/aio/test_sessions_mtls.py +++ b/packages/google-auth/tests/transport/aio/test_sessions_mtls.py @@ -558,8 +558,66 @@ def mock_check_side_effect(cached_cert, callback): await session.close() @pytest.mark.asyncio + async def test_cert_rotation_lock_contention_no_cert_change(self): + mock_creds = mock.AsyncMock(spec=credentials.Credentials) + mock_creds.before_request = mock.AsyncMock(return_value=None) + mock_creds.refresh = mock.AsyncMock(return_value=None) + + mock_resp_401 = mock.Mock() + mock_resp_401.status_code = http_client.UNAUTHORIZED + mock_resp_200 = mock.Mock() + mock_resp_200.status_code = http_client.OK + + # Return 401 for the first 3 requests, then 200 for the retries. + mock_auth_req = mock.AsyncMock( + side_effect=[mock_resp_401] * 3 + [mock_resp_200] * 3 + ) + + session = sessions.AsyncAuthorizedSession( + mock_creds, auth_request=mock_auth_req + ) + session._is_mtls = True + session._cached_cert = b"old_cert" + + def mock_check_side_effect(cached_cert, callback): + import time + # Introduce a small delay to guarantee lock contention from asyncio.gather tasks + time.sleep(0.01) + # Return matching fingerprints to skip reconfiguration + return (b"old_cert", b"old_key", b"old_fp", b"old_fp") + + with mock.patch( + "google.auth.transport._mtls_helper.check_parameters_for_unauthorized_response" + ) as mock_check, mock.patch.object( + session, "configure_mtls_channel", new_callable=mock.AsyncMock + ) as mock_conf: + mock_check.side_effect = mock_check_side_effect + + # Launch multiple concurrent requests triggering 401s + tasks = [ + session.request("GET", "https://pubsub.mtls.googleapis.com/test") + for _ in range(3) + ] + responses = await asyncio.gather(*tasks) + + for resp in responses: + assert resp == mock_resp_200 + + # Confirm that the cert parameters were only checked once across the concurrent requests + mock_check.assert_called_once() + # Because fingerprints match, reconfiguration should not happen + mock_conf.assert_not_called() + # Credentials should still be refreshed for each task as they fall back to normal refresh logic + assert mock_creds.refresh.call_count == 3 + + await session.close() + + @pytest.mark.asyncio async def test_non_mtls_url_bypasses_rotation(self): mock_creds = mock.AsyncMock(spec=credentials.Credentials) + mock_creds.before_request = mock.AsyncMock(return_value=None) + mock_creds.refresh = mock.AsyncMock(return_value=None) + mock_resp_401 = mock.Mock() mock_resp_401.status_code = http_client.UNAUTHORIZED mock_auth_req = mock.AsyncMock(return_value=mock_resp_401) @@ -577,15 +635,20 @@ async def test_non_mtls_url_bypasses_rotation(self): ) as mock_check, mock.patch.object( session, "configure_mtls_channel", new_callable=mock.AsyncMock ) as mock_conf: - # ...a 401 on a regular domain bypasses checks and just returns the 401 locally + # Will attempt to refresh credentials and retry the request up to 2 times (3 requests total). resp = await session.request("GET", "https://pubsub.googleapis.com/test") assert resp == mock_resp_401 mock_check.assert_not_called() mock_conf.assert_not_called() + + # Verify the standard retry behavior executed + assert mock_creds.refresh.call_count == 2 + assert mock_auth_req.call_count == 3 await session.close() + @pytest.mark.asyncio async def test_cert_rotation_skips_retry_for_streaming(self): mock_creds = mock.AsyncMock(spec=credentials.Credentials) From 8c867af847f33a31e10a3df6a9aa64f9455f95a8 Mon Sep 17 00:00:00 2001 From: agrawalradhika-cell Date: Tue, 1 Sep 2026 08:11:03 -0700 Subject: [PATCH 45/79] fix: Refactor mTLS configuration and error handling fix: Refactor mTLS configuration and error handling --- .../google/auth/aio/transport/sessions.py | 37 +++++++++++-------- 1 file changed, 22 insertions(+), 15 deletions(-) diff --git a/packages/google-auth/google/auth/aio/transport/sessions.py b/packages/google-auth/google/auth/aio/transport/sessions.py index 882c0bbc9fde..181e1f0d7775 100644 --- a/packages/google-auth/google/auth/aio/transport/sessions.py +++ b/packages/google-auth/google/auth/aio/transport/sessions.py @@ -162,7 +162,6 @@ def __init__( self._mtls_rotation_lock = None # type: Optional[asyncio.Lock] self._mtls_check_counter = 0 - async def configure_mtls_channel(self, client_cert_callback=None): """Configure the client certificate and key for SSL connection. @@ -332,20 +331,21 @@ async def request( if response.status_code == http_client.UNAUTHORIZED: if _auth_retry_count < 2: if max_allowed_time is not None: - elapsed = time.monotonic() - start_time - remaining_time = max(0.0, max_allowed_time - elapsed) - if remaining_time == 0.0: - raise google.auth.exceptions.TimeoutError( - "Timeout exceeded before credential refresh could begin" - ) + elapsed = time.monotonic() - start_time + remaining_time = max(0.0, max_allowed_time - elapsed) + if remaining_time == 0.0: + raise google.auth.exceptions.TimeoutError( + "Timeout exceeded before credential refresh could begin" + ) else: - remaining_time = None + remaining_time = None is_streaming = data is not None and ( isinstance( data, (collections.abc.Iterator, collections.abc.AsyncIterable) ) or hasattr(data, "read") ) + async def _recover_auth_state(): is_mtls_endpoint = False if getattr(self, "is_mtls", False): @@ -359,14 +359,14 @@ async def _recover_auth_state(): # This represents the cert that caused the 401 rejection. if is_mtls_endpoint: stale_cert = self._cached_cert - + if self._mtls_rotation_lock is None: self._mtls_rotation_lock = asyncio.Lock() # Snapshot the counter state BEFORE acquiring the lock. check_counter_at_error = self._mtls_check_counter - + async with self._mtls_rotation_lock: - # Check if another coroutine already reconfigured mTLS or + # Check if another coroutine already reconfigured mTLS or # ran the validation check. if self._mtls_check_counter > check_counter_at_error: pass @@ -395,7 +395,10 @@ async def _recover_auth_state(): ) return response else: - if cached_fingerprint != current_cert_fingerprint: + if ( + cached_fingerprint + != current_cert_fingerprint + ): try: _LOGGER.info( "Client certificate has changed, reconfiguring mTLS " @@ -445,12 +448,14 @@ async def _recover_auth_state(): e, ) return response - + # Return None to explicitly signal successful recovery return None async with timeout_guard(remaining_time) as auth_with_timeout: - early_return_response = await auth_with_timeout(_recover_auth_state()) + early_return_response = await auth_with_timeout( + _recover_auth_state() + ) # If it returned a response (meaning streaming or error), bail out if early_return_response is not None: return early_return_response @@ -462,7 +467,9 @@ async def _recover_auth_state(): response.close() if max_allowed_time is not None: - remaining_time = max(0.0, max_allowed_time - (time.monotonic() - start_time)) + remaining_time = max( + 0.0, max_allowed_time - (time.monotonic() - start_time) + ) if remaining_time == 0.0: raise google.auth.exceptions.TimeoutError( "Timeout exceeded before retrying the request" From 8942a127477a8752a08b9815278147d820f54837 Mon Sep 17 00:00:00 2001 From: agrawalradhika-cell Date: Tue, 1 Sep 2026 08:18:04 -0700 Subject: [PATCH 46/79] fix: Fix indentation for asyncio test decorator fix: Fix indentation for asyncio test decorator --- packages/google-auth/tests/transport/aio/test_sessions_mtls.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/google-auth/tests/transport/aio/test_sessions_mtls.py b/packages/google-auth/tests/transport/aio/test_sessions_mtls.py index 4e0078df4028..6f641467323e 100644 --- a/packages/google-auth/tests/transport/aio/test_sessions_mtls.py +++ b/packages/google-auth/tests/transport/aio/test_sessions_mtls.py @@ -612,7 +612,7 @@ def mock_check_side_effect(cached_cert, callback): await session.close() - @pytest.mark.asyncio + @pytest.mark.asyncio async def test_non_mtls_url_bypasses_rotation(self): mock_creds = mock.AsyncMock(spec=credentials.Credentials) mock_creds.before_request = mock.AsyncMock(return_value=None) From eb95f3c033ae0c6ea45e1fb43a430fbb4cfaaab6 Mon Sep 17 00:00:00 2001 From: agrawalradhika-cell Date: Tue, 1 Sep 2026 08:21:06 -0700 Subject: [PATCH 47/79] fix: Enhance tests for MTLS session certificate rotation Added a delay in mock_check to ensure lock contention during asyncio.gather tasks. Adjusted assertions to verify behavior when non-mtls URL is used. --- .../google-auth/tests/transport/aio/test_sessions_mtls.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/packages/google-auth/tests/transport/aio/test_sessions_mtls.py b/packages/google-auth/tests/transport/aio/test_sessions_mtls.py index 6f641467323e..b3e458d338b5 100644 --- a/packages/google-auth/tests/transport/aio/test_sessions_mtls.py +++ b/packages/google-auth/tests/transport/aio/test_sessions_mtls.py @@ -581,6 +581,7 @@ async def test_cert_rotation_lock_contention_no_cert_change(self): def mock_check_side_effect(cached_cert, callback): import time + # Introduce a small delay to guarantee lock contention from asyncio.gather tasks time.sleep(0.01) # Return matching fingerprints to skip reconfiguration @@ -617,7 +618,7 @@ async def test_non_mtls_url_bypasses_rotation(self): mock_creds = mock.AsyncMock(spec=credentials.Credentials) mock_creds.before_request = mock.AsyncMock(return_value=None) mock_creds.refresh = mock.AsyncMock(return_value=None) - + mock_resp_401 = mock.Mock() mock_resp_401.status_code = http_client.UNAUTHORIZED mock_auth_req = mock.AsyncMock(return_value=mock_resp_401) @@ -641,14 +642,13 @@ async def test_non_mtls_url_bypasses_rotation(self): assert resp == mock_resp_401 mock_check.assert_not_called() mock_conf.assert_not_called() - + # Verify the standard retry behavior executed assert mock_creds.refresh.call_count == 2 assert mock_auth_req.call_count == 3 await session.close() - @pytest.mark.asyncio async def test_cert_rotation_skips_retry_for_streaming(self): mock_creds = mock.AsyncMock(spec=credentials.Credentials) From 50ba4854ecc6aa5ebe6bd68f4dfe02b928384860 Mon Sep 17 00:00:00 2001 From: agrawalradhika-cell Date: Tue, 1 Sep 2026 09:12:18 -0700 Subject: [PATCH 48/79] fix: Remove stale_cert assignment in sessions.py Remove assignment of stale_cert when is_mtls_endpoint is true. --- packages/google-auth/google/auth/aio/transport/sessions.py | 1 - 1 file changed, 1 deletion(-) diff --git a/packages/google-auth/google/auth/aio/transport/sessions.py b/packages/google-auth/google/auth/aio/transport/sessions.py index 181e1f0d7775..cb95029553c0 100644 --- a/packages/google-auth/google/auth/aio/transport/sessions.py +++ b/packages/google-auth/google/auth/aio/transport/sessions.py @@ -358,7 +358,6 @@ async def _recover_auth_state(): # Snapshot the stale certificate state BEFORE acquiring the lock. # This represents the cert that caused the 401 rejection. if is_mtls_endpoint: - stale_cert = self._cached_cert if self._mtls_rotation_lock is None: self._mtls_rotation_lock = asyncio.Lock() From 0e8950fec8e2373b3c1bbf94d62f6d827b85539c Mon Sep 17 00:00:00 2001 From: agrawalradhika-cell Date: Tue, 1 Sep 2026 09:14:41 -0700 Subject: [PATCH 49/79] fix: Update sessions.py for lint From bc8bebb57d1a370fde82bb2d3e37a8854bb029e6 Mon Sep 17 00:00:00 2001 From: Radhika Agrawal Date: Tue, 1 Sep 2026 16:23:18 +0000 Subject: [PATCH 50/79] fix: fix lint errors in sessions.py Signed-off-by: Radhika Agrawal --- packages/google-auth/google/auth/aio/transport/sessions.py | 1 - 1 file changed, 1 deletion(-) diff --git a/packages/google-auth/google/auth/aio/transport/sessions.py b/packages/google-auth/google/auth/aio/transport/sessions.py index cb95029553c0..91bb6c3073b3 100644 --- a/packages/google-auth/google/auth/aio/transport/sessions.py +++ b/packages/google-auth/google/auth/aio/transport/sessions.py @@ -358,7 +358,6 @@ async def _recover_auth_state(): # Snapshot the stale certificate state BEFORE acquiring the lock. # This represents the cert that caused the 401 rejection. if is_mtls_endpoint: - if self._mtls_rotation_lock is None: self._mtls_rotation_lock = asyncio.Lock() # Snapshot the counter state BEFORE acquiring the lock. From 2b57c1141046ba2f12681711bbabf4bdac3ffcd1 Mon Sep 17 00:00:00 2001 From: agrawalradhika-cell Date: Tue, 1 Sep 2026 14:06:51 -0700 Subject: [PATCH 51/79] Rename MTLS_URL_PREFIXES to _MTLS_URL_PREFIXES and typecasting fix --- packages/google-auth/google/auth/aio/transport/sessions.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/packages/google-auth/google/auth/aio/transport/sessions.py b/packages/google-auth/google/auth/aio/transport/sessions.py index 91bb6c3073b3..d22518b40dd2 100644 --- a/packages/google-auth/google/auth/aio/transport/sessions.py +++ b/packages/google-auth/google/auth/aio/transport/sessions.py @@ -42,7 +42,7 @@ ClientTimeout = None _LOGGER = logging.getLogger(__name__) -MTLS_URL_PREFIXES = ["mtls.googleapis.com", "mtls.sandbox.googleapis.com"] +_MTLS_URL_PREFIXES = ["mtls.googleapis.com", "mtls.sandbox.googleapis.com"] # Tracks the internal aiohttp installation and usage @@ -159,7 +159,7 @@ def __init__( "`auth_request` must either be configured or the external package `aiohttp` must be installed to use the default value." ) self._auth_request = _auth_request - self._mtls_rotation_lock = None # type: Optional[asyncio.Lock] + self._mtls_rotation_lock: Optional[asyncio.Lock] = None self._mtls_check_counter = 0 async def configure_mtls_channel(self, client_cert_callback=None): @@ -353,7 +353,7 @@ async def _recover_auth_state(): if hostname: is_mtls_endpoint = any( hostname == prefix or hostname.endswith("." + prefix) - for prefix in MTLS_URL_PREFIXES + for prefix in _MTLS_URL_PREFIXES ) # Snapshot the stale certificate state BEFORE acquiring the lock. # This represents the cert that caused the 401 rejection. From a149eb7bbd6f571b37a43d08a8adcd6530196664 Mon Sep 17 00:00:00 2001 From: agrawalradhika-cell Date: Wed, 2 Sep 2026 08:20:30 -0700 Subject: [PATCH 52/79] fix: Refactor request headers handling in sessions.py fix: Refactor request headers handling in sessions.py --- packages/google-auth/google/auth/aio/transport/sessions.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/packages/google-auth/google/auth/aio/transport/sessions.py b/packages/google-auth/google/auth/aio/transport/sessions.py index d22518b40dd2..5ebcf380a14e 100644 --- a/packages/google-auth/google/auth/aio/transport/sessions.py +++ b/packages/google-auth/google/auth/aio/transport/sessions.py @@ -301,14 +301,13 @@ async def request( retries = _exponential_backoff.AsyncExponentialBackoff( total_attempts=total_attempts, ) - if headers is None: - headers = {} + request_headers = dict(headers) if headers is not None else {} start_time = time.monotonic() async with timeout_guard(max_allowed_time) as with_timeout: await with_timeout( # Note: before_request will attempt to refresh credentials if expired. self._credentials.before_request( - self._auth_request, method, url, headers + self._auth_request, method, url, request_headers ) ) actual_timeout: float = 0.0 @@ -321,7 +320,7 @@ async def request( async for _ in retries: # pragma: no branch response = await with_timeout( self._auth_request( - url, method, data, headers, actual_timeout, **kwargs + url, method, data, request_headers, actual_timeout, **kwargs ) ) From e66a7233fd5969f1f0e7b444f904941a8aefa02c Mon Sep 17 00:00:00 2001 From: agrawalradhika-cell Date: Wed, 2 Sep 2026 09:51:26 -0700 Subject: [PATCH 53/79] fix: Fix client certificate callback handling in mTLS --- packages/google-auth/google/auth/aio/transport/sessions.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/packages/google-auth/google/auth/aio/transport/sessions.py b/packages/google-auth/google/auth/aio/transport/sessions.py index 5ebcf380a14e..637322df5f4e 100644 --- a/packages/google-auth/google/auth/aio/transport/sessions.py +++ b/packages/google-auth/google/auth/aio/transport/sessions.py @@ -396,6 +396,7 @@ async def _recover_auth_state(): cached_fingerprint != current_cert_fingerprint ): + saved_callback = self._client_cert_callback try: _LOGGER.info( "Client certificate has changed, reconfiguring mTLS " @@ -407,7 +408,7 @@ async def _recover_auth_state(): ): self._mtls_init_task = None await self.configure_mtls_channel( - self._client_cert_callback + lambda: (call_cert_bytes, call_key_bytes) ) except Exception as e: _LOGGER.error( @@ -424,6 +425,8 @@ async def _recover_auth_state(): raise exceptions.MutualTLSChannelError( "Failed to reconfigure mTLS channel" ) from e + finally: + self._client_cert_callback = saved_callback else: _LOGGER.info( "Skipping reconfiguration of mTLS channel because the client" From 314d974f545a890e6deeb2984921bd3818c35052 Mon Sep 17 00:00:00 2001 From: agrawalradhika-cell Date: Wed, 2 Sep 2026 11:19:11 -0700 Subject: [PATCH 54/79] feat: Implement mTLS parameter check and fingerprinting Add async helper to check parameters for mTLS rotation and compute fingerprints. --- .../google/auth/aio/transport/mtls.py | 30 ++++++++++++++++++- 1 file changed, 29 insertions(+), 1 deletion(-) diff --git a/packages/google-auth/google/auth/aio/transport/mtls.py b/packages/google-auth/google/auth/aio/transport/mtls.py index a7d1baf7355d..b13787476dec 100644 --- a/packages/google-auth/google/auth/aio/transport/mtls.py +++ b/packages/google-auth/google/auth/aio/transport/mtls.py @@ -22,7 +22,8 @@ import ssl from typing import Optional -from google.auth import exceptions +from google.auth import _agent_identity_utils, exceptions +from google.auth.transport import _mtls_helper from google.auth.transport._mtls_helper import secure_cert_key_paths import google.auth.transport.mtls @@ -177,3 +178,30 @@ async def get_client_cert_and_key(client_cert_callback=None): has_cert, cert, key, _ = await get_client_ssl_credentials() return has_cert, cert, key + + +async def check_parameters_for_unauthorized_response(client_cert_callback, cached_cert): + """Async helper to retrieve certs and compute fingerprints for mTLS rotation.""" + is_mtls, call_cert_bytes, call_key_bytes = await get_client_cert_and_key( + client_cert_callback + ) + if not is_mtls or not call_cert_bytes: + return None, None, None, None + + def _fetch_fingerprints(): + cert_obj = _agent_identity_utils.parse_certificate(call_cert_bytes) + current_fingerprint = _agent_identity_utils.calculate_certificate_fingerprint( + cert_obj + ) + if cached_cert: + cached_fingerprint = _agent_identity_utils.get_cached_cert_fingerprint( + cached_cert + ) + else: + cached_fingerprint = current_fingerprint + return cached_fingerprint, current_fingerprint + + cached_fingerprint, current_cert_fingerprint = await _run_in_executor( + _fetch_fingerprints + ) + return call_cert_bytes, call_key_bytes, cached_fingerprint, current_cert_fingerprint From cf308b8bc5218816086c05079e740241092c13e3 Mon Sep 17 00:00:00 2001 From: agrawalradhika-cell Date: Wed, 2 Sep 2026 11:25:02 -0700 Subject: [PATCH 55/79] fix: Refactor MTLS parameter checking in sessions.py --- .../google-auth/google/auth/aio/transport/sessions.py | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/packages/google-auth/google/auth/aio/transport/sessions.py b/packages/google-auth/google/auth/aio/transport/sessions.py index 637322df5f4e..809f545f7462 100644 --- a/packages/google-auth/google/auth/aio/transport/sessions.py +++ b/packages/google-auth/google/auth/aio/transport/sessions.py @@ -374,10 +374,9 @@ async def _recover_auth_state(): call_key_bytes, cached_fingerprint, current_cert_fingerprint, - ) = await mtls._run_in_executor( - google.auth.transport._mtls_helper.check_parameters_for_unauthorized_response, - self._cached_cert, - self._client_cert_callback, + ) = await mtls.check_parameters_for_unauthorized_response( + self._client_cert_callback, + self._cached_cert, ) except ( exceptions.ClientCertError, @@ -393,8 +392,8 @@ async def _recover_auth_state(): return response else: if ( - cached_fingerprint - != current_cert_fingerprint + current_cert_fingerprint is not None + and cached_fingerprint != current_cert_fingerprint ): saved_callback = self._client_cert_callback try: From 77945063b500f4d82c6f25f0d46366b7dac0296d Mon Sep 17 00:00:00 2001 From: agrawalradhika-cell Date: Wed, 2 Sep 2026 12:01:32 -0700 Subject: [PATCH 56/79] fix: Update mock patch for MTLS check parameters --- .../tests/transport/aio/test_sessions_mtls.py | 30 ++++++++++++------- 1 file changed, 20 insertions(+), 10 deletions(-) diff --git a/packages/google-auth/tests/transport/aio/test_sessions_mtls.py b/packages/google-auth/tests/transport/aio/test_sessions_mtls.py index b3e458d338b5..34c1f27f0ebc 100644 --- a/packages/google-auth/tests/transport/aio/test_sessions_mtls.py +++ b/packages/google-auth/tests/transport/aio/test_sessions_mtls.py @@ -366,7 +366,8 @@ async def test_cert_rotation_failure_raises_error(self): new_key = b"new_key" with mock.patch( - "google.auth.transport._mtls_helper.check_parameters_for_unauthorized_response" + "google.auth.aio.transport.mtls.check_parameters_for_unauthorized_response", + new_callable=mock.AsyncMock, ) as mock_check, mock.patch.object( session, "configure_mtls_channel", new_callable=mock.AsyncMock ) as mock_conf: @@ -398,7 +399,8 @@ async def test_cert_rotation_check_params_fails(self): session._cached_cert = b"old_cert" with mock.patch( - "google.auth.transport._mtls_helper.check_parameters_for_unauthorized_response" + "google.auth.aio.transport.mtls.check_parameters_for_unauthorized_response", + new_callable=mock.AsyncMock, ) as mock_check, mock.patch.object( session, "configure_mtls_channel", new_callable=mock.AsyncMock ) as mock_conf: @@ -436,7 +438,8 @@ async def test_no_cert_rotation_when_cert_matches_and_mtls_enabled(self): new_key = b"new_key" with mock.patch( - "google.auth.transport._mtls_helper.check_parameters_for_unauthorized_response" + "google.auth.aio.transport.mtls.check_parameters_for_unauthorized_response", + new_callable=mock.AsyncMock, ) as mock_check, mock.patch.object( session, "configure_mtls_channel", new_callable=mock.AsyncMock ) as mock_conf: @@ -478,7 +481,8 @@ async def test_cert_rotation_success_and_retry(self): new_key = b"new_key" with mock.patch( - "google.auth.transport._mtls_helper.check_parameters_for_unauthorized_response" + "google.auth.aio.transport.mtls.check_parameters_for_unauthorized_response", + new_callable=mock.AsyncMock, ) as mock_check, mock.patch.object( session, "configure_mtls_channel", new_callable=mock.AsyncMock ) as mock_conf: @@ -535,7 +539,8 @@ def mock_check_side_effect(cached_cert, callback): return (new_cert, new_key, b"new_fp", b"new_fp") with mock.patch( - "google.auth.transport._mtls_helper.check_parameters_for_unauthorized_response" + "google.auth.aio.transport.mtls.check_parameters_for_unauthorized_response", + new_callable=mock.AsyncMock, ) as mock_check, mock.patch.object( session, "configure_mtls_channel", new_callable=mock.AsyncMock ) as mock_conf: @@ -588,7 +593,8 @@ def mock_check_side_effect(cached_cert, callback): return (b"old_cert", b"old_key", b"old_fp", b"old_fp") with mock.patch( - "google.auth.transport._mtls_helper.check_parameters_for_unauthorized_response" + "google.auth.aio.transport.mtls.check_parameters_for_unauthorized_response", + new_callable=mock.AsyncMock, ) as mock_check, mock.patch.object( session, "configure_mtls_channel", new_callable=mock.AsyncMock ) as mock_conf: @@ -632,7 +638,8 @@ async def test_non_mtls_url_bypasses_rotation(self): session._cached_cert = b"old_cert" with mock.patch( - "google.auth.transport._mtls_helper.check_parameters_for_unauthorized_response" + "google.auth.aio.transport.mtls.check_parameters_for_unauthorized_response", + new_callable=mock.AsyncMock, ) as mock_check, mock.patch.object( session, "configure_mtls_channel", new_callable=mock.AsyncMock ) as mock_conf: @@ -671,7 +678,8 @@ def read(self): pass with mock.patch( - "google.auth.transport._mtls_helper.check_parameters_for_unauthorized_response" + "google.auth.aio.transport.mtls.check_parameters_for_unauthorized_response", + new_callable=mock.AsyncMock, ) as mock_check, mock.patch.object( session, "configure_mtls_channel", new_callable=mock.AsyncMock ) as mock_conf: @@ -709,7 +717,8 @@ async def test_cert_rotation_credential_refresh_fails(self): session._cached_cert = b"old_cert" with mock.patch( - "google.auth.transport._mtls_helper.check_parameters_for_unauthorized_response" + "google.auth.aio.transport.mtls.check_parameters_for_unauthorized_response", + new_callable=mock.AsyncMock, ) as mock_check, mock.patch.object( session, "configure_mtls_channel", new_callable=mock.AsyncMock ): @@ -747,7 +756,8 @@ async def test_cert_rotation_max_retries_exceeded(self): session._cached_cert = b"old_cert" with mock.patch( - "google.auth.transport._mtls_helper.check_parameters_for_unauthorized_response" + "google.auth.aio.transport.mtls.check_parameters_for_unauthorized_response", + new_callable=mock.AsyncMock, ) as mock_check, mock.patch.object( session, "configure_mtls_channel", new_callable=mock.AsyncMock ): From 931bdfd17d31e11003b8db8c99fbdc6d4b3e468c Mon Sep 17 00:00:00 2001 From: agrawalradhika-cell Date: Wed, 2 Sep 2026 12:05:35 -0700 Subject: [PATCH 57/79] fix: Refactor mTLS channel reconfiguration logic for lint --- .../google/auth/aio/transport/sessions.py | 16 +++++++++++----- 1 file changed, 11 insertions(+), 5 deletions(-) diff --git a/packages/google-auth/google/auth/aio/transport/sessions.py b/packages/google-auth/google/auth/aio/transport/sessions.py index 809f545f7462..fe218e2bf10a 100644 --- a/packages/google-auth/google/auth/aio/transport/sessions.py +++ b/packages/google-auth/google/auth/aio/transport/sessions.py @@ -375,8 +375,8 @@ async def _recover_auth_state(): cached_fingerprint, current_cert_fingerprint, ) = await mtls.check_parameters_for_unauthorized_response( - self._client_cert_callback, - self._cached_cert, + self._client_cert_callback, + self._cached_cert, ) except ( exceptions.ClientCertError, @@ -393,7 +393,8 @@ async def _recover_auth_state(): else: if ( current_cert_fingerprint is not None - and cached_fingerprint != current_cert_fingerprint + and cached_fingerprint + != current_cert_fingerprint ): saved_callback = self._client_cert_callback try: @@ -407,7 +408,10 @@ async def _recover_auth_state(): ): self._mtls_init_task = None await self.configure_mtls_channel( - lambda: (call_cert_bytes, call_key_bytes) + lambda: ( + call_cert_bytes, + call_key_bytes, + ) ) except Exception as e: _LOGGER.error( @@ -425,7 +429,9 @@ async def _recover_auth_state(): "Failed to reconfigure mTLS channel" ) from e finally: - self._client_cert_callback = saved_callback + self._client_cert_callback = ( + saved_callback + ) else: _LOGGER.info( "Skipping reconfiguration of mTLS channel because the client" From 9bd8e8ffdfc0ff8ebe569b18b8b6fdd5ebd39deb Mon Sep 17 00:00:00 2001 From: Radhika Agrawal Date: Wed, 2 Sep 2026 19:28:26 +0000 Subject: [PATCH 58/79] fix: Fix lint errors Signed-off-by: Radhika Agrawal --- packages/google-auth/google/auth/aio/transport/mtls.py | 1 - 1 file changed, 1 deletion(-) diff --git a/packages/google-auth/google/auth/aio/transport/mtls.py b/packages/google-auth/google/auth/aio/transport/mtls.py index b13787476dec..654385a9515b 100644 --- a/packages/google-auth/google/auth/aio/transport/mtls.py +++ b/packages/google-auth/google/auth/aio/transport/mtls.py @@ -23,7 +23,6 @@ from typing import Optional from google.auth import _agent_identity_utils, exceptions -from google.auth.transport import _mtls_helper from google.auth.transport._mtls_helper import secure_cert_key_paths import google.auth.transport.mtls From e8ee7b640a9369f10800f7b2e53dbd8ff154bc52 Mon Sep 17 00:00:00 2001 From: Radhika Agrawal Date: Wed, 2 Sep 2026 19:52:15 +0000 Subject: [PATCH 59/79] fix: fix the unit tests based on the code changes Signed-off-by: Radhika Agrawal --- .../tests/transport/aio/test_sessions_mtls.py | 21 ++++--------------- 1 file changed, 4 insertions(+), 17 deletions(-) diff --git a/packages/google-auth/tests/transport/aio/test_sessions_mtls.py b/packages/google-auth/tests/transport/aio/test_sessions_mtls.py index 34c1f27f0ebc..67e91101670a 100644 --- a/packages/google-auth/tests/transport/aio/test_sessions_mtls.py +++ b/packages/google-auth/tests/transport/aio/test_sessions_mtls.py @@ -527,13 +527,10 @@ async def test_cert_rotation_lock_contention(self): new_key = b"new_key" async def mock_configure_mtls_channel(*args, **kwargs): - # Introduce a small delay to guarantee lock contention from asyncio.gather tasks await asyncio.sleep(0.01) - # Simulate the channel update to ensure following tasks take the skip branch session._cached_cert = new_cert - def mock_check_side_effect(cached_cert, callback): - # Return new fingerprints on first check, but matching fingerprints on retries + async def mock_check_side_effect(callback, cached_cert): if cached_cert == b"old_cert": return (new_cert, new_key, b"old_fp", b"new_fp") return (new_cert, new_key, b"new_fp", b"new_fp") @@ -547,7 +544,6 @@ def mock_check_side_effect(cached_cert, callback): mock_check.side_effect = mock_check_side_effect mock_conf.side_effect = mock_configure_mtls_channel - # Launch multiple concurrent requests triggering 401s tasks = [ session.request("GET", "https://pubsub.mtls.googleapis.com/test") for _ in range(3) @@ -557,7 +553,6 @@ def mock_check_side_effect(cached_cert, callback): for resp in responses: assert resp == mock_resp_401 - # Confirm the channel is only reconfigured once across all concurrent requests mock_conf.assert_called_once() await session.close() @@ -573,7 +568,6 @@ async def test_cert_rotation_lock_contention_no_cert_change(self): mock_resp_200 = mock.Mock() mock_resp_200.status_code = http_client.OK - # Return 401 for the first 3 requests, then 200 for the retries. mock_auth_req = mock.AsyncMock( side_effect=[mock_resp_401] * 3 + [mock_resp_200] * 3 ) @@ -584,12 +578,9 @@ async def test_cert_rotation_lock_contention_no_cert_change(self): session._is_mtls = True session._cached_cert = b"old_cert" - def mock_check_side_effect(cached_cert, callback): - import time - - # Introduce a small delay to guarantee lock contention from asyncio.gather tasks - time.sleep(0.01) - # Return matching fingerprints to skip reconfiguration + # FIX: async def + await asyncio.sleep to allow concurrent lock contention + async def mock_check_side_effect(callback, cached_cert): + await asyncio.sleep(0.01) return (b"old_cert", b"old_key", b"old_fp", b"old_fp") with mock.patch( @@ -600,7 +591,6 @@ def mock_check_side_effect(cached_cert, callback): ) as mock_conf: mock_check.side_effect = mock_check_side_effect - # Launch multiple concurrent requests triggering 401s tasks = [ session.request("GET", "https://pubsub.mtls.googleapis.com/test") for _ in range(3) @@ -610,11 +600,8 @@ def mock_check_side_effect(cached_cert, callback): for resp in responses: assert resp == mock_resp_200 - # Confirm that the cert parameters were only checked once across the concurrent requests mock_check.assert_called_once() - # Because fingerprints match, reconfiguration should not happen mock_conf.assert_not_called() - # Credentials should still be refreshed for each task as they fall back to normal refresh logic assert mock_creds.refresh.call_count == 3 await session.close() From 7379d8a8bb558febd346a715b7c674833cb50995 Mon Sep 17 00:00:00 2001 From: agrawalradhika-cell Date: Wed, 2 Sep 2026 14:18:14 -0700 Subject: [PATCH 60/79] fix: Log when credentials do not implement refresh method Added debug logging for unimplemented refresh method in credentials. --- packages/google-auth/google/auth/aio/transport/sessions.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/packages/google-auth/google/auth/aio/transport/sessions.py b/packages/google-auth/google/auth/aio/transport/sessions.py index fe218e2bf10a..019e707d8c43 100644 --- a/packages/google-auth/google/auth/aio/transport/sessions.py +++ b/packages/google-auth/google/auth/aio/transport/sessions.py @@ -444,6 +444,8 @@ async def _recover_auth_state(): return response try: await self._credentials.refresh(self._auth_request) + except NotImplementedError: + _LOGGER.debug("Credentials do not implement refresh().") except ( exceptions.RefreshError, getattr(exceptions, "InvalidOperation", Exception), @@ -454,7 +456,7 @@ async def _recover_auth_state(): ) return response - # Return None to explicitly signal successful recovery + # Return None to explicitly signal successful recovery & trigger retry if needed return None async with timeout_guard(remaining_time) as auth_with_timeout: From 48cf5d214a6631cae298e8b91550454368dfbee6 Mon Sep 17 00:00:00 2001 From: agrawalradhika-cell Date: Wed, 2 Sep 2026 19:48:48 -0700 Subject: [PATCH 61/79] chore: Refactor authentication retry logic in sessions.py --- .../google/auth/aio/transport/sessions.py | 258 +++++++++--------- 1 file changed, 128 insertions(+), 130 deletions(-) diff --git a/packages/google-auth/google/auth/aio/transport/sessions.py b/packages/google-auth/google/auth/aio/transport/sessions.py index 019e707d8c43..d5faaea74ca3 100644 --- a/packages/google-auth/google/auth/aio/transport/sessions.py +++ b/packages/google-auth/google/auth/aio/transport/sessions.py @@ -329,150 +329,149 @@ async def request( if response.status_code == http_client.UNAUTHORIZED: if _auth_retry_count < 2: - if max_allowed_time is not None: - elapsed = time.monotonic() - start_time - remaining_time = max(0.0, max_allowed_time - elapsed) - if remaining_time == 0.0: - raise google.auth.exceptions.TimeoutError( - "Timeout exceeded before credential refresh could begin" + try: + if max_allowed_time is not None: + elapsed = time.monotonic() - start_time + remaining_time = max(0.0, max_allowed_time - elapsed) + if remaining_time == 0.0: + raise google.auth.exceptions.TimeoutError( + "Timeout exceeded before credential refresh could begin" + ) + else: + remaining_time = None + is_streaming = data is not None and ( + isinstance( + data, (collections.abc.Iterator, collections.abc.AsyncIterable) ) - else: - remaining_time = None - is_streaming = data is not None and ( - isinstance( - data, (collections.abc.Iterator, collections.abc.AsyncIterable) + or hasattr(data, "read") ) - or hasattr(data, "read") - ) - - async def _recover_auth_state(): - is_mtls_endpoint = False - if getattr(self, "is_mtls", False): - hostname = urllib.parse.urlsplit(url).hostname - if hostname: - is_mtls_endpoint = any( - hostname == prefix or hostname.endswith("." + prefix) - for prefix in _MTLS_URL_PREFIXES - ) - # Snapshot the stale certificate state BEFORE acquiring the lock. - # This represents the cert that caused the 401 rejection. - if is_mtls_endpoint: - if self._mtls_rotation_lock is None: - self._mtls_rotation_lock = asyncio.Lock() - # Snapshot the counter state BEFORE acquiring the lock. - check_counter_at_error = self._mtls_check_counter - - async with self._mtls_rotation_lock: - # Check if another coroutine already reconfigured mTLS or - # ran the validation check. - if self._mtls_check_counter > check_counter_at_error: - pass - else: - try: - ( - call_cert_bytes, - call_key_bytes, - cached_fingerprint, - current_cert_fingerprint, - ) = await mtls.check_parameters_for_unauthorized_response( - self._client_cert_callback, - self._cached_cert, - ) - except ( - exceptions.ClientCertError, - exceptions.MutualTLSChannelError, - OSError, - ValueError, - ImportError, - ) as e: - _LOGGER.warning( - "Failed to check client certificate parameters: %s. Proceeding with original response.", - e, - ) - return response + async def _recover_auth_state(): + is_mtls_endpoint = False + if getattr(self, "is_mtls", False): + hostname = urllib.parse.urlsplit(url).hostname + if hostname: + is_mtls_endpoint = any( + hostname == prefix or hostname.endswith("." + prefix) + for prefix in _MTLS_URL_PREFIXES + ) + # Snapshot the stale certificate state BEFORE acquiring the lock. + # This represents the cert that caused the 401 rejection. + if is_mtls_endpoint: + if self._mtls_rotation_lock is None: + self._mtls_rotation_lock = asyncio.Lock() + # Snapshot the counter state BEFORE acquiring the lock. + check_counter_at_error = self._mtls_check_counter + async with self._mtls_rotation_lock: + # Check if another coroutine already reconfigured mTLS or + # ran the validation check. + if self._mtls_check_counter > check_counter_at_error: + pass else: - if ( - current_cert_fingerprint is not None - and cached_fingerprint - != current_cert_fingerprint - ): - saved_callback = self._client_cert_callback - try: - _LOGGER.info( - "Client certificate has changed, reconfiguring mTLS " - "channel." - ) - if ( - self._mtls_init_task - and self._mtls_init_task.done() - ): - self._mtls_init_task = None - await self.configure_mtls_channel( - lambda: ( - call_cert_bytes, - call_key_bytes, + try: + ( + call_cert_bytes, + call_key_bytes, + cached_fingerprint, + current_cert_fingerprint, + ) = await mtls.check_parameters_for_unauthorized_response( + self._client_cert_callback, + self._cached_cert, + ) + except ( + exceptions.ClientCertError, + exceptions.MutualTLSChannelError, + OSError, + ValueError, + ImportError, + ) as e: + _LOGGER.warning( + "Failed to check client certificate parameters: %s. Proceeding with original response.", + e, + ) + return response + else: + if ( + current_cert_fingerprint is not None + and cached_fingerprint + != current_cert_fingerprint + ): + saved_callback = self._client_cert_callback + try: + _LOGGER.info( + "Client certificate has changed, reconfiguring mTLS " + "channel." ) - ) - except Exception as e: - _LOGGER.error( - "Failed to reconfigure mTLS channel: %s", - e, - ) - if hasattr(response, "close"): - if asyncio.iscoroutinefunction( - response.close + if ( + self._mtls_init_task + and self._mtls_init_task.done() ): - await response.close() - else: - response.close() - raise exceptions.MutualTLSChannelError( - "Failed to reconfigure mTLS channel" - ) from e - finally: - self._client_cert_callback = ( - saved_callback + self._mtls_init_task = None + await self.configure_mtls_channel( + lambda: ( + call_cert_bytes, + call_key_bytes, + ) + ) + except Exception as e: + _LOGGER.error( + "Failed to reconfigure mTLS channel: %s", + e, + ) + raise exceptions.MutualTLSChannelError( + "Failed to reconfigure mTLS channel" + ) from e + finally: + self._client_cert_callback = ( + saved_callback + ) + else: + _LOGGER.info( + "Skipping reconfiguration of mTLS channel because the client" + " certificate has not changed." ) - else: - _LOGGER.info( - "Skipping reconfiguration of mTLS channel because the client" - " certificate has not changed." - ) - finally: - # Always increment so waiting tasks skip the check block - self._mtls_check_counter += 1 - if is_streaming: - return response - try: - await self._credentials.refresh(self._auth_request) - except NotImplementedError: - _LOGGER.debug("Credentials do not implement refresh().") - except ( - exceptions.RefreshError, - getattr(exceptions, "InvalidOperation", Exception), - ) as e: - _LOGGER.debug( - "Credential refresh failed, returning 401 response. Error: %s", - e, + finally: + # Always increment so waiting tasks skip the check block + self._mtls_check_counter += 1 + try: + await self._credentials.refresh(self._auth_request) + except NotImplementedError: + _LOGGER.debug("Credentials do not implement refresh().") + except ( + exceptions.RefreshError, + getattr(exceptions, "InvalidOperation", Exception), + ) as e: + _LOGGER.debug( + "Credential refresh failed, returning 401 response. Error: %s", + e, + ) + return response + + if is_streaming: + return response + # Return None to explicitly signal successful recovery & trigger retry if needed + return None + async with timeout_guard(remaining_time) as auth_with_timeout: + early_return_response = await auth_with_timeout( + _recover_auth_state() ) - return response - - # Return None to explicitly signal successful recovery & trigger retry if needed - return None - - async with timeout_guard(remaining_time) as auth_with_timeout: - early_return_response = await auth_with_timeout( - _recover_auth_state() - ) + except (Exception, asyncio.CancelledError): + if hasattr(response, "close"): + try: + if asyncio.iscoroutinefunction(response.close): + await response.close() + else: + response.close() + except Exception: + pass + raise # If it returned a response (meaning streaming or error), bail out if early_return_response is not None: return early_return_response - if hasattr(response, "close"): if asyncio.iscoroutinefunction(response.close): await response.close() else: response.close() - if max_allowed_time is not None: remaining_time = max( 0.0, max_allowed_time - (time.monotonic() - start_time) @@ -481,7 +480,6 @@ async def _recover_auth_state(): raise google.auth.exceptions.TimeoutError( "Timeout exceeded before retrying the request" ) - kwargs["_auth_retry_count"] = _auth_retry_count + 1 return await self.request( method, From 8b5d35d12aeba6a5a723ed7beebbbf41fa54488f Mon Sep 17 00:00:00 2001 From: agrawalradhika-cell Date: Wed, 2 Sep 2026 20:15:13 -0700 Subject: [PATCH 62/79] fix: Implement concurrent credential refresh management Added refresh lock and counter to manage concurrent credential refreshes. --- .../google/auth/aio/transport/sessions.py | 43 +++++++++++++------ 1 file changed, 29 insertions(+), 14 deletions(-) diff --git a/packages/google-auth/google/auth/aio/transport/sessions.py b/packages/google-auth/google/auth/aio/transport/sessions.py index d5faaea74ca3..184ac6eb63e7 100644 --- a/packages/google-auth/google/auth/aio/transport/sessions.py +++ b/packages/google-auth/google/auth/aio/transport/sessions.py @@ -161,6 +161,9 @@ def __init__( self._auth_request = _auth_request self._mtls_rotation_lock: Optional[asyncio.Lock] = None self._mtls_check_counter = 0 + self._refresh_lock: Optional[asyncio.Lock] = None + self._refresh_counter = 0 + async def configure_mtls_channel(self, client_cert_callback=None): """Configure the client certificate and key for SSL connection. @@ -429,22 +432,34 @@ async def _recover_auth_state(): "Skipping reconfiguration of mTLS channel because the client" " certificate has not changed." ) - finally: # Always increment so waiting tasks skip the check block self._mtls_check_counter += 1 - try: - await self._credentials.refresh(self._auth_request) - except NotImplementedError: - _LOGGER.debug("Credentials do not implement refresh().") - except ( - exceptions.RefreshError, - getattr(exceptions, "InvalidOperation", Exception), - ) as e: - _LOGGER.debug( - "Credential refresh failed, returning 401 response. Error: %s", - e, - ) - return response + if self._refresh_lock is None: + self._refresh_lock = asyncio.Lock() + refresh_counter_at_error = self._refresh_counter + + async with self._refresh_lock: + # Check if another task already refreshed credentials while we were waiting + if self._refresh_counter > refresh_counter_at_error: + _LOGGER.debug( + "Credentials were already refreshed by a concurrent task. Skipping duplicate refresh." + ) + else: + try: + await self._credentials.refresh(self._auth_request) + except NotImplementedError: + _LOGGER.debug("Credentials do not implement refresh().") + except ( + exceptions.RefreshError, + getattr(exceptions, "InvalidOperation", Exception), + ) as e: + _LOGGER.debug( + "Credential refresh failed, returning 401 response. Error: %s", + e, + ) + return response + else: + self._refresh_counter += 1 if is_streaming: return response From c7746a4824312d5c787a4b90a92ae4cd8cb9e392 Mon Sep 17 00:00:00 2001 From: agrawalradhika-cell Date: Wed, 2 Sep 2026 20:39:34 -0700 Subject: [PATCH 63/79] fix: Limit old auth requests to 2 and close oldest requests Limit the number of old authentication requests to 2 and ensure proper closure of the oldest requests. --- .../google-auth/google/auth/aio/transport/sessions.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/packages/google-auth/google/auth/aio/transport/sessions.py b/packages/google-auth/google/auth/aio/transport/sessions.py index 184ac6eb63e7..5b639341971b 100644 --- a/packages/google-auth/google/auth/aio/transport/sessions.py +++ b/packages/google-auth/google/auth/aio/transport/sessions.py @@ -221,6 +221,17 @@ async def _do_configure(): old_auth_request = self._auth_request self._auth_request = AiohttpRequest(session=new_session) + while len(self._old_auth_requests) >= 2: + oldest_auth_request = self._old_auth_requests.pop(0) + try: + if hasattr(oldest_auth_request, "close"): + if asyncio.iscoroutinefunction(oldest_auth_request.close): + await oldest_auth_request.close() + else: + oldest_auth_request.close() + except Exception: + pass + self._old_auth_requests.append(old_auth_request) else: From d5c674696909907c86089e8947c285a7bed3a6fe Mon Sep 17 00:00:00 2001 From: agrawalradhika-cell Date: Wed, 2 Sep 2026 20:49:08 -0700 Subject: [PATCH 64/79] fix: Enhance tests for 401 response handling Added tests for handling 401 responses with timeout and cancellation scenarios. --- .../tests/transport/aio/test_sessions_mtls.py | 135 +++++++++++++++++- 1 file changed, 129 insertions(+), 6 deletions(-) diff --git a/packages/google-auth/tests/transport/aio/test_sessions_mtls.py b/packages/google-auth/tests/transport/aio/test_sessions_mtls.py index 67e91101670a..d1b3b581ff37 100644 --- a/packages/google-auth/tests/transport/aio/test_sessions_mtls.py +++ b/packages/google-auth/tests/transport/aio/test_sessions_mtls.py @@ -22,6 +22,7 @@ import pytest from google.auth import exceptions +from google.auth.exceptions import TimeoutError from google.auth.aio import credentials from google.auth.aio import transport from google.auth.aio.transport import sessions @@ -423,10 +424,17 @@ async def test_cert_rotation_check_params_fails(self): async def test_no_cert_rotation_when_cert_matches_and_mtls_enabled(self): mock_creds = mock.AsyncMock(spec=credentials.Credentials) mock_creds.before_request = mock.AsyncMock(return_value=None) + mock_creds.refresh = mock.AsyncMock(return_value=None) - mock_resp = mock.Mock() - mock_resp.status_code = http_client.UNAUTHORIZED - mock_auth_req = mock.AsyncMock(return_value=mock_resp) + mock_resp_401 = mock.Mock() + mock_resp_401.status_code = http_client.UNAUTHORIZED + mock_resp_401.close = mock.AsyncMock() + + mock_resp_200 = mock.Mock() + mock_resp_200.status_code = http_client.OK + + # 401 on initial request, 200 on retry after refresh + mock_auth_req = mock.AsyncMock(side_effect=[mock_resp_401, mock_resp_200]) session = sessions.AsyncAuthorizedSession( mock_creds, auth_request=mock_auth_req @@ -443,16 +451,19 @@ async def test_no_cert_rotation_when_cert_matches_and_mtls_enabled(self): ) as mock_check, mock.patch.object( session, "configure_mtls_channel", new_callable=mock.AsyncMock ) as mock_conf: - # Matching fingerprints mean no layout rotation is needed + # Matching fingerprints mean no mTLS rotation is needed mock_check.return_value = (new_cert, new_key, b"old_fp", b"old_fp") resp = await session.request( "GET", "https://pubsub.mtls.googleapis.com/test" ) - assert resp == mock_resp - assert mock_check.call_count >= 1 + assert resp == mock_resp_200 + mock_check.assert_called_once() mock_conf.assert_not_called() + mock_creds.refresh.assert_called_once() + assert mock_auth_req.call_count == 2 + mock_resp_401.close.assert_awaited_once() await session.close() @@ -793,3 +804,115 @@ async def test_session_close_cleans_old_auth_requests(self): mock_old_req_3_fails.close.assert_called_once() # Ensure the list was cleared assert len(session._old_auth_requests) == 0 + + @pytest.mark.asyncio + async def test_request_401_closes_response_on_timeout_before_recovery(session): + """Verifies that response is closed and TimeoutError raised when max_allowed_time elapses before refresh.""" + mock_resp_401 = mock.AsyncMock() + mock_resp_401.status_code = 401 + + session._auth_request = mock.AsyncMock(return_value=mock_resp_401) + + # Set max_allowed_time to 0 so remaining_time == 0.0 immediately + with pytest.raises( + TimeoutError, match="Timeout exceeded before credential refresh" + ): + await session.request("GET", "https://example.com", max_allowed_time=0.0) + + mock_resp_401.close.assert_awaited_once() + + @pytest.mark.asyncio + async def test_request_401_closes_response_on_timeout_during_recovery(session): + """Verifies that response is closed when auth_with_timeout times out during _recover_auth_state.""" + mock_resp_401 = mock.AsyncMock() + mock_resp_401.status_code = 401 + session._auth_request = mock.AsyncMock(return_value=mock_resp_401) + + async def slow_refresh(*args, **kwargs): + await asyncio.sleep(10) + + session._credentials.refresh = mock.AsyncMock(side_effect=slow_refresh) + + # Set short max_allowed_time so timeout_guard triggers during credentials.refresh + with pytest.raises(TimeoutError): + await session.request("GET", "https://example.com", max_allowed_time=0.01) + + mock_resp_401.close.assert_awaited_once() + + @pytest.mark.asyncio + async def test_request_401_closes_response_on_cancellation(session): + """Verifies that response is closed and CancelledError propagated if task is cancelled during recovery.""" + mock_resp_401 = mock.AsyncMock() + mock_resp_401.status_code = 401 + session._auth_request = mock.AsyncMock(return_value=mock_resp_401) + + refresh_started = asyncio.Event() + + async def cancel_on_refresh(*args, **kwargs): + refresh_started.set() + await asyncio.sleep(10) + + session._credentials.refresh = mock.AsyncMock(side_effect=cancel_on_refresh) + + task = asyncio.create_task(session.request("GET", "https://example.com")) + await refresh_started.wait() + task.cancel() + + with pytest.raises(asyncio.CancelledError): + await task + + mock_resp_401.close.assert_awaited_once() + + @pytest.mark.asyncio + async def test_request_401_streaming_refreshes_creds_and_returns_open_response( + session, + ): + """Verifies that streaming requests refresh credentials but return the unclosed response to the caller.""" + mock_resp_401 = mock.AsyncMock() + mock_resp_401.status_code = 401 + session._auth_request = mock.AsyncMock(return_value=mock_resp_401) + session._credentials.refresh = mock.AsyncMock() + + # Generator simulating streaming body + streaming_data = (chunk for chunk in [b"chunk1", b"chunk2"]) + + response = await session.request( + "POST", "https://example.com", data=streaming_data + ) + + assert response is mock_resp_401 + session._credentials.refresh.assert_awaited_once() + # Response must NOT be closed so caller can consume the body + mock_resp_401.close.assert_not_called() + + @pytest.mark.asyncio + async def test_request_401_concurrent_refreshes_are_deduplicated(session): + """Verifies that concurrent 401s execute only one credentials.refresh call.""" + mock_resp_401 = mock.AsyncMock() + mock_resp_401.status_code = 401 + mock_resp_200 = mock.AsyncMock() + mock_resp_200.status_code = 200 + + # Return 401 on initial calls, then 200 on retries + session._auth_request = mock.AsyncMock( + side_effect=[mock_resp_401, mock_resp_401, mock_resp_200, mock_resp_200] + ) + + refresh_count = 0 + + async def slow_refresh(*args, **kwargs): + nonlocal refresh_count + refresh_count += 1 + await asyncio.sleep(0.05) + + session._credentials.refresh = mock.AsyncMock(side_effect=slow_refresh) + + # Launch two concurrent requests encountering 401 + results = await asyncio.gather( + session.request("GET", "https://example.com/1"), + session.request("GET", "https://example.com/2"), + ) + + assert all(r.status_code == 200 for r in results) + # Refresh should only have been called ONCE across both requests + assert refresh_count == 1 From 83095eaa54fe5e73f2209a3198ceac3ee3193fd0 Mon Sep 17 00:00:00 2001 From: Radhika Agrawal Date: Thu, 3 Sep 2026 05:37:02 +0000 Subject: [PATCH 65/79] Fix: Fix the lint and unit tests Signed-off-by: Radhika Agrawal --- .../google/auth/aio/transport/sessions.py | 44 +++--- .../tests/transport/aio/test_sessions_mtls.py | 125 ++++++++++-------- 2 files changed, 98 insertions(+), 71 deletions(-) diff --git a/packages/google-auth/google/auth/aio/transport/sessions.py b/packages/google-auth/google/auth/aio/transport/sessions.py index 5b639341971b..26576527540a 100644 --- a/packages/google-auth/google/auth/aio/transport/sessions.py +++ b/packages/google-auth/google/auth/aio/transport/sessions.py @@ -164,7 +164,6 @@ def __init__( self._refresh_lock: Optional[asyncio.Lock] = None self._refresh_counter = 0 - async def configure_mtls_channel(self, client_cert_callback=None): """Configure the client certificate and key for SSL connection. @@ -221,16 +220,18 @@ async def _do_configure(): old_auth_request = self._auth_request self._auth_request = AiohttpRequest(session=new_session) - while len(self._old_auth_requests) >= 2: - oldest_auth_request = self._old_auth_requests.pop(0) - try: - if hasattr(oldest_auth_request, "close"): - if asyncio.iscoroutinefunction(oldest_auth_request.close): - await oldest_auth_request.close() - else: - oldest_auth_request.close() - except Exception: - pass + while len(self._old_auth_requests) >= 2: + oldest_auth_request = self._old_auth_requests.pop(0) + try: + if hasattr(oldest_auth_request, "close"): + if asyncio.iscoroutinefunction( + oldest_auth_request.close + ): + await oldest_auth_request.close() + else: + oldest_auth_request.close() + except Exception: + pass self._old_auth_requests.append(old_auth_request) @@ -355,17 +356,20 @@ async def request( remaining_time = None is_streaming = data is not None and ( isinstance( - data, (collections.abc.Iterator, collections.abc.AsyncIterable) + data, + (collections.abc.Iterator, collections.abc.AsyncIterable), ) or hasattr(data, "read") ) + async def _recover_auth_state(): is_mtls_endpoint = False if getattr(self, "is_mtls", False): hostname = urllib.parse.urlsplit(url).hostname if hostname: is_mtls_endpoint = any( - hostname == prefix or hostname.endswith("." + prefix) + hostname == prefix + or hostname.endswith("." + prefix) for prefix in _MTLS_URL_PREFIXES ) # Snapshot the stale certificate state BEFORE acquiring the lock. @@ -378,7 +382,10 @@ async def _recover_auth_state(): async with self._mtls_rotation_lock: # Check if another coroutine already reconfigured mTLS or # ran the validation check. - if self._mtls_check_counter > check_counter_at_error: + if ( + self._mtls_check_counter + > check_counter_at_error + ): pass else: try: @@ -409,7 +416,9 @@ async def _recover_auth_state(): and cached_fingerprint != current_cert_fingerprint ): - saved_callback = self._client_cert_callback + saved_callback = ( + self._client_cert_callback + ) try: _LOGGER.info( "Client certificate has changed, reconfiguring mTLS " @@ -459,7 +468,9 @@ async def _recover_auth_state(): try: await self._credentials.refresh(self._auth_request) except NotImplementedError: - _LOGGER.debug("Credentials do not implement refresh().") + _LOGGER.debug( + "Credentials do not implement refresh()." + ) except ( exceptions.RefreshError, getattr(exceptions, "InvalidOperation", Exception), @@ -476,6 +487,7 @@ async def _recover_auth_state(): return response # Return None to explicitly signal successful recovery & trigger retry if needed return None + async with timeout_guard(remaining_time) as auth_with_timeout: early_return_response = await auth_with_timeout( _recover_auth_state() diff --git a/packages/google-auth/tests/transport/aio/test_sessions_mtls.py b/packages/google-auth/tests/transport/aio/test_sessions_mtls.py index d1b3b581ff37..1a61c00f3862 100644 --- a/packages/google-auth/tests/transport/aio/test_sessions_mtls.py +++ b/packages/google-auth/tests/transport/aio/test_sessions_mtls.py @@ -22,10 +22,10 @@ import pytest from google.auth import exceptions -from google.auth.exceptions import TimeoutError from google.auth.aio import credentials from google.auth.aio import transport from google.auth.aio.transport import sessions +from google.auth.exceptions import TimeoutError # This is the valid "workload" format the library expects VALID_WORKLOAD_CONFIG = { @@ -690,7 +690,7 @@ def read(self): assert resp == mock_resp mock_conf.assert_called_once() # Because it is streaming, it skips credentials refresh and retry - mock_creds.refresh.assert_not_called() + mock_creds.refresh.assert_called_once() await session.close() @@ -806,45 +806,64 @@ async def test_session_close_cleans_old_auth_requests(self): assert len(session._old_auth_requests) == 0 @pytest.mark.asyncio - async def test_request_401_closes_response_on_timeout_before_recovery(session): - """Verifies that response is closed and TimeoutError raised when max_allowed_time elapses before refresh.""" - mock_resp_401 = mock.AsyncMock() - mock_resp_401.status_code = 401 + async def test_request_401_streaming_refreshes_creds_and_returns_open_response( + self, + ): + """Verifies that streaming requests refresh credentials but return the unclosed response.""" + mock_creds = mock.AsyncMock(spec=credentials.Credentials) + mock_creds.before_request = mock.AsyncMock(return_value=None) + mock_creds.refresh = mock.AsyncMock(return_value=None) - session._auth_request = mock.AsyncMock(return_value=mock_resp_401) + mock_resp_401 = mock.Mock() + mock_resp_401.status_code = http_client.UNAUTHORIZED + mock_resp_401.close = mock.AsyncMock() - # Set max_allowed_time to 0 so remaining_time == 0.0 immediately - with pytest.raises( - TimeoutError, match="Timeout exceeded before credential refresh" - ): - await session.request("GET", "https://example.com", max_allowed_time=0.0) + mock_auth_req = mock.AsyncMock(return_value=mock_resp_401) + session = sessions.AsyncAuthorizedSession( + mock_creds, auth_request=mock_auth_req + ) - mock_resp_401.close.assert_awaited_once() + streaming_data = (chunk for chunk in [b"chunk1", b"chunk2"]) + response = await session.request( + "POST", "https://example.com", data=streaming_data + ) + + assert response == mock_resp_401 + mock_creds.refresh.assert_awaited_once() + mock_resp_401.close.assert_not_called() + await session.close() @pytest.mark.asyncio - async def test_request_401_closes_response_on_timeout_during_recovery(session): + async def test_request_401_closes_response_on_timeout_during_recovery(self): """Verifies that response is closed when auth_with_timeout times out during _recover_auth_state.""" - mock_resp_401 = mock.AsyncMock() - mock_resp_401.status_code = 401 - session._auth_request = mock.AsyncMock(return_value=mock_resp_401) + mock_creds = mock.AsyncMock(spec=credentials.Credentials) + mock_creds.before_request = mock.AsyncMock(return_value=None) async def slow_refresh(*args, **kwargs): await asyncio.sleep(10) - session._credentials.refresh = mock.AsyncMock(side_effect=slow_refresh) + mock_creds.refresh = mock.AsyncMock(side_effect=slow_refresh) + + mock_resp_401 = mock.Mock() + mock_resp_401.status_code = http_client.UNAUTHORIZED + mock_resp_401.close = mock.AsyncMock() + + mock_auth_req = mock.AsyncMock(return_value=mock_resp_401) + session = sessions.AsyncAuthorizedSession( + mock_creds, auth_request=mock_auth_req + ) - # Set short max_allowed_time so timeout_guard triggers during credentials.refresh with pytest.raises(TimeoutError): await session.request("GET", "https://example.com", max_allowed_time=0.01) mock_resp_401.close.assert_awaited_once() + await session.close() @pytest.mark.asyncio - async def test_request_401_closes_response_on_cancellation(session): + async def test_request_401_closes_response_on_cancellation(self): """Verifies that response is closed and CancelledError propagated if task is cancelled during recovery.""" - mock_resp_401 = mock.AsyncMock() - mock_resp_401.status_code = 401 - session._auth_request = mock.AsyncMock(return_value=mock_resp_401) + mock_creds = mock.AsyncMock(spec=credentials.Credentials) + mock_creds.before_request = mock.AsyncMock(return_value=None) refresh_started = asyncio.Event() @@ -852,7 +871,16 @@ async def cancel_on_refresh(*args, **kwargs): refresh_started.set() await asyncio.sleep(10) - session._credentials.refresh = mock.AsyncMock(side_effect=cancel_on_refresh) + mock_creds.refresh = mock.AsyncMock(side_effect=cancel_on_refresh) + + mock_resp_401 = mock.Mock() + mock_resp_401.status_code = http_client.UNAUTHORIZED + mock_resp_401.close = mock.AsyncMock() + + mock_auth_req = mock.AsyncMock(return_value=mock_resp_401) + session = sessions.AsyncAuthorizedSession( + mock_creds, auth_request=mock_auth_req + ) task = asyncio.create_task(session.request("GET", "https://example.com")) await refresh_started.wait() @@ -862,41 +890,29 @@ async def cancel_on_refresh(*args, **kwargs): await task mock_resp_401.close.assert_awaited_once() + await session.close() @pytest.mark.asyncio - async def test_request_401_streaming_refreshes_creds_and_returns_open_response( - session, - ): - """Verifies that streaming requests refresh credentials but return the unclosed response to the caller.""" - mock_resp_401 = mock.AsyncMock() - mock_resp_401.status_code = 401 - session._auth_request = mock.AsyncMock(return_value=mock_resp_401) - session._credentials.refresh = mock.AsyncMock() - - # Generator simulating streaming body - streaming_data = (chunk for chunk in [b"chunk1", b"chunk2"]) - - response = await session.request( - "POST", "https://example.com", data=streaming_data - ) + async def test_request_401_concurrent_refreshes_are_deduplicated(self): + """Verifies that concurrent 401s execute only one credentials.refresh call.""" + mock_creds = mock.AsyncMock(spec=credentials.Credentials) + mock_creds.before_request = mock.AsyncMock(return_value=None) - assert response is mock_resp_401 - session._credentials.refresh.assert_awaited_once() - # Response must NOT be closed so caller can consume the body - mock_resp_401.close.assert_not_called() + mock_resp_401 = mock.Mock() + mock_resp_401.status_code = http_client.UNAUTHORIZED + mock_resp_401.close = mock.AsyncMock() - @pytest.mark.asyncio - async def test_request_401_concurrent_refreshes_are_deduplicated(session): - """Verifies that concurrent 401s execute only one credentials.refresh call.""" - mock_resp_401 = mock.AsyncMock() - mock_resp_401.status_code = 401 - mock_resp_200 = mock.AsyncMock() - mock_resp_200.status_code = 200 + mock_resp_200 = mock.Mock() + mock_resp_200.status_code = http_client.OK + mock_resp_200.close = mock.AsyncMock() - # Return 401 on initial calls, then 200 on retries - session._auth_request = mock.AsyncMock( + # Both concurrent requests get 401 initially, then 200 on retry + mock_auth_req = mock.AsyncMock( side_effect=[mock_resp_401, mock_resp_401, mock_resp_200, mock_resp_200] ) + session = sessions.AsyncAuthorizedSession( + mock_creds, auth_request=mock_auth_req + ) refresh_count = 0 @@ -905,14 +921,13 @@ async def slow_refresh(*args, **kwargs): refresh_count += 1 await asyncio.sleep(0.05) - session._credentials.refresh = mock.AsyncMock(side_effect=slow_refresh) + mock_creds.refresh = mock.AsyncMock(side_effect=slow_refresh) - # Launch two concurrent requests encountering 401 results = await asyncio.gather( session.request("GET", "https://example.com/1"), session.request("GET", "https://example.com/2"), ) assert all(r.status_code == 200 for r in results) - # Refresh should only have been called ONCE across both requests assert refresh_count == 1 + await session.close() From da00b4093c3a4a53978d87389716f0e346c6efba Mon Sep 17 00:00:00 2001 From: agrawalradhika-cell Date: Thu, 3 Sep 2026 09:52:37 -0700 Subject: [PATCH 66/79] fix: Add .p.googleapis.com to MTLS URL prefixes --- packages/google-auth/google/auth/aio/transport/sessions.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/packages/google-auth/google/auth/aio/transport/sessions.py b/packages/google-auth/google/auth/aio/transport/sessions.py index 26576527540a..6e8b9817b0fe 100644 --- a/packages/google-auth/google/auth/aio/transport/sessions.py +++ b/packages/google-auth/google/auth/aio/transport/sessions.py @@ -42,8 +42,7 @@ ClientTimeout = None _LOGGER = logging.getLogger(__name__) -_MTLS_URL_PREFIXES = ["mtls.googleapis.com", "mtls.sandbox.googleapis.com"] - +_MTLS_URL_PREFIXES = ["mtls.googleapis.com", "mtls.sandbox.googleapis.com", ".p.googleapis.com"] # Tracks the internal aiohttp installation and usage try: @@ -364,6 +363,7 @@ async def request( async def _recover_auth_state(): is_mtls_endpoint = False + refresh_counter_at_error = self._refresh_counter if getattr(self, "is_mtls", False): hostname = urllib.parse.urlsplit(url).hostname if hostname: @@ -456,7 +456,6 @@ async def _recover_auth_state(): self._mtls_check_counter += 1 if self._refresh_lock is None: self._refresh_lock = asyncio.Lock() - refresh_counter_at_error = self._refresh_counter async with self._refresh_lock: # Check if another task already refreshed credentials while we were waiting From a355bc7b4a647950888130e82fb31f7db937bac1 Mon Sep 17 00:00:00 2001 From: agrawalradhika-cell Date: Thu, 3 Sep 2026 10:00:36 -0700 Subject: [PATCH 67/79] fix: Refactor close method calls to handle awaitables --- .../google/auth/aio/transport/sessions.py | 47 ++++++++++--------- 1 file changed, 26 insertions(+), 21 deletions(-) diff --git a/packages/google-auth/google/auth/aio/transport/sessions.py b/packages/google-auth/google/auth/aio/transport/sessions.py index 6e8b9817b0fe..117fd513bdc8 100644 --- a/packages/google-auth/google/auth/aio/transport/sessions.py +++ b/packages/google-auth/google/auth/aio/transport/sessions.py @@ -17,6 +17,7 @@ from contextlib import asynccontextmanager import functools import http.client as http_client +import inspect import logging import time from typing import Mapping, Optional, TYPE_CHECKING, Union @@ -220,17 +221,14 @@ async def _do_configure(): old_auth_request = self._auth_request self._auth_request = AiohttpRequest(session=new_session) while len(self._old_auth_requests) >= 2: - oldest_auth_request = self._old_auth_requests.pop(0) - try: - if hasattr(oldest_auth_request, "close"): - if asyncio.iscoroutinefunction( - oldest_auth_request.close - ): - await oldest_auth_request.close() - else: - oldest_auth_request.close() - except Exception: - pass + oldest_auth_request = self._old_auth_requests.pop(0) + try: + if hasattr(oldest_auth_request, "close"): + res = oldest_auth_request.close() + if inspect.isawaitable(res): + await res + except Exception: + pass self._old_auth_requests.append(old_auth_request) @@ -494,10 +492,9 @@ async def _recover_auth_state(): except (Exception, asyncio.CancelledError): if hasattr(response, "close"): try: - if asyncio.iscoroutinefunction(response.close): - await response.close() - else: - response.close() + res = response.close() + if inspect.isawaitable(res): + await res except Exception: pass raise @@ -505,10 +502,12 @@ async def _recover_auth_state(): if early_return_response is not None: return early_return_response if hasattr(response, "close"): - if asyncio.iscoroutinefunction(response.close): - await response.close() - else: - response.close() + try: + res = response.close() + if inspect.isawaitable(res): + await res + except Exception: + pass if max_allowed_time is not None: remaining_time = max( 0.0, max_allowed_time - (time.monotonic() - start_time) @@ -811,11 +810,17 @@ async def close(self) -> None: except asyncio.CancelledError: pass try: - await self._auth_request.close() + if hasattr(self._auth_request, "close"): + res = self._auth_request.close() + if inspect.isawaitable(res): + await res finally: for old_request in self._old_auth_requests: try: - await old_request.close() + if hasattr(old_request, "close"): + res = old_request.close() + if inspect.isawaitable(res): + await res except Exception: pass self._old_auth_requests.clear() From 0ae183f74c9a1dabedaee18001f1a3c9aed53b9c Mon Sep 17 00:00:00 2001 From: agrawalradhika-cell Date: Thu, 3 Sep 2026 10:04:30 -0700 Subject: [PATCH 68/79] fix: Change parameter order in check_parameters_for_unauthorized_response Reordered parameters in check_parameters_for_unauthorized_response function. --- packages/google-auth/google/auth/aio/transport/mtls.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/google-auth/google/auth/aio/transport/mtls.py b/packages/google-auth/google/auth/aio/transport/mtls.py index 654385a9515b..1d76fb744d3d 100644 --- a/packages/google-auth/google/auth/aio/transport/mtls.py +++ b/packages/google-auth/google/auth/aio/transport/mtls.py @@ -179,7 +179,7 @@ async def get_client_cert_and_key(client_cert_callback=None): return has_cert, cert, key -async def check_parameters_for_unauthorized_response(client_cert_callback, cached_cert): +async def check_parameters_for_unauthorized_response(cached_cert, client_cert_callback): """Async helper to retrieve certs and compute fingerprints for mTLS rotation.""" is_mtls, call_cert_bytes, call_key_bytes = await get_client_cert_and_key( client_cert_callback From c976c76b363923f552e5e4e227e32b793c934a9d Mon Sep 17 00:00:00 2001 From: agrawalradhika-cell Date: Thu, 3 Sep 2026 10:07:44 -0700 Subject: [PATCH 69/79] fix: Fix client cert callback assignment in sessions.py --- packages/google-auth/google/auth/aio/transport/sessions.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/google-auth/google/auth/aio/transport/sessions.py b/packages/google-auth/google/auth/aio/transport/sessions.py index 117fd513bdc8..dff59da9b112 100644 --- a/packages/google-auth/google/auth/aio/transport/sessions.py +++ b/packages/google-auth/google/auth/aio/transport/sessions.py @@ -393,8 +393,8 @@ async def _recover_auth_state(): cached_fingerprint, current_cert_fingerprint, ) = await mtls.check_parameters_for_unauthorized_response( - self._client_cert_callback, self._cached_cert, + self._client_cert_callback, ) except ( exceptions.ClientCertError, From 3390024534f53316fe9e3d54e2c7ede788aef1c0 Mon Sep 17 00:00:00 2001 From: agrawalradhika-cell Date: Thu, 3 Sep 2026 10:12:29 -0700 Subject: [PATCH 70/79] fix: Update the tests based on chnages in sessions.py --- .../tests/transport/aio/test_sessions_mtls.py | 1685 +++++++++-------- 1 file changed, 881 insertions(+), 804 deletions(-) diff --git a/packages/google-auth/tests/transport/aio/test_sessions_mtls.py b/packages/google-auth/tests/transport/aio/test_sessions_mtls.py index 1a61c00f3862..e1bfddbcc235 100644 --- a/packages/google-auth/tests/transport/aio/test_sessions_mtls.py +++ b/packages/google-auth/tests/transport/aio/test_sessions_mtls.py @@ -19,915 +19,992 @@ import ssl from unittest import mock -import pytest - from google.auth import exceptions from google.auth.aio import credentials from google.auth.aio import transport from google.auth.aio.transport import sessions from google.auth.exceptions import TimeoutError +import pytest # This is the valid "workload" format the library expects VALID_WORKLOAD_CONFIG = { "version": 1, "cert_configs": { - "workload": {"cert_path": "/tmp/mock_cert.pem", "key_path": "/tmp/mock_key.pem"} + "workload": { + "cert_path": "/tmp/mock_cert.pem", + "key_path": "/tmp/mock_key.pem", + } }, } class TestSessionsMtls: - @pytest.mark.asyncio - async def test_configure_mtls_channel(self): - """ - Tests that the mTLS channel configures correctly when a - valid workload config is mocked. - """ - with mock.patch.dict( + + @pytest.mark.asyncio + async def test_configure_mtls_channel(self): + """Tests that the mTLS channel configures correctly when a valid workload config is mocked.""" + with ( + mock.patch.dict( os.environ, {"GOOGLE_API_USE_CLIENT_CERTIFICATE": "true"} - ), mock.patch("os.path.exists") as mock_exists, mock.patch( - "builtins.open", mock.mock_open(read_data=json.dumps(VALID_WORKLOAD_CONFIG)) - ), mock.patch( + ), + mock.patch("os.path.exists") as mock_exists, + mock.patch( + "builtins.open", + mock.mock_open(read_data=json.dumps(VALID_WORKLOAD_CONFIG)), + ), + mock.patch( "google.auth.aio.transport.mtls.get_client_cert_and_key" - ) as mock_helper, mock.patch( + ) as mock_helper, + mock.patch( "google.auth.aio.transport.mtls.make_client_cert_ssl_context" - ) as mock_make_context, mock.patch( - "aiohttp.TCPConnector" - ) as mock_connector, mock.patch( - "aiohttp.ClientSession" - ) as mock_session: - mock_session.return_value.close = mock.AsyncMock() - mock_exists.return_value = True - mock_helper.return_value = (True, b"fake_cert_data", b"fake_key_data") - - mock_context = mock.Mock(spec=ssl.SSLContext) - mock_make_context.return_value = mock_context - - mock_creds = mock.AsyncMock(spec=credentials.Credentials) - session = sessions.AsyncAuthorizedSession(mock_creds) - - await session.configure_mtls_channel() - - assert session._is_mtls is True - mock_make_context.assert_called_once_with( - b"fake_cert_data", b"fake_key_data" - ) - mock_connector.assert_called_once_with(ssl=mock_context) - mock_session.assert_called_once_with(connector=mock_connector.return_value) - await session.close() - - @pytest.mark.asyncio - async def test_configure_mtls_channel_disabled(self): - """ - Tests behavior when the config file does not exist. - """ - with mock.patch.dict( + ) as mock_make_context, + mock.patch("aiohttp.TCPConnector") as mock_connector, + mock.patch("aiohttp.ClientSession") as mock_session, + ): + mock_session.return_value.close = mock.AsyncMock() + mock_exists.return_value = True + mock_helper.return_value = (True, b"fake_cert_data", b"fake_key_data") + + mock_context = mock.Mock(spec=ssl.SSLContext) + mock_make_context.return_value = mock_context + + mock_creds = mock.AsyncMock(spec=credentials.Credentials) + session = sessions.AsyncAuthorizedSession(mock_creds) + + await session.configure_mtls_channel() + + assert session._is_mtls is True + mock_make_context.assert_called_once_with( + b"fake_cert_data", b"fake_key_data" + ) + mock_connector.assert_called_once_with(ssl=mock_context) + mock_session.assert_called_once_with( + connector=mock_connector.return_value + ) + await session.close() + + @pytest.mark.asyncio + async def test_configure_mtls_channel_disabled(self): + """Tests behavior when the config file does not exist.""" + with ( + mock.patch.dict( os.environ, {"GOOGLE_API_USE_CLIENT_CERTIFICATE": "true"} - ), mock.patch("os.path.exists") as mock_exists: - mock_exists.return_value = False - mock_creds = mock.AsyncMock(spec=credentials.Credentials) - session = sessions.AsyncAuthorizedSession(mock_creds) - await session.configure_mtls_channel() - assert session._is_mtls is False - await session.close() - - @pytest.mark.asyncio - async def test_configure_mtls_channel_invalid_format(self): - """ - Verifies that the MutualTLSChannelError is raised for bad formats. - """ - with mock.patch.dict( + ), + mock.patch("os.path.exists") as mock_exists, + ): + mock_exists.return_value = False + mock_creds = mock.AsyncMock(spec=credentials.Credentials) + session = sessions.AsyncAuthorizedSession(mock_creds) + await session.configure_mtls_channel() + assert session._is_mtls is False + await session.close() + + @pytest.mark.asyncio + async def test_configure_mtls_channel_invalid_format(self): + """Verifies that the MutualTLSChannelError is raised for bad formats.""" + with ( + mock.patch.dict( os.environ, {"GOOGLE_API_USE_CLIENT_CERTIFICATE": "true"} - ), mock.patch("os.path.exists") as mock_exists, mock.patch( + ), + mock.patch("os.path.exists") as mock_exists, + mock.patch( "builtins.open", mock.mock_open(read_data='{"invalid": "format"}') - ): - mock_exists.return_value = True - mock_creds = mock.AsyncMock(spec=credentials.Credentials) - session = sessions.AsyncAuthorizedSession(mock_creds) - - with pytest.raises(exceptions.MutualTLSChannelError): - await session.configure_mtls_channel() - await session.close() - - @pytest.mark.asyncio - async def test_configure_mtls_channel_invalud_fields(self): - """ - If cert is missing expected keys, it should fail gracefully - """ - with mock.patch.dict( + ), + ): + mock_exists.return_value = True + mock_creds = mock.AsyncMock(spec=credentials.Credentials) + session = sessions.AsyncAuthorizedSession(mock_creds) + + with pytest.raises(exceptions.MutualTLSChannelError): + await session.configure_mtls_channel() + await session.close() + + @pytest.mark.asyncio + async def test_configure_mtls_channel_invalid_fields(self): + """If cert is missing expected keys, it should fail gracefully.""" + with ( + mock.patch.dict( os.environ, {"GOOGLE_API_USE_CLIENT_CERTIFICATE": "true"} - ), mock.patch("os.path.exists") as mock_exists, mock.patch( + ), + mock.patch("os.path.exists") as mock_exists, + mock.patch( "builtins.open", mock.mock_open(read_data='{"cert_configs": {}}') - ): - mock_exists.return_value = True - mock_creds = mock.AsyncMock(spec=credentials.Credentials) - session = sessions.AsyncAuthorizedSession(mock_creds) - await session.configure_mtls_channel() - assert session._is_mtls is False - await session.close() - - @pytest.mark.asyncio - async def test_configure_mtls_channel_mock_callback(self): - """ - Tests mTLS configuration using bytes-returning callback. - """ - - def mock_callback(): - return (b"fake_cert_bytes", b"fake_key_bytes") - - with mock.patch.dict( + ), + ): + mock_exists.return_value = True + mock_creds = mock.AsyncMock(spec=credentials.Credentials) + session = sessions.AsyncAuthorizedSession(mock_creds) + await session.configure_mtls_channel() + assert session._is_mtls is False + await session.close() + + @pytest.mark.asyncio + async def test_configure_mtls_channel_mock_callback(self): + """Tests mTLS configuration using bytes-returning callback.""" + + def mock_callback(): + return (b"fake_cert_bytes", b"fake_key_bytes") + + with ( + mock.patch.dict( os.environ, {"GOOGLE_API_USE_CLIENT_CERTIFICATE": "true"} - ), mock.patch( + ), + mock.patch( "google.auth.transport.mtls.has_default_client_cert_source", return_value=True, - ), mock.patch( + ), + mock.patch( "google.auth.aio.transport.mtls.make_client_cert_ssl_context" - ) as mock_make_context, mock.patch( - "aiohttp.TCPConnector" - ) as mock_connector, mock.patch( - "aiohttp.ClientSession" - ) as mock_session: - mock_session.return_value.close = mock.AsyncMock() - mock_context = mock.Mock(spec=ssl.SSLContext) - mock_make_context.return_value = mock_context - - mock_creds = mock.AsyncMock(spec=credentials.Credentials) - session = sessions.AsyncAuthorizedSession(mock_creds) - - await session.configure_mtls_channel(client_cert_callback=mock_callback) - - assert session._is_mtls is True - mock_make_context.assert_called_once_with( - b"fake_cert_bytes", b"fake_key_bytes" - ) - mock_connector.assert_called_once_with(ssl=mock_context) - mock_session.assert_called_once_with(connector=mock_connector.return_value) - await session.close() - - @pytest.mark.asyncio - async def test_configure_mtls_channel_custom_request(self): - """Tests that if _auth_request is not an AiohttpRequest, _is_mtls is set to False - because we can't configure the custom request with mTLS. - """ - with mock.patch.dict( + ) as mock_make_context, + mock.patch("aiohttp.TCPConnector") as mock_connector, + mock.patch("aiohttp.ClientSession") as mock_session, + ): + mock_session.return_value.close = mock.AsyncMock() + mock_context = mock.Mock(spec=ssl.SSLContext) + mock_make_context.return_value = mock_context + + mock_creds = mock.AsyncMock(spec=credentials.Credentials) + session = sessions.AsyncAuthorizedSession(mock_creds) + + await session.configure_mtls_channel(client_cert_callback=mock_callback) + + assert session._is_mtls is True + mock_make_context.assert_called_once_with( + b"fake_cert_bytes", b"fake_key_bytes" + ) + mock_connector.assert_called_once_with(ssl=mock_context) + mock_session.assert_called_once_with( + connector=mock_connector.return_value + ) + await session.close() + + @pytest.mark.asyncio + async def test_configure_mtls_channel_custom_request(self): + """Tests that if _auth_request is not an AiohttpRequest, _is_mtls is set to False.""" + with ( + mock.patch.dict( os.environ, {"GOOGLE_API_USE_CLIENT_CERTIFICATE": "true"} - ), mock.patch("os.path.exists") as mock_exists, mock.patch( - "builtins.open", mock.mock_open(read_data=json.dumps(VALID_WORKLOAD_CONFIG)) - ), mock.patch( + ), + mock.patch("os.path.exists") as mock_exists, + mock.patch( + "builtins.open", + mock.mock_open(read_data=json.dumps(VALID_WORKLOAD_CONFIG)), + ), + mock.patch( "google.auth.aio.transport.mtls.get_client_cert_and_key" - ) as mock_helper, mock.patch( + ) as mock_helper, + mock.patch( "google.auth.aio.transport.mtls.make_client_cert_ssl_context" - ) as mock_make_context: - mock_exists.return_value = True - mock_helper.return_value = (True, b"fake_cert_data", b"fake_key_data") - - mock_context = mock.Mock(spec=ssl.SSLContext) - mock_make_context.return_value = mock_context - - mock_creds = mock.AsyncMock(spec=credentials.Credentials) - mock_auth_request = mock.AsyncMock(spec=transport.Request) - session = sessions.AsyncAuthorizedSession( - mock_creds, auth_request=mock_auth_request - ) - - with pytest.warns(UserWarning, match="Attempted to establish mTLS"): - await session.configure_mtls_channel() - - # If the request handler is not an AiohttpRequest, the library cannot configure - # the connection to use mTLS, so _is_mtls must be False to reflect this unconfigured state. - assert session._is_mtls is False - mock_make_context.assert_not_called() - await session.close() - - @pytest.mark.asyncio - async def test_configure_mtls_channel_exception_resets_flag(self): - """ - Tests that self._is_mtls is reset to False if an exception is raised - during configuration. - """ - with mock.patch.dict( + ) as mock_make_context, + ): + mock_exists.return_value = True + mock_helper.return_value = (True, b"fake_cert_data", b"fake_key_data") + + mock_context = mock.Mock(spec=ssl.SSLContext) + mock_make_context.return_value = mock_context + + mock_creds = mock.AsyncMock(spec=credentials.Credentials) + mock_auth_request = mock.AsyncMock(spec=transport.Request) + session = sessions.AsyncAuthorizedSession( + mock_creds, auth_request=mock_auth_request + ) + + with pytest.warns(UserWarning, match="Attempted to establish mTLS"): + await session.configure_mtls_channel() + + assert session._is_mtls is False + mock_make_context.assert_not_called() + await session.close() + + @pytest.mark.asyncio + async def test_configure_mtls_channel_exception_resets_flag(self): + """Tests that self._is_mtls is reset to False if an exception is raised during configuration.""" + with ( + mock.patch.dict( os.environ, {"GOOGLE_API_USE_CLIENT_CERTIFICATE": "true"} - ), mock.patch("os.path.exists") as mock_exists, mock.patch( - "builtins.open", mock.mock_open(read_data=json.dumps(VALID_WORKLOAD_CONFIG)) - ), mock.patch( + ), + mock.patch("os.path.exists") as mock_exists, + mock.patch( + "builtins.open", + mock.mock_open(read_data=json.dumps(VALID_WORKLOAD_CONFIG)), + ), + mock.patch( "google.auth.aio.transport.mtls.get_client_cert_and_key" - ) as mock_helper, mock.patch( + ) as mock_helper, + mock.patch( "google.auth.aio.transport.mtls.make_client_cert_ssl_context" - ) as mock_make_context: - mock_exists.return_value = True - mock_helper.return_value = (True, b"fake_cert_data", b"fake_key_data") - mock_make_context.side_effect = exceptions.ClientCertError("Mock error") - - mock_creds = mock.AsyncMock(spec=credentials.Credentials) - session = sessions.AsyncAuthorizedSession(mock_creds) - - with pytest.raises(exceptions.MutualTLSChannelError): - await session.configure_mtls_channel() - - assert session._is_mtls is False - await session.close() - - @pytest.mark.asyncio - async def test_configure_mtls_channel_transport_error_resets_flag(self): - """ - Tests that self._is_mtls is reset to False if a TransportError is raised - during configuration. - """ - with mock.patch.dict( + ) as mock_make_context, + ): + mock_exists.return_value = True + mock_helper.return_value = (True, b"fake_cert_data", b"fake_key_data") + mock_make_context.side_effect = exceptions.ClientCertError("Mock error") + + mock_creds = mock.AsyncMock(spec=credentials.Credentials) + session = sessions.AsyncAuthorizedSession(mock_creds) + + with pytest.raises(exceptions.MutualTLSChannelError): + await session.configure_mtls_channel() + + assert session._is_mtls is False + await session.close() + + @pytest.mark.asyncio + async def test_configure_mtls_channel_transport_error_resets_flag(self): + """Tests that self._is_mtls is reset to False if a TransportError is raised.""" + with ( + mock.patch.dict( os.environ, {"GOOGLE_API_USE_CLIENT_CERTIFICATE": "true"} - ), mock.patch("os.path.exists") as mock_exists, mock.patch( - "builtins.open", mock.mock_open(read_data=json.dumps(VALID_WORKLOAD_CONFIG)) - ), mock.patch( + ), + mock.patch("os.path.exists") as mock_exists, + mock.patch( + "builtins.open", + mock.mock_open(read_data=json.dumps(VALID_WORKLOAD_CONFIG)), + ), + mock.patch( "google.auth.aio.transport.mtls.get_client_cert_and_key" - ) as mock_helper, mock.patch( + ) as mock_helper, + mock.patch( "google.auth.aio.transport.mtls.make_client_cert_ssl_context" - ) as mock_make_context: - mock_exists.return_value = True - mock_helper.return_value = (True, b"fake_cert_data", b"fake_key_data") - mock_make_context.side_effect = exceptions.TransportError("Mock error") - - mock_creds = mock.AsyncMock(spec=credentials.Credentials) - session = sessions.AsyncAuthorizedSession(mock_creds) - - with pytest.raises(exceptions.MutualTLSChannelError): - await session.configure_mtls_channel() - - assert session._is_mtls is False - await session.close() - - @pytest.mark.asyncio - async def test_configure_mtls_channel_atomic_on_exception(self): - """ - Tests that if configure_mtls_channel has already successfully configured mTLS, - a subsequent attempt that raises an exception will preserve the original mTLS state. - """ - # Step 1: Successful configuration - with mock.patch.dict( + ) as mock_make_context, + ): + mock_exists.return_value = True + mock_helper.return_value = (True, b"fake_cert_data", b"fake_key_data") + mock_make_context.side_effect = exceptions.TransportError("Mock error") + + mock_creds = mock.AsyncMock(spec=credentials.Credentials) + session = sessions.AsyncAuthorizedSession(mock_creds) + + with pytest.raises(exceptions.MutualTLSChannelError): + await session.configure_mtls_channel() + + assert session._is_mtls is False + await session.close() + + @pytest.mark.asyncio + async def test_configure_mtls_channel_atomic_on_exception(self): + """Tests that if configure_mtls_channel already succeeded, a subsequent failure preserves state.""" + # Step 1: Successful configuration + with ( + mock.patch.dict( os.environ, {"GOOGLE_API_USE_CLIENT_CERTIFICATE": "true"} - ), mock.patch("os.path.exists") as mock_exists, mock.patch( - "builtins.open", mock.mock_open(read_data=json.dumps(VALID_WORKLOAD_CONFIG)) - ), mock.patch( + ), + mock.patch("os.path.exists") as mock_exists, + mock.patch( + "builtins.open", + mock.mock_open(read_data=json.dumps(VALID_WORKLOAD_CONFIG)), + ), + mock.patch( "google.auth.aio.transport.mtls.get_client_cert_and_key" - ) as mock_helper, mock.patch( + ) as mock_helper, + mock.patch( "google.auth.aio.transport.mtls.make_client_cert_ssl_context" - ) as mock_make_context, mock.patch( - "aiohttp.TCPConnector" - ), mock.patch( - "aiohttp.ClientSession" - ) as mock_session: - mock_session.return_value.close = mock.AsyncMock() - mock_exists.return_value = True - mock_helper.return_value = (True, b"fake_cert_data_1", b"fake_key_data_1") - - mock_context = mock.Mock(spec=ssl.SSLContext) - mock_make_context.return_value = mock_context - - mock_creds = mock.AsyncMock(spec=credentials.Credentials) - session = sessions.AsyncAuthorizedSession(mock_creds) - - await session.configure_mtls_channel() - assert session._is_mtls is True - assert session._cached_cert == b"fake_cert_data_1" - first_auth_request = session._auth_request - - # Step 2: Failed subsequent configuration attempt - # Reset task so we trigger a new configuration run - session._mtls_init_task = None - - # Patch context generator to fail this time - mock_make_context.side_effect = exceptions.ClientCertError("Mock error") - - with pytest.raises(exceptions.MutualTLSChannelError): - await session.configure_mtls_channel() - - # Verify that the state remains unchanged from the first successful configuration - assert session._is_mtls is True - assert session._cached_cert == b"fake_cert_data_1" - assert session._auth_request is first_auth_request - await session.close() - - @pytest.mark.asyncio - async def test_configure_mtls_channel_close_exception_does_not_abort(self): - """ - Tests that if old_auth_request.close() raises an exception, the mTLS - configuration is still considered successful, and is_mtls remains True - without raising MutualTLSChannelError. - """ - with mock.patch.dict( + ) as mock_make_context, + mock.patch("aiohttp.TCPConnector"), + mock.patch("aiohttp.ClientSession") as mock_session, + ): + mock_session.return_value.close = mock.AsyncMock() + mock_exists.return_value = True + mock_helper.return_value = ( + True, + b"fake_cert_data_1", + b"fake_key_data_1", + ) + + mock_context = mock.Mock(spec=ssl.SSLContext) + mock_make_context.return_value = mock_context + + mock_creds = mock.AsyncMock(spec=credentials.Credentials) + session = sessions.AsyncAuthorizedSession(mock_creds) + + await session.configure_mtls_channel() + assert session._is_mtls is True + assert session._cached_cert == b"fake_cert_data_1" + first_auth_request = session._auth_request + + # Step 2: Failed subsequent configuration attempt + session._mtls_init_task = None + mock_make_context.side_effect = exceptions.ClientCertError("Mock error") + + with pytest.raises(exceptions.MutualTLSChannelError): + await session.configure_mtls_channel() + + assert session._is_mtls is True + assert session._cached_cert == b"fake_cert_data_1" + assert session._auth_request is first_auth_request + await session.close() + + @pytest.mark.asyncio + async def test_configure_mtls_channel_close_exception_does_not_abort(self): + """Tests that an exception in old_auth_request.close() does not abort configuration.""" + with ( + mock.patch.dict( os.environ, {"GOOGLE_API_USE_CLIENT_CERTIFICATE": "true"} - ), mock.patch("os.path.exists") as mock_exists, mock.patch( - "builtins.open", mock.mock_open(read_data=json.dumps(VALID_WORKLOAD_CONFIG)) - ), mock.patch( + ), + mock.patch("os.path.exists") as mock_exists, + mock.patch( + "builtins.open", + mock.mock_open(read_data=json.dumps(VALID_WORKLOAD_CONFIG)), + ), + mock.patch( "google.auth.aio.transport.mtls.get_client_cert_and_key" - ) as mock_helper, mock.patch( + ) as mock_helper, + mock.patch( "google.auth.aio.transport.mtls.make_client_cert_ssl_context" - ) as mock_make_context, mock.patch( - "aiohttp.TCPConnector" - ), mock.patch( - "aiohttp.ClientSession" - ) as mock_session: - mock_session.return_value.close = mock.AsyncMock() - mock_exists.return_value = True - mock_helper.return_value = (True, b"fake_cert_data", b"fake_key_data") - - mock_context = mock.Mock(spec=ssl.SSLContext) - mock_make_context.return_value = mock_context - - mock_creds = mock.AsyncMock(spec=credentials.Credentials) - session = sessions.AsyncAuthorizedSession(mock_creds) - - # Mock close() of the initial self._auth_request to raise an exception - session._auth_request.close = mock.AsyncMock( - side_effect=Exception("Mock close error") - ) - - # Should complete successfully without raising MutualTLSChannelError - await session.configure_mtls_channel() - - assert session._is_mtls is True - assert session._cached_cert == b"fake_cert_data" - await session.close() - - @pytest.mark.asyncio - async def test_cert_rotation_failure_raises_error(self): - mock_creds = mock.AsyncMock(spec=credentials.Credentials) - mock_creds.before_request = mock.AsyncMock(return_value=None) - - mock_resp = mock.Mock() - mock_resp.status_code = http_client.UNAUTHORIZED - mock_auth_req = mock.AsyncMock(return_value=mock_resp) - - session = sessions.AsyncAuthorizedSession( - mock_creds, auth_request=mock_auth_req - ) - session._is_mtls = True - session._cached_cert = b"old_cert" - - new_cert = b"new_cert" - new_key = b"new_key" - - with mock.patch( + ) as mock_make_context, + mock.patch("aiohttp.TCPConnector"), + mock.patch("aiohttp.ClientSession") as mock_session, + ): + mock_session.return_value.close = mock.AsyncMock() + mock_exists.return_value = True + mock_helper.return_value = (True, b"fake_cert_data", b"fake_key_data") + + mock_context = mock.Mock(spec=ssl.SSLContext) + mock_make_context.return_value = mock_context + + mock_creds = mock.AsyncMock(spec=credentials.Credentials) + session = sessions.AsyncAuthorizedSession(mock_creds) + + session._auth_request.close = mock.AsyncMock( + side_effect=Exception("Mock close error") + ) + + await session.configure_mtls_channel() + + assert session._is_mtls is True + assert session._cached_cert == b"fake_cert_data" + await session.close() + + @pytest.mark.asyncio + async def test_cert_rotation_failure_raises_error(self): + mock_creds = mock.AsyncMock(spec=credentials.Credentials) + mock_creds.before_request = mock.AsyncMock(return_value=None) + + mock_resp = mock.Mock() + mock_resp.status_code = http_client.UNAUTHORIZED + mock_auth_req = mock.AsyncMock(return_value=mock_resp) + + session = sessions.AsyncAuthorizedSession( + mock_creds, auth_request=mock_auth_req + ) + session._is_mtls = True + session._cached_cert = b"old_cert" + + new_cert = b"new_cert" + new_key = b"new_key" + + with ( + mock.patch( "google.auth.aio.transport.mtls.check_parameters_for_unauthorized_response", new_callable=mock.AsyncMock, - ) as mock_check, mock.patch.object( + ) as mock_check, + mock.patch.object( session, "configure_mtls_channel", new_callable=mock.AsyncMock - ) as mock_conf: - mock_check.return_value = (new_cert, new_key, b"old_fp", b"new_fp") - mock_conf.side_effect = Exception("Failed to reconfigure") + ) as mock_conf, + ): + mock_check.return_value = (new_cert, new_key, b"old_fp", b"new_fp") + mock_conf.side_effect = Exception("Failed to reconfigure") - with pytest.raises(exceptions.MutualTLSChannelError): - await session.request("GET", "https://pubsub.mtls.googleapis.com/test") + with pytest.raises(exceptions.MutualTLSChannelError): + await session.request("GET", "https://pubsub.mtls.googleapis.com/test") - mock_check.assert_called_once() - mock_conf.assert_called_once() + mock_check.assert_called_once() + mock_conf.assert_called_once() - await session.close() + await session.close() - @pytest.mark.asyncio - async def test_cert_rotation_check_params_fails(self): - mock_creds = mock.AsyncMock(spec=credentials.Credentials) - mock_creds.before_request = mock.AsyncMock(return_value=None) - mock_creds.refresh = mock.AsyncMock() + @pytest.mark.asyncio + async def test_cert_rotation_check_params_fails(self): + mock_creds = mock.AsyncMock(spec=credentials.Credentials) + mock_creds.before_request = mock.AsyncMock(return_value=None) + mock_creds.refresh = mock.AsyncMock() - mock_resp = mock.Mock() - mock_resp.status_code = http_client.UNAUTHORIZED - mock_auth_req = mock.AsyncMock(return_value=mock_resp) + mock_resp = mock.Mock() + mock_resp.status_code = http_client.UNAUTHORIZED + mock_auth_req = mock.AsyncMock(return_value=mock_resp) - session = sessions.AsyncAuthorizedSession( - mock_creds, auth_request=mock_auth_req - ) - session._is_mtls = True - session._cached_cert = b"old_cert" + session = sessions.AsyncAuthorizedSession( + mock_creds, auth_request=mock_auth_req + ) + session._is_mtls = True + session._cached_cert = b"old_cert" - with mock.patch( + with ( + mock.patch( "google.auth.aio.transport.mtls.check_parameters_for_unauthorized_response", new_callable=mock.AsyncMock, - ) as mock_check, mock.patch.object( + ) as mock_check, + mock.patch.object( session, "configure_mtls_channel", new_callable=mock.AsyncMock - ) as mock_conf: - mock_check.side_effect = exceptions.MutualTLSChannelError( - "Failed to check params" - ) + ) as mock_conf, + ): + mock_check.side_effect = exceptions.MutualTLSChannelError( + "Failed to check params" + ) - resp = await session.request( - "GET", "https://pubsub.mtls.googleapis.com/test" - ) + resp = await session.request( + "GET", "https://pubsub.mtls.googleapis.com/test" + ) - assert resp == mock_resp - mock_check.assert_called_once() - mock_creds.refresh.assert_not_called() - mock_conf.assert_not_called() + assert resp == mock_resp + mock_check.assert_called_once() + mock_creds.refresh.assert_not_called() + mock_conf.assert_not_called() - await session.close() + await session.close() - @pytest.mark.asyncio - async def test_no_cert_rotation_when_cert_matches_and_mtls_enabled(self): - mock_creds = mock.AsyncMock(spec=credentials.Credentials) - mock_creds.before_request = mock.AsyncMock(return_value=None) - mock_creds.refresh = mock.AsyncMock(return_value=None) + @pytest.mark.asyncio + async def test_no_cert_rotation_when_cert_matches_and_mtls_enabled(self): + mock_creds = mock.AsyncMock(spec=credentials.Credentials) + mock_creds.before_request = mock.AsyncMock(return_value=None) + mock_creds.refresh = mock.AsyncMock(return_value=None) - mock_resp_401 = mock.Mock() - mock_resp_401.status_code = http_client.UNAUTHORIZED - mock_resp_401.close = mock.AsyncMock() + mock_resp_401 = mock.Mock() + mock_resp_401.status_code = http_client.UNAUTHORIZED + mock_resp_401.close = mock.AsyncMock() - mock_resp_200 = mock.Mock() - mock_resp_200.status_code = http_client.OK + mock_resp_200 = mock.Mock() + mock_resp_200.status_code = http_client.OK - # 401 on initial request, 200 on retry after refresh - mock_auth_req = mock.AsyncMock(side_effect=[mock_resp_401, mock_resp_200]) + # 401 on initial request, 200 on retry after refresh + mock_auth_req = mock.AsyncMock(side_effect=[mock_resp_401, mock_resp_200]) - session = sessions.AsyncAuthorizedSession( - mock_creds, auth_request=mock_auth_req - ) - session._is_mtls = True - session._cached_cert = b"old_cert" + session = sessions.AsyncAuthorizedSession( + mock_creds, auth_request=mock_auth_req + ) + session._is_mtls = True + session._cached_cert = b"old_cert" - new_cert = b"new_cert" - new_key = b"new_key" + new_cert = b"new_cert" + new_key = b"new_key" - with mock.patch( + with ( + mock.patch( "google.auth.aio.transport.mtls.check_parameters_for_unauthorized_response", new_callable=mock.AsyncMock, - ) as mock_check, mock.patch.object( + ) as mock_check, + mock.patch.object( session, "configure_mtls_channel", new_callable=mock.AsyncMock - ) as mock_conf: - # Matching fingerprints mean no mTLS rotation is needed - mock_check.return_value = (new_cert, new_key, b"old_fp", b"old_fp") - - resp = await session.request( - "GET", "https://pubsub.mtls.googleapis.com/test" - ) - - assert resp == mock_resp_200 - mock_check.assert_called_once() - mock_conf.assert_not_called() - mock_creds.refresh.assert_called_once() - assert mock_auth_req.call_count == 2 - mock_resp_401.close.assert_awaited_once() - - await session.close() - - @pytest.mark.asyncio - async def test_cert_rotation_success_and_retry(self): - mock_creds = mock.AsyncMock(spec=credentials.Credentials) - mock_creds.before_request = mock.AsyncMock(return_value=None) - mock_creds.refresh = mock.AsyncMock(return_value=None) - - # Initial request fails natively with 401. Retry succeeds with 200. - mock_resp_401 = mock.Mock() - mock_resp_401.status_code = http_client.UNAUTHORIZED - mock_resp_200 = mock.Mock() - mock_resp_200.status_code = http_client.OK - - # Use side_effect to dynamically yield responses - mock_auth_req = mock.AsyncMock(side_effect=[mock_resp_401, mock_resp_200]) - - session = sessions.AsyncAuthorizedSession( - mock_creds, auth_request=mock_auth_req - ) - session._is_mtls = True - session._cached_cert = b"old_cert" - - new_cert = b"new_cert" - new_key = b"new_key" - - with mock.patch( + ) as mock_conf, + ): + # Matching fingerprints mean no mTLS rotation is needed + mock_check.return_value = (new_cert, new_key, b"old_fp", b"old_fp") + + resp = await session.request( + "GET", "https://pubsub.mtls.googleapis.com/test" + ) + + assert resp == mock_resp_200 + mock_check.assert_called_once() + mock_conf.assert_not_called() + mock_creds.refresh.assert_called_once() + assert mock_auth_req.call_count == 2 + mock_resp_401.close.assert_awaited_once() + + await session.close() + + @pytest.mark.asyncio + async def test_cert_rotation_success_and_retry(self): + mock_creds = mock.AsyncMock(spec=credentials.Credentials) + mock_creds.before_request = mock.AsyncMock(return_value=None) + mock_creds.refresh = mock.AsyncMock(return_value=None) + + mock_resp_401 = mock.Mock() + mock_resp_401.status_code = http_client.UNAUTHORIZED + mock_resp_200 = mock.Mock() + mock_resp_200.status_code = http_client.OK + + mock_auth_req = mock.AsyncMock(side_effect=[mock_resp_401, mock_resp_200]) + + session = sessions.AsyncAuthorizedSession( + mock_creds, auth_request=mock_auth_req + ) + session._is_mtls = True + session._cached_cert = b"old_cert" + + new_cert = b"new_cert" + new_key = b"new_key" + + with ( + mock.patch( "google.auth.aio.transport.mtls.check_parameters_for_unauthorized_response", new_callable=mock.AsyncMock, - ) as mock_check, mock.patch.object( + ) as mock_check, + mock.patch.object( session, "configure_mtls_channel", new_callable=mock.AsyncMock - ) as mock_conf: - mock_check.return_value = (new_cert, new_key, b"old_fp", b"new_fp") - - resp = await session.request( - "GET", "https://pubsub.mtls.googleapis.com/test" - ) - - # 1. Assert the retried 200 response is successfully returned to the user - assert resp == mock_resp_200 - - # 2. Assert rotation logic correctly executed - mock_check.assert_called_once() - mock_conf.assert_called_once_with(mock.ANY) + ) as mock_conf, + ): + mock_check.return_value = (new_cert, new_key, b"old_fp", b"new_fp") + + resp = await session.request( + "GET", "https://pubsub.mtls.googleapis.com/test" + ) + + assert resp == mock_resp_200 + mock_check.assert_called_once() + mock_conf.assert_called_once_with(mock.ANY) + mock_creds.refresh.assert_called_once() + assert mock_creds.before_request.call_count == 2 + + await session.close() + + @pytest.mark.asyncio + async def test_cert_rotation_lock_contention(self): + mock_creds = mock.AsyncMock(spec=credentials.Credentials) + mock_creds.before_request = mock.AsyncMock(return_value=None) + mock_creds.refresh = mock.AsyncMock(return_value=None) + + mock_resp_401 = mock.Mock() + mock_resp_401.status_code = http_client.UNAUTHORIZED + mock_auth_req = mock.AsyncMock(return_value=mock_resp_401) + + session = sessions.AsyncAuthorizedSession( + mock_creds, auth_request=mock_auth_req + ) + session._is_mtls = True + session._cached_cert = b"old_cert" + + new_cert = b"new_cert" + new_key = b"new_key" + + async def mock_configure_mtls_channel(*args, **kwargs): + await asyncio.sleep(0.01) + session._cached_cert = new_cert + + async def mock_check_side_effect(cached_cert, callback=None): + if cached_cert == b"old_cert": + return (new_cert, new_key, b"old_fp", b"new_fp") + return (new_cert, new_key, b"new_fp", b"new_fp") + + with ( + mock.patch( + "google.auth.aio.transport.mtls.check_parameters_for_unauthorized_response", + new_callable=mock.AsyncMock, + ) as mock_check, + mock.patch.object( + session, "configure_mtls_channel", new_callable=mock.AsyncMock + ) as mock_conf, + ): + mock_check.side_effect = mock_check_side_effect + mock_conf.side_effect = mock_configure_mtls_channel - # 3. Assert credentials were explicitly refreshed - mock_creds.refresh.assert_called_once() + tasks = [ + session.request("GET", "https://pubsub.mtls.googleapis.com/test") + for _ in range(3) + ] + responses = await asyncio.gather(*tasks) - # 4. Assert headers were explicitly rebound on the recursive retry (2 invocations) - assert mock_creds.before_request.call_count == 2 + for resp in responses: + assert resp == mock_resp_401 - await session.close() + mock_conf.assert_called_once() - @pytest.mark.asyncio - async def test_cert_rotation_lock_contention(self): - mock_creds = mock.AsyncMock(spec=credentials.Credentials) - mock_creds.before_request = mock.AsyncMock(return_value=None) - mock_creds.refresh = mock.AsyncMock(return_value=None) + await session.close() - mock_resp_401 = mock.Mock() - mock_resp_401.status_code = http_client.UNAUTHORIZED - mock_auth_req = mock.AsyncMock(return_value=mock_resp_401) + @pytest.mark.asyncio + async def test_cert_rotation_lock_contention_no_cert_change(self): + mock_creds = mock.AsyncMock(spec=credentials.Credentials) + mock_creds.before_request = mock.AsyncMock(return_value=None) + mock_creds.refresh = mock.AsyncMock(return_value=None) - session = sessions.AsyncAuthorizedSession( - mock_creds, auth_request=mock_auth_req - ) - session._is_mtls = True - session._cached_cert = b"old_cert" + mock_resp_401 = mock.Mock() + mock_resp_401.status_code = http_client.UNAUTHORIZED + mock_resp_200 = mock.Mock() + mock_resp_200.status_code = http_client.OK - new_cert = b"new_cert" - new_key = b"new_key" + mock_auth_req = mock.AsyncMock( + side_effect=[mock_resp_401] * 3 + [mock_resp_200] * 3 + ) - async def mock_configure_mtls_channel(*args, **kwargs): - await asyncio.sleep(0.01) - session._cached_cert = new_cert + session = sessions.AsyncAuthorizedSession( + mock_creds, auth_request=mock_auth_req + ) + session._is_mtls = True + session._cached_cert = b"old_cert" - async def mock_check_side_effect(callback, cached_cert): - if cached_cert == b"old_cert": - return (new_cert, new_key, b"old_fp", b"new_fp") - return (new_cert, new_key, b"new_fp", b"new_fp") + async def mock_check_side_effect(cached_cert, callback=None): + await asyncio.sleep(0.01) + return (b"old_cert", b"old_key", b"old_fp", b"old_fp") - with mock.patch( + with ( + mock.patch( "google.auth.aio.transport.mtls.check_parameters_for_unauthorized_response", new_callable=mock.AsyncMock, - ) as mock_check, mock.patch.object( + ) as mock_check, + mock.patch.object( session, "configure_mtls_channel", new_callable=mock.AsyncMock - ) as mock_conf: - mock_check.side_effect = mock_check_side_effect - mock_conf.side_effect = mock_configure_mtls_channel - - tasks = [ - session.request("GET", "https://pubsub.mtls.googleapis.com/test") - for _ in range(3) - ] - responses = await asyncio.gather(*tasks) - - for resp in responses: - assert resp == mock_resp_401 - - mock_conf.assert_called_once() - - await session.close() - - @pytest.mark.asyncio - async def test_cert_rotation_lock_contention_no_cert_change(self): - mock_creds = mock.AsyncMock(spec=credentials.Credentials) - mock_creds.before_request = mock.AsyncMock(return_value=None) - mock_creds.refresh = mock.AsyncMock(return_value=None) - - mock_resp_401 = mock.Mock() - mock_resp_401.status_code = http_client.UNAUTHORIZED - mock_resp_200 = mock.Mock() - mock_resp_200.status_code = http_client.OK - - mock_auth_req = mock.AsyncMock( - side_effect=[mock_resp_401] * 3 + [mock_resp_200] * 3 - ) - - session = sessions.AsyncAuthorizedSession( - mock_creds, auth_request=mock_auth_req - ) - session._is_mtls = True - session._cached_cert = b"old_cert" - - # FIX: async def + await asyncio.sleep to allow concurrent lock contention - async def mock_check_side_effect(callback, cached_cert): - await asyncio.sleep(0.01) - return (b"old_cert", b"old_key", b"old_fp", b"old_fp") - - with mock.patch( + ) as mock_conf, + ): + mock_check.side_effect = mock_check_side_effect + + tasks = [ + session.request("GET", "https://pubsub.mtls.googleapis.com/test") + for _ in range(3) + ] + responses = await asyncio.gather(*tasks) + + for resp in responses: + assert resp == mock_resp_200 + + mock_check.assert_called_once() + mock_conf.assert_not_called() + # Concurrent 401s properly deduplicate to 1 refresh + assert mock_creds.refresh.call_count == 1 + + await session.close() + + @pytest.mark.asyncio + async def test_psc_endpoint_triggers_cert_rotation(self): + """Verifies that PSC endpoints (*.p.googleapis.com) are recognized as mTLS endpoints.""" + mock_creds = mock.AsyncMock(spec=credentials.Credentials) + mock_creds.before_request = mock.AsyncMock(return_value=None) + mock_creds.refresh = mock.AsyncMock(return_value=None) + + mock_resp_401 = mock.Mock() + mock_resp_401.status_code = http_client.UNAUTHORIZED + mock_resp_200 = mock.Mock() + mock_resp_200.status_code = http_client.OK + + mock_auth_req = mock.AsyncMock(side_effect=[mock_resp_401, mock_resp_200]) + + session = sessions.AsyncAuthorizedSession( + mock_creds, auth_request=mock_auth_req + ) + session._is_mtls = True + session._cached_cert = b"old_cert" + + new_cert = b"new_cert" + new_key = b"new_key" + + with ( + mock.patch( "google.auth.aio.transport.mtls.check_parameters_for_unauthorized_response", new_callable=mock.AsyncMock, - ) as mock_check, mock.patch.object( + ) as mock_check, + mock.patch.object( session, "configure_mtls_channel", new_callable=mock.AsyncMock - ) as mock_conf: - mock_check.side_effect = mock_check_side_effect - - tasks = [ - session.request("GET", "https://pubsub.mtls.googleapis.com/test") - for _ in range(3) - ] - responses = await asyncio.gather(*tasks) + ) as mock_conf, + ): + mock_check.return_value = (new_cert, new_key, b"old_fp", b"new_fp") - for resp in responses: - assert resp == mock_resp_200 + resp = await session.request( + "GET", "https://pubsub.p.googleapis.com/test" + ) - mock_check.assert_called_once() - mock_conf.assert_not_called() - assert mock_creds.refresh.call_count == 3 + assert resp == mock_resp_200 + mock_check.assert_called_once() + mock_conf.assert_called_once_with(mock.ANY) - await session.close() + await session.close() - @pytest.mark.asyncio - async def test_non_mtls_url_bypasses_rotation(self): - mock_creds = mock.AsyncMock(spec=credentials.Credentials) - mock_creds.before_request = mock.AsyncMock(return_value=None) - mock_creds.refresh = mock.AsyncMock(return_value=None) + @pytest.mark.asyncio + async def test_non_mtls_url_bypasses_rotation(self): + mock_creds = mock.AsyncMock(spec=credentials.Credentials) + mock_creds.before_request = mock.AsyncMock(return_value=None) + mock_creds.refresh = mock.AsyncMock(return_value=None) - mock_resp_401 = mock.Mock() - mock_resp_401.status_code = http_client.UNAUTHORIZED - mock_auth_req = mock.AsyncMock(return_value=mock_resp_401) + mock_resp_401 = mock.Mock() + mock_resp_401.status_code = http_client.UNAUTHORIZED + mock_auth_req = mock.AsyncMock(return_value=mock_resp_401) - session = sessions.AsyncAuthorizedSession( - mock_creds, auth_request=mock_auth_req - ) + session = sessions.AsyncAuthorizedSession( + mock_creds, auth_request=mock_auth_req + ) - # Even if mTLS is enabled globally... - session._is_mtls = True - session._cached_cert = b"old_cert" + session._is_mtls = True + session._cached_cert = b"old_cert" - with mock.patch( + with ( + mock.patch( "google.auth.aio.transport.mtls.check_parameters_for_unauthorized_response", new_callable=mock.AsyncMock, - ) as mock_check, mock.patch.object( + ) as mock_check, + mock.patch.object( session, "configure_mtls_channel", new_callable=mock.AsyncMock - ) as mock_conf: - # Will attempt to refresh credentials and retry the request up to 2 times (3 requests total). - resp = await session.request("GET", "https://pubsub.googleapis.com/test") - - assert resp == mock_resp_401 - mock_check.assert_not_called() - mock_conf.assert_not_called() - - # Verify the standard retry behavior executed - assert mock_creds.refresh.call_count == 2 - assert mock_auth_req.call_count == 3 - - await session.close() - - @pytest.mark.asyncio - async def test_cert_rotation_skips_retry_for_streaming(self): - mock_creds = mock.AsyncMock(spec=credentials.Credentials) - mock_creds.before_request = mock.AsyncMock(return_value=None) - mock_creds.refresh = mock.AsyncMock() - - mock_resp = mock.Mock() - mock_resp.status_code = http_client.UNAUTHORIZED - mock_auth_req = mock.AsyncMock(return_value=mock_resp) - - session = sessions.AsyncAuthorizedSession( - mock_creds, auth_request=mock_auth_req - ) - session._is_mtls = True - session._cached_cert = b"old_cert" - - # Simulate a streaming object (hasattr(data, 'read')) - class MockStream: - def read(self): - pass - - with mock.patch( + ) as mock_conf, + ): + resp = await session.request("GET", "https://example.com/test") + + assert resp == mock_resp_401 + mock_check.assert_not_called() + mock_conf.assert_not_called() + assert mock_creds.refresh.call_count == 2 + assert mock_auth_req.call_count == 3 + + await session.close() + + @pytest.mark.asyncio + async def test_cert_rotation_skips_retry_for_streaming(self): + mock_creds = mock.AsyncMock(spec=credentials.Credentials) + mock_creds.before_request = mock.AsyncMock(return_value=None) + mock_creds.refresh = mock.AsyncMock() + + mock_resp = mock.Mock() + mock_resp.status_code = http_client.UNAUTHORIZED + mock_auth_req = mock.AsyncMock(return_value=mock_resp) + + session = sessions.AsyncAuthorizedSession( + mock_creds, auth_request=mock_auth_req + ) + session._is_mtls = True + session._cached_cert = b"old_cert" + + class MockStream: + + def read(self): + pass + + with ( + mock.patch( "google.auth.aio.transport.mtls.check_parameters_for_unauthorized_response", new_callable=mock.AsyncMock, - ) as mock_check, mock.patch.object( + ) as mock_check, + mock.patch.object( session, "configure_mtls_channel", new_callable=mock.AsyncMock - ) as mock_conf: - mock_check.return_value = (b"new", b"new", b"old_fp", b"new_fp") - - resp = await session.request( - "GET", "https://pubsub.mtls.googleapis.com/test", data=MockStream() - ) - - assert resp == mock_resp - mock_conf.assert_called_once() - # Because it is streaming, it skips credentials refresh and retry - mock_creds.refresh.assert_called_once() - - await session.close() - - @pytest.mark.asyncio - async def test_cert_rotation_credential_refresh_fails(self): - """Covers the except block for RefreshError when credentials fail to refresh.""" - mock_creds = mock.AsyncMock(spec=credentials.Credentials) - mock_creds.before_request = mock.AsyncMock(return_value=None) - # Simulate a credentials refresh failure - mock_creds.refresh = mock.AsyncMock( - side_effect=exceptions.RefreshError("Refresh failed") - ) - - mock_resp = mock.Mock() - mock_resp.status_code = http_client.UNAUTHORIZED - mock_auth_req = mock.AsyncMock(return_value=mock_resp) - - session = sessions.AsyncAuthorizedSession( - mock_creds, auth_request=mock_auth_req - ) - session._is_mtls = True - session._cached_cert = b"old_cert" - - with mock.patch( + ) as mock_conf, + ): + mock_check.return_value = (b"new", b"new", b"old_fp", b"new_fp") + + resp = await session.request( + "GET", "https://pubsub.mtls.googleapis.com/test", data=MockStream() + ) + + assert resp == mock_resp + mock_conf.assert_called_once() + + mock_creds.refresh.assert_called_once() + await session.close() + + @pytest.mark.asyncio + async def test_cert_rotation_credential_refresh_fails(self): + """Covers the except block for RefreshError when credentials fail to refresh.""" + mock_creds = mock.AsyncMock(spec=credentials.Credentials) + mock_creds.before_request = mock.AsyncMock(return_value=None) + mock_creds.refresh = mock.AsyncMock( + side_effect=exceptions.RefreshError("Refresh failed") + ) + + mock_resp = mock.Mock() + mock_resp.status_code = http_client.UNAUTHORIZED + mock_auth_req = mock.AsyncMock(return_value=mock_resp) + + session = sessions.AsyncAuthorizedSession( + mock_creds, auth_request=mock_auth_req + ) + session._is_mtls = True + session._cached_cert = b"old_cert" + + with ( + mock.patch( "google.auth.aio.transport.mtls.check_parameters_for_unauthorized_response", new_callable=mock.AsyncMock, - ) as mock_check, mock.patch.object( + ) as mock_check, + mock.patch.object( session, "configure_mtls_channel", new_callable=mock.AsyncMock - ): - mock_check.return_value = (b"new", b"new", b"old_fp", b"new_fp") - - resp = await session.request( - "GET", "https://pubsub.mtls.googleapis.com/test" - ) - - # It should catch the RefreshError and return the 401 response directly - assert resp == mock_resp - mock_creds.refresh.assert_called_once() - # Ensure it didn't retry the request by checking auth_request was only called once - mock_auth_req.assert_called_once() - - await session.close() - - @pytest.mark.asyncio - async def test_cert_rotation_max_retries_exceeded(self): - """Covers the `if _auth_retry_count < 2:` max retry limit.""" - mock_creds = mock.AsyncMock(spec=credentials.Credentials) - mock_creds.before_request = mock.AsyncMock(return_value=None) - mock_creds.refresh = mock.AsyncMock(return_value=None) - - mock_resp = mock.Mock() - mock_resp.status_code = http_client.UNAUTHORIZED - mock_resp.close = mock.AsyncMock() - # Always return 401 - mock_auth_req = mock.AsyncMock(return_value=mock_resp) - - session = sessions.AsyncAuthorizedSession( - mock_creds, auth_request=mock_auth_req - ) - session._is_mtls = True - session._cached_cert = b"old_cert" - - with mock.patch( + ), + ): + mock_check.return_value = (b"new", b"new", b"old_fp", b"new_fp") + + resp = await session.request( + "GET", "https://pubsub.mtls.googleapis.com/test" + ) + + assert resp == mock_resp + mock_creds.refresh.assert_called_once() + mock_auth_req.assert_called_once() + + await session.close() + + @pytest.mark.asyncio + async def test_cert_rotation_max_retries_exceeded(self): + """Covers the `if _auth_retry_count < 2:` max retry limit.""" + mock_creds = mock.AsyncMock(spec=credentials.Credentials) + mock_creds.before_request = mock.AsyncMock(return_value=None) + mock_creds.refresh = mock.AsyncMock(return_value=None) + + mock_resp = mock.Mock() + mock_resp.status_code = http_client.UNAUTHORIZED + mock_resp.close = mock.AsyncMock() + mock_auth_req = mock.AsyncMock(return_value=mock_resp) + + session = sessions.AsyncAuthorizedSession( + mock_creds, auth_request=mock_auth_req + ) + session._is_mtls = True + session._cached_cert = b"old_cert" + + with ( + mock.patch( "google.auth.aio.transport.mtls.check_parameters_for_unauthorized_response", new_callable=mock.AsyncMock, - ) as mock_check, mock.patch.object( + ) as mock_check, + mock.patch.object( session, "configure_mtls_channel", new_callable=mock.AsyncMock - ): - # Yield new fingerprints to trigger reconfiguration branches - mock_check.return_value = (b"new", b"new", b"old_fp", b"new_fp") - - resp = await session.request( - "GET", "https://pubsub.mtls.googleapis.com/test" - ) - - assert resp == mock_resp - # Expecting exactly 3 requests (Initial try, Retry 1, Retry 2). - # When _auth_retry_count reaches 2, it drops out of the retry condition. - assert mock_auth_req.call_count == 3 - # It should have checked the cert twice before hitting the retry cap - assert mock_check.call_count == 2 - # Verify the stale response is closed before the retry - assert mock_resp.close.call_count == 2 - - await session.close() - - @pytest.mark.asyncio - async def test_session_close_cleans_old_auth_requests(self): - """Covers the loop in the `close()` method that drains `_old_auth_requests`.""" - mock_creds = mock.AsyncMock(spec=credentials.Credentials) - session = sessions.AsyncAuthorizedSession( - mock_creds, auth_request=mock.AsyncMock() - ) - - # Manually inject mocked old requests simulating past channel reconfigurations - mock_old_req_1 = mock.AsyncMock() - mock_old_req_2 = mock.AsyncMock() - mock_old_req_3_fails = mock.AsyncMock() - mock_old_req_3_fails.close.side_effect = Exception("Close error") - - session._old_auth_requests.extend( - [mock_old_req_1, mock_old_req_2, mock_old_req_3_fails] - ) - - # Triggers `await old_request.close()` for each item - await session.close() - - # Ensure all were called - mock_old_req_1.close.assert_called_once() - mock_old_req_2.close.assert_called_once() - mock_old_req_3_fails.close.assert_called_once() - # Ensure the list was cleared - assert len(session._old_auth_requests) == 0 - - @pytest.mark.asyncio - async def test_request_401_streaming_refreshes_creds_and_returns_open_response( - self, + ), ): - """Verifies that streaming requests refresh credentials but return the unclosed response.""" - mock_creds = mock.AsyncMock(spec=credentials.Credentials) - mock_creds.before_request = mock.AsyncMock(return_value=None) - mock_creds.refresh = mock.AsyncMock(return_value=None) - - mock_resp_401 = mock.Mock() - mock_resp_401.status_code = http_client.UNAUTHORIZED - mock_resp_401.close = mock.AsyncMock() - - mock_auth_req = mock.AsyncMock(return_value=mock_resp_401) - session = sessions.AsyncAuthorizedSession( - mock_creds, auth_request=mock_auth_req - ) - - streaming_data = (chunk for chunk in [b"chunk1", b"chunk2"]) - response = await session.request( - "POST", "https://example.com", data=streaming_data - ) - - assert response == mock_resp_401 - mock_creds.refresh.assert_awaited_once() - mock_resp_401.close.assert_not_called() - await session.close() - - @pytest.mark.asyncio - async def test_request_401_closes_response_on_timeout_during_recovery(self): - """Verifies that response is closed when auth_with_timeout times out during _recover_auth_state.""" - mock_creds = mock.AsyncMock(spec=credentials.Credentials) - mock_creds.before_request = mock.AsyncMock(return_value=None) - - async def slow_refresh(*args, **kwargs): - await asyncio.sleep(10) - - mock_creds.refresh = mock.AsyncMock(side_effect=slow_refresh) - - mock_resp_401 = mock.Mock() - mock_resp_401.status_code = http_client.UNAUTHORIZED - mock_resp_401.close = mock.AsyncMock() - - mock_auth_req = mock.AsyncMock(return_value=mock_resp_401) - session = sessions.AsyncAuthorizedSession( - mock_creds, auth_request=mock_auth_req - ) - - with pytest.raises(TimeoutError): - await session.request("GET", "https://example.com", max_allowed_time=0.01) - - mock_resp_401.close.assert_awaited_once() - await session.close() - - @pytest.mark.asyncio - async def test_request_401_closes_response_on_cancellation(self): - """Verifies that response is closed and CancelledError propagated if task is cancelled during recovery.""" - mock_creds = mock.AsyncMock(spec=credentials.Credentials) - mock_creds.before_request = mock.AsyncMock(return_value=None) - - refresh_started = asyncio.Event() - - async def cancel_on_refresh(*args, **kwargs): - refresh_started.set() - await asyncio.sleep(10) - - mock_creds.refresh = mock.AsyncMock(side_effect=cancel_on_refresh) - - mock_resp_401 = mock.Mock() - mock_resp_401.status_code = http_client.UNAUTHORIZED - mock_resp_401.close = mock.AsyncMock() - - mock_auth_req = mock.AsyncMock(return_value=mock_resp_401) - session = sessions.AsyncAuthorizedSession( - mock_creds, auth_request=mock_auth_req - ) - - task = asyncio.create_task(session.request("GET", "https://example.com")) - await refresh_started.wait() - task.cancel() - - with pytest.raises(asyncio.CancelledError): - await task - - mock_resp_401.close.assert_awaited_once() - await session.close() - - @pytest.mark.asyncio - async def test_request_401_concurrent_refreshes_are_deduplicated(self): - """Verifies that concurrent 401s execute only one credentials.refresh call.""" - mock_creds = mock.AsyncMock(spec=credentials.Credentials) - mock_creds.before_request = mock.AsyncMock(return_value=None) - - mock_resp_401 = mock.Mock() - mock_resp_401.status_code = http_client.UNAUTHORIZED - mock_resp_401.close = mock.AsyncMock() - - mock_resp_200 = mock.Mock() - mock_resp_200.status_code = http_client.OK - mock_resp_200.close = mock.AsyncMock() - - # Both concurrent requests get 401 initially, then 200 on retry - mock_auth_req = mock.AsyncMock( - side_effect=[mock_resp_401, mock_resp_401, mock_resp_200, mock_resp_200] - ) - session = sessions.AsyncAuthorizedSession( - mock_creds, auth_request=mock_auth_req - ) - - refresh_count = 0 - - async def slow_refresh(*args, **kwargs): - nonlocal refresh_count - refresh_count += 1 - await asyncio.sleep(0.05) - - mock_creds.refresh = mock.AsyncMock(side_effect=slow_refresh) - - results = await asyncio.gather( - session.request("GET", "https://example.com/1"), - session.request("GET", "https://example.com/2"), - ) - - assert all(r.status_code == 200 for r in results) - assert refresh_count == 1 - await session.close() + mock_check.return_value = (b"new", b"new", b"old_fp", b"new_fp") + + resp = await session.request( + "GET", "https://pubsub.mtls.googleapis.com/test" + ) + + assert resp == mock_resp + assert mock_auth_req.call_count == 3 + assert mock_check.call_count == 2 + assert mock_resp.close.call_count == 2 + + await session.close() + + @pytest.mark.asyncio + async def test_session_close_cleans_old_auth_requests(self): + """Covers the loop in the `close()` method that drains `_old_auth_requests`.""" + mock_creds = mock.AsyncMock(spec=credentials.Credentials) + session = sessions.AsyncAuthorizedSession( + mock_creds, auth_request=mock.AsyncMock() + ) + + mock_old_req_1 = mock.AsyncMock() + mock_old_req_2 = mock.AsyncMock() + mock_old_req_3_fails = mock.AsyncMock() + mock_old_req_3_fails.close.side_effect = Exception("Close error") + + session._old_auth_requests.extend( + [mock_old_req_1, mock_old_req_2, mock_old_req_3_fails] + ) + + await session.close() + + mock_old_req_1.close.assert_called_once() + mock_old_req_2.close.assert_called_once() + mock_old_req_3_fails.close.assert_called_once() + assert len(session._old_auth_requests) == 0 + + @pytest.mark.asyncio + async def test_request_401_streaming_refreshes_creds_and_returns_open_response( + self, + ): + """Verifies that streaming requests refresh credentials but return the unclosed response.""" + mock_creds = mock.AsyncMock(spec=credentials.Credentials) + mock_creds.before_request = mock.AsyncMock(return_value=None) + mock_creds.refresh = mock.AsyncMock(return_value=None) + + mock_resp_401 = mock.Mock() + mock_resp_401.status_code = http_client.UNAUTHORIZED + mock_resp_401.close = mock.AsyncMock() + + mock_auth_req = mock.AsyncMock(return_value=mock_resp_401) + session = sessions.AsyncAuthorizedSession( + mock_creds, auth_request=mock_auth_req + ) + + streaming_data = (chunk for chunk in [b"chunk1", b"chunk2"]) + response = await session.request( + "POST", "https://example.com", data=streaming_data + ) + + assert response == mock_resp_401 + mock_creds.refresh.assert_awaited_once() + mock_resp_401.close.assert_not_called() + await session.close() + + @pytest.mark.asyncio + async def test_request_401_closes_response_on_timeout_during_recovery(self): + """Verifies that response is closed when auth_with_timeout times out during _recover_auth_state.""" + mock_creds = mock.AsyncMock(spec=credentials.Credentials) + mock_creds.before_request = mock.AsyncMock(return_value=None) + + async def slow_refresh(*args, **kwargs): + await asyncio.sleep(10) + + mock_creds.refresh = mock.AsyncMock(side_effect=slow_refresh) + + mock_resp_401 = mock.Mock() + mock_resp_401.status_code = http_client.UNAUTHORIZED + mock_resp_401.close = mock.AsyncMock() + + mock_auth_req = mock.AsyncMock(return_value=mock_resp_401) + session = sessions.AsyncAuthorizedSession( + mock_creds, auth_request=mock_auth_req + ) + + with pytest.raises(TimeoutError): + await session.request("GET", "https://example.com", max_allowed_time=0.01) + + mock_resp_401.close.assert_awaited_once() + await session.close() + + @pytest.mark.asyncio + async def test_request_401_closes_response_on_cancellation(self): + """Verifies that response is closed and CancelledError propagated if task is cancelled.""" + mock_creds = mock.AsyncMock(spec=credentials.Credentials) + mock_creds.before_request = mock.AsyncMock(return_value=None) + + refresh_started = asyncio.Event() + + async def cancel_on_refresh(*args, **kwargs): + refresh_started.set() + await asyncio.sleep(10) + + mock_creds.refresh = mock.AsyncMock(side_effect=cancel_on_refresh) + + mock_resp_401 = mock.Mock() + mock_resp_401.status_code = http_client.UNAUTHORIZED + mock_resp_401.close = mock.AsyncMock() + + mock_auth_req = mock.AsyncMock(return_value=mock_resp_401) + session = sessions.AsyncAuthorizedSession( + mock_creds, auth_request=mock_auth_req + ) + + task = asyncio.create_task(session.request("GET", "https://example.com")) + await refresh_started.wait() + task.cancel() + + with pytest.raises(asyncio.CancelledError): + await task + + mock_resp_401.close.assert_awaited_once() + await session.close() + + @pytest.mark.asyncio + async def test_request_401_concurrent_refreshes_are_deduplicated(self): + """Verifies that concurrent 401s execute only one credentials.refresh call.""" + mock_creds = mock.AsyncMock(spec=credentials.Credentials) + mock_creds.before_request = mock.AsyncMock(return_value=None) + + mock_resp_401 = mock.Mock() + mock_resp_401.status_code = http_client.UNAUTHORIZED + mock_resp_401.close = mock.AsyncMock() + + mock_resp_200 = mock.Mock() + mock_resp_200.status_code = http_client.OK + mock_resp_200.close = mock.AsyncMock() + + mock_auth_req = mock.AsyncMock( + side_effect=[ + mock_resp_401, + mock_resp_401, + mock_resp_200, + mock_resp_200, + ] + ) + session = sessions.AsyncAuthorizedSession( + mock_creds, auth_request=mock_auth_req + ) + + refresh_count = 0 + + async def slow_refresh(*args, **kwargs): + nonlocal refresh_count + refresh_count += 1 + await asyncio.sleep(0.05) + + mock_creds.refresh = mock.AsyncMock(side_effect=slow_refresh) + + results = await asyncio.gather( + session.request("GET", "https://example.com/1"), + session.request("GET", "https://example.com/2"), + ) + + assert all(r.status_code == 200 for r in results) + assert refresh_count == 1 + await session.close() From 377f6be17e57c78bca607397c46f232fd5b52abe Mon Sep 17 00:00:00 2001 From: agrawalradhika-cell Date: Thu, 3 Sep 2026 10:14:26 -0700 Subject: [PATCH 71/79] fix: Add unit tests for MTLS parameter checking --- .../tests/transport/aio/test_mtls.py | 196 ++++++++++++++++++ 1 file changed, 196 insertions(+) create mode 100644 packages/google-auth/tests/transport/aio/test_mtls.py diff --git a/packages/google-auth/tests/transport/aio/test_mtls.py b/packages/google-auth/tests/transport/aio/test_mtls.py new file mode 100644 index 000000000000..2c977c0e6fdb --- /dev/null +++ b/packages/google-auth/tests/transport/aio/test_mtls.py @@ -0,0 +1,196 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from unittest import mock +from google.auth import exceptions +from google.auth.aio.transport import mtls +import pytest + +CERT_BYTES = b"-----BEGIN CERTIFICATE-----\nMIID...CERT1...=\n-----END CERTIFICATE-----\n" +KEY_BYTES = ( + b"-----BEGIN PRIVATE KEY-----\nMIIE...KEY1...==\n-----END PRIVATE KEY-----\n" +) +NEW_CERT_BYTES = b"-----BEGIN CERTIFICATE-----\nMIID...CERT2...=\n-----END CERTIFICATE-----\n" +NEW_KEY_BYTES = ( + b"-----BEGIN PRIVATE KEY-----\nMIIE...KEY2...==\n-----END PRIVATE KEY-----\n" +) + + +@pytest.mark.asyncio +async def test_check_parameters_no_client_cert(): + """Test when no certificate is discovered (has_cert is False).""" + with mock.patch.object( + mtls, "get_client_cert_and_key", return_value=(False, None, None) + ): + cert, key, cached_fp, current_fp = ( + await mtls.check_parameters_for_unauthorized_response( + cached_cert=b"stale_cert", client_cert_callback=None + ) + ) + + assert cert is None + assert key is None + assert cached_fp is None + assert current_fp is None + + +@pytest.mark.asyncio +async def test_check_parameters_cert_matched(): + """Test when newly retrieved certificate matches the cached certificate.""" + + def callback(): + return CERT_BYTES, KEY_BYTES + + with mock.patch( + "google.auth.transport._agent_identity_utils.parse_certificate" + ) as mock_parse, mock.patch( + "google.auth.transport._agent_identity_utils.calculate_certificate_fingerprint", + return_value="FINGERPRINT_A", + ), mock.patch( + "google.auth.transport._agent_identity_utils.get_cached_cert_fingerprint", + return_value="FINGERPRINT_A", + ): + + cert, key, cached_fp, current_fp = ( + await mtls.check_parameters_for_unauthorized_response( + cached_cert=CERT_BYTES, client_cert_callback=callback + ) + ) + + assert cert == CERT_BYTES + assert key == KEY_BYTES + assert cached_fp == "FINGERPRINT_A" + assert current_fp == "FINGERPRINT_A" + assert cached_fp == current_fp # Indicates no cert rotation needed + + +@pytest.mark.asyncio +async def test_check_parameters_cert_mismatch_rotation(): + """Test when newly retrieved certificate differs from the cached certificate (rotation occurred).""" + + def callback(): + return NEW_CERT_BYTES, NEW_KEY_BYTES + + with mock.patch( + "google.auth.transport._agent_identity_utils.parse_certificate" + ) as mock_parse, mock.patch( + "google.auth.transport._agent_identity_utils.calculate_certificate_fingerprint", + return_value="FINGERPRINT_NEW", + ), mock.patch( + "google.auth.transport._agent_identity_utils.get_cached_cert_fingerprint", + return_value="FINGERPRINT_OLD", + ): + + cert, key, cached_fp, current_fp = ( + await mtls.check_parameters_for_unauthorized_response( + cached_cert=CERT_BYTES, client_cert_callback=callback + ) + ) + + assert cert == NEW_CERT_BYTES + assert key == NEW_KEY_BYTES + assert cached_fp == "FINGERPRINT_OLD" + assert current_fp == "FINGERPRINT_NEW" + assert cached_fp != current_fp # Indicates cert rotation required + + +@pytest.mark.asyncio +async def test_check_parameters_without_cached_cert(): + """Test when cached_cert is None.""" + + def callback(): + return CERT_BYTES, KEY_BYTES + + with mock.patch( + "google.auth.transport._agent_identity_utils.parse_certificate" + ), mock.patch( + "google.auth.transport._agent_identity_utils.calculate_certificate_fingerprint", + return_value="FINGERPRINT_CURRENT", + ), mock.patch( + "google.auth.transport._agent_identity_utils.get_cached_cert_fingerprint" + ) as mock_get_cached: + + cert, key, cached_fp, current_fp = ( + await mtls.check_parameters_for_unauthorized_response( + cached_cert=None, client_cert_callback=callback + ) + ) + + assert cert == CERT_BYTES + assert key == KEY_BYTES + assert cached_fp == "FINGERPRINT_CURRENT" + assert current_fp == "FINGERPRINT_CURRENT" + # Should not attempt to parse a None cached cert + mock_get_cached.assert_not_called() + + +@pytest.mark.asyncio +async def test_check_parameters_executor_fingerprint_computation(): + """Test that fingerprint computation is properly offloaded to the executor.""" + + def callback(): + return CERT_BYTES, KEY_BYTES + + with mock.patch.object( + mtls, "_run_in_executor", wraps=mtls._run_in_executor + ) as mock_run_in_executor, mock.patch( + "google.auth.transport._agent_identity_utils.parse_certificate" + ), mock.patch( + "google.auth.transport._agent_identity_utils.calculate_certificate_fingerprint", + return_value="FP_CURRENT", + ), mock.patch( + "google.auth.transport._agent_identity_utils.get_cached_cert_fingerprint", + return_value="FP_CACHED", + ): + + cert, key, cached_fp, current_fp = ( + await mtls.check_parameters_for_unauthorized_response( + cached_cert=CERT_BYTES, client_cert_callback=callback + ) + ) + + assert mock_run_in_executor.called + assert cert == CERT_BYTES + assert cached_fp == "FP_CACHED" + assert current_fp == "FP_CURRENT" + + +@pytest.mark.asyncio +async def test_check_parameters_callback_exception_propagation(): + """Test that exceptions raised by client_cert_callback propagate cleanly.""" + + def failing_callback(): + raise exceptions.ClientCertError("Client cert provider failed") + + with pytest.raises( + exceptions.ClientCertError, match="Client cert provider failed" + ): + await mtls.check_parameters_for_unauthorized_response( + cached_cert=CERT_BYTES, client_cert_callback=failing_callback + ) + + +@pytest.mark.asyncio +async def test_check_parameters_async_callback_exception_propagation(): + """Test that exceptions raised in an async client_cert_callback propagate cleanly.""" + + async def failing_async_callback(): + raise OSError("Disk read error while loading certificates") + + with pytest.raises( + OSError, match="Disk read error while loading certificates" + ): + await mtls.check_parameters_for_unauthorized_response( + cached_cert=CERT_BYTES, client_cert_callback=failing_async_callback + ) From 86640b28923148d548b42a42051b2943a928a7c9 Mon Sep 17 00:00:00 2001 From: agrawalradhika-cell Date: Thu, 3 Sep 2026 10:42:32 -0700 Subject: [PATCH 72/79] fix: Refactor mock patches in test_mtls.py --- .../tests/transport/aio/test_mtls.py | 87 ++++++++++--------- 1 file changed, 46 insertions(+), 41 deletions(-) diff --git a/packages/google-auth/tests/transport/aio/test_mtls.py b/packages/google-auth/tests/transport/aio/test_mtls.py index 2c977c0e6fdb..20a66380b2bd 100644 --- a/packages/google-auth/tests/transport/aio/test_mtls.py +++ b/packages/google-auth/tests/transport/aio/test_mtls.py @@ -13,6 +13,7 @@ # limitations under the License. from unittest import mock + from google.auth import exceptions from google.auth.aio.transport import mtls import pytest @@ -52,16 +53,17 @@ async def test_check_parameters_cert_matched(): def callback(): return CERT_BYTES, KEY_BYTES - with mock.patch( - "google.auth.transport._agent_identity_utils.parse_certificate" - ) as mock_parse, mock.patch( - "google.auth.transport._agent_identity_utils.calculate_certificate_fingerprint", - return_value="FINGERPRINT_A", - ), mock.patch( - "google.auth.transport._agent_identity_utils.get_cached_cert_fingerprint", - return_value="FINGERPRINT_A", + with ( + mock.patch("google.auth._agent_identity_utils.parse_certificate"), + mock.patch( + "google.auth._agent_identity_utils.calculate_certificate_fingerprint", + return_value="FINGERPRINT_A", + ), + mock.patch( + "google.auth._agent_identity_utils.get_cached_cert_fingerprint", + return_value="FINGERPRINT_A", + ), ): - cert, key, cached_fp, current_fp = ( await mtls.check_parameters_for_unauthorized_response( cached_cert=CERT_BYTES, client_cert_callback=callback @@ -72,7 +74,7 @@ def callback(): assert key == KEY_BYTES assert cached_fp == "FINGERPRINT_A" assert current_fp == "FINGERPRINT_A" - assert cached_fp == current_fp # Indicates no cert rotation needed + assert cached_fp == current_fp @pytest.mark.asyncio @@ -82,16 +84,17 @@ async def test_check_parameters_cert_mismatch_rotation(): def callback(): return NEW_CERT_BYTES, NEW_KEY_BYTES - with mock.patch( - "google.auth.transport._agent_identity_utils.parse_certificate" - ) as mock_parse, mock.patch( - "google.auth.transport._agent_identity_utils.calculate_certificate_fingerprint", - return_value="FINGERPRINT_NEW", - ), mock.patch( - "google.auth.transport._agent_identity_utils.get_cached_cert_fingerprint", - return_value="FINGERPRINT_OLD", + with ( + mock.patch("google.auth._agent_identity_utils.parse_certificate"), + mock.patch( + "google.auth._agent_identity_utils.calculate_certificate_fingerprint", + return_value="FINGERPRINT_NEW", + ), + mock.patch( + "google.auth._agent_identity_utils.get_cached_cert_fingerprint", + return_value="FINGERPRINT_OLD", + ), ): - cert, key, cached_fp, current_fp = ( await mtls.check_parameters_for_unauthorized_response( cached_cert=CERT_BYTES, client_cert_callback=callback @@ -102,7 +105,7 @@ def callback(): assert key == NEW_KEY_BYTES assert cached_fp == "FINGERPRINT_OLD" assert current_fp == "FINGERPRINT_NEW" - assert cached_fp != current_fp # Indicates cert rotation required + assert cached_fp != current_fp @pytest.mark.asyncio @@ -112,15 +115,16 @@ async def test_check_parameters_without_cached_cert(): def callback(): return CERT_BYTES, KEY_BYTES - with mock.patch( - "google.auth.transport._agent_identity_utils.parse_certificate" - ), mock.patch( - "google.auth.transport._agent_identity_utils.calculate_certificate_fingerprint", - return_value="FINGERPRINT_CURRENT", - ), mock.patch( - "google.auth.transport._agent_identity_utils.get_cached_cert_fingerprint" - ) as mock_get_cached: - + with ( + mock.patch("google.auth._agent_identity_utils.parse_certificate"), + mock.patch( + "google.auth._agent_identity_utils.calculate_certificate_fingerprint", + return_value="FINGERPRINT_CURRENT", + ), + mock.patch( + "google.auth._agent_identity_utils.get_cached_cert_fingerprint" + ) as mock_get_cached, + ): cert, key, cached_fp, current_fp = ( await mtls.check_parameters_for_unauthorized_response( cached_cert=None, client_cert_callback=callback @@ -131,7 +135,6 @@ def callback(): assert key == KEY_BYTES assert cached_fp == "FINGERPRINT_CURRENT" assert current_fp == "FINGERPRINT_CURRENT" - # Should not attempt to parse a None cached cert mock_get_cached.assert_not_called() @@ -142,18 +145,20 @@ async def test_check_parameters_executor_fingerprint_computation(): def callback(): return CERT_BYTES, KEY_BYTES - with mock.patch.object( - mtls, "_run_in_executor", wraps=mtls._run_in_executor - ) as mock_run_in_executor, mock.patch( - "google.auth.transport._agent_identity_utils.parse_certificate" - ), mock.patch( - "google.auth.transport._agent_identity_utils.calculate_certificate_fingerprint", - return_value="FP_CURRENT", - ), mock.patch( - "google.auth.transport._agent_identity_utils.get_cached_cert_fingerprint", - return_value="FP_CACHED", + with ( + mock.patch.object( + mtls, "_run_in_executor", wraps=mtls._run_in_executor + ) as mock_run_in_executor, + mock.patch("google.auth._agent_identity_utils.parse_certificate"), + mock.patch( + "google.auth._agent_identity_utils.calculate_certificate_fingerprint", + return_value="FP_CURRENT", + ), + mock.patch( + "google.auth._agent_identity_utils.get_cached_cert_fingerprint", + return_value="FP_CACHED", + ), ): - cert, key, cached_fp, current_fp = ( await mtls.check_parameters_for_unauthorized_response( cached_cert=CERT_BYTES, client_cert_callback=callback From fbb9e475fa60539b14ea9d78f747c98f0162e491 Mon Sep 17 00:00:00 2001 From: Radhika Agrawal Date: Thu, 3 Sep 2026 17:43:27 +0000 Subject: [PATCH 73/79] Fix: fix lint and unit tests Signed-off-by: Radhika Agrawal --- .../google/auth/aio/transport/sessions.py | 22 +- .../tests/transport/aio/test_sessions_mtls.py | 1911 ++++++++--------- 2 files changed, 954 insertions(+), 979 deletions(-) diff --git a/packages/google-auth/google/auth/aio/transport/sessions.py b/packages/google-auth/google/auth/aio/transport/sessions.py index dff59da9b112..cb7fdf9361b9 100644 --- a/packages/google-auth/google/auth/aio/transport/sessions.py +++ b/packages/google-auth/google/auth/aio/transport/sessions.py @@ -43,7 +43,11 @@ ClientTimeout = None _LOGGER = logging.getLogger(__name__) -_MTLS_URL_PREFIXES = ["mtls.googleapis.com", "mtls.sandbox.googleapis.com", ".p.googleapis.com"] +_MTLS_URL_PREFIXES = [ + "mtls.googleapis.com", + "mtls.sandbox.googleapis.com", + ".p.googleapis.com", +] # Tracks the internal aiohttp installation and usage try: @@ -221,14 +225,14 @@ async def _do_configure(): old_auth_request = self._auth_request self._auth_request = AiohttpRequest(session=new_session) while len(self._old_auth_requests) >= 2: - oldest_auth_request = self._old_auth_requests.pop(0) - try: - if hasattr(oldest_auth_request, "close"): - res = oldest_auth_request.close() - if inspect.isawaitable(res): - await res - except Exception: - pass + oldest_auth_request = self._old_auth_requests.pop(0) + try: + if hasattr(oldest_auth_request, "close"): + res = oldest_auth_request.close() + if inspect.isawaitable(res): + await res + except Exception: + pass self._old_auth_requests.append(old_auth_request) diff --git a/packages/google-auth/tests/transport/aio/test_sessions_mtls.py b/packages/google-auth/tests/transport/aio/test_sessions_mtls.py index e1bfddbcc235..0e7f3bff1787 100644 --- a/packages/google-auth/tests/transport/aio/test_sessions_mtls.py +++ b/packages/google-auth/tests/transport/aio/test_sessions_mtls.py @@ -18,13 +18,12 @@ import os import ssl from unittest import mock - +import pytest from google.auth import exceptions from google.auth.aio import credentials from google.auth.aio import transport from google.auth.aio.transport import sessions from google.auth.exceptions import TimeoutError -import pytest # This is the valid "workload" format the library expects VALID_WORKLOAD_CONFIG = { @@ -39,972 +38,944 @@ class TestSessionsMtls: - - @pytest.mark.asyncio - async def test_configure_mtls_channel(self): - """Tests that the mTLS channel configures correctly when a valid workload config is mocked.""" - with ( - mock.patch.dict( - os.environ, {"GOOGLE_API_USE_CLIENT_CERTIFICATE": "true"} - ), - mock.patch("os.path.exists") as mock_exists, - mock.patch( - "builtins.open", - mock.mock_open(read_data=json.dumps(VALID_WORKLOAD_CONFIG)), - ), - mock.patch( - "google.auth.aio.transport.mtls.get_client_cert_and_key" - ) as mock_helper, - mock.patch( - "google.auth.aio.transport.mtls.make_client_cert_ssl_context" - ) as mock_make_context, - mock.patch("aiohttp.TCPConnector") as mock_connector, - mock.patch("aiohttp.ClientSession") as mock_session, - ): - mock_session.return_value.close = mock.AsyncMock() - mock_exists.return_value = True - mock_helper.return_value = (True, b"fake_cert_data", b"fake_key_data") - - mock_context = mock.Mock(spec=ssl.SSLContext) - mock_make_context.return_value = mock_context - - mock_creds = mock.AsyncMock(spec=credentials.Credentials) - session = sessions.AsyncAuthorizedSession(mock_creds) - - await session.configure_mtls_channel() - - assert session._is_mtls is True - mock_make_context.assert_called_once_with( - b"fake_cert_data", b"fake_key_data" - ) - mock_connector.assert_called_once_with(ssl=mock_context) - mock_session.assert_called_once_with( - connector=mock_connector.return_value - ) - await session.close() - - @pytest.mark.asyncio - async def test_configure_mtls_channel_disabled(self): - """Tests behavior when the config file does not exist.""" - with ( - mock.patch.dict( - os.environ, {"GOOGLE_API_USE_CLIENT_CERTIFICATE": "true"} - ), - mock.patch("os.path.exists") as mock_exists, - ): - mock_exists.return_value = False - mock_creds = mock.AsyncMock(spec=credentials.Credentials) - session = sessions.AsyncAuthorizedSession(mock_creds) - await session.configure_mtls_channel() - assert session._is_mtls is False - await session.close() - - @pytest.mark.asyncio - async def test_configure_mtls_channel_invalid_format(self): - """Verifies that the MutualTLSChannelError is raised for bad formats.""" - with ( - mock.patch.dict( - os.environ, {"GOOGLE_API_USE_CLIENT_CERTIFICATE": "true"} - ), - mock.patch("os.path.exists") as mock_exists, - mock.patch( - "builtins.open", mock.mock_open(read_data='{"invalid": "format"}') - ), - ): - mock_exists.return_value = True - mock_creds = mock.AsyncMock(spec=credentials.Credentials) - session = sessions.AsyncAuthorizedSession(mock_creds) - - with pytest.raises(exceptions.MutualTLSChannelError): - await session.configure_mtls_channel() - await session.close() - - @pytest.mark.asyncio - async def test_configure_mtls_channel_invalid_fields(self): - """If cert is missing expected keys, it should fail gracefully.""" - with ( - mock.patch.dict( - os.environ, {"GOOGLE_API_USE_CLIENT_CERTIFICATE": "true"} - ), - mock.patch("os.path.exists") as mock_exists, - mock.patch( - "builtins.open", mock.mock_open(read_data='{"cert_configs": {}}') - ), - ): - mock_exists.return_value = True - mock_creds = mock.AsyncMock(spec=credentials.Credentials) - session = sessions.AsyncAuthorizedSession(mock_creds) - await session.configure_mtls_channel() - assert session._is_mtls is False - await session.close() - - @pytest.mark.asyncio - async def test_configure_mtls_channel_mock_callback(self): - """Tests mTLS configuration using bytes-returning callback.""" - - def mock_callback(): - return (b"fake_cert_bytes", b"fake_key_bytes") - - with ( - mock.patch.dict( - os.environ, {"GOOGLE_API_USE_CLIENT_CERTIFICATE": "true"} - ), - mock.patch( - "google.auth.transport.mtls.has_default_client_cert_source", - return_value=True, - ), - mock.patch( - "google.auth.aio.transport.mtls.make_client_cert_ssl_context" - ) as mock_make_context, - mock.patch("aiohttp.TCPConnector") as mock_connector, - mock.patch("aiohttp.ClientSession") as mock_session, - ): - mock_session.return_value.close = mock.AsyncMock() - mock_context = mock.Mock(spec=ssl.SSLContext) - mock_make_context.return_value = mock_context - - mock_creds = mock.AsyncMock(spec=credentials.Credentials) - session = sessions.AsyncAuthorizedSession(mock_creds) - - await session.configure_mtls_channel(client_cert_callback=mock_callback) - - assert session._is_mtls is True - mock_make_context.assert_called_once_with( - b"fake_cert_bytes", b"fake_key_bytes" - ) - mock_connector.assert_called_once_with(ssl=mock_context) - mock_session.assert_called_once_with( - connector=mock_connector.return_value - ) - await session.close() - - @pytest.mark.asyncio - async def test_configure_mtls_channel_custom_request(self): - """Tests that if _auth_request is not an AiohttpRequest, _is_mtls is set to False.""" - with ( - mock.patch.dict( - os.environ, {"GOOGLE_API_USE_CLIENT_CERTIFICATE": "true"} - ), - mock.patch("os.path.exists") as mock_exists, - mock.patch( - "builtins.open", - mock.mock_open(read_data=json.dumps(VALID_WORKLOAD_CONFIG)), - ), - mock.patch( - "google.auth.aio.transport.mtls.get_client_cert_and_key" - ) as mock_helper, - mock.patch( - "google.auth.aio.transport.mtls.make_client_cert_ssl_context" - ) as mock_make_context, - ): - mock_exists.return_value = True - mock_helper.return_value = (True, b"fake_cert_data", b"fake_key_data") - - mock_context = mock.Mock(spec=ssl.SSLContext) - mock_make_context.return_value = mock_context - - mock_creds = mock.AsyncMock(spec=credentials.Credentials) - mock_auth_request = mock.AsyncMock(spec=transport.Request) - session = sessions.AsyncAuthorizedSession( - mock_creds, auth_request=mock_auth_request - ) - - with pytest.warns(UserWarning, match="Attempted to establish mTLS"): - await session.configure_mtls_channel() - - assert session._is_mtls is False - mock_make_context.assert_not_called() - await session.close() - - @pytest.mark.asyncio - async def test_configure_mtls_channel_exception_resets_flag(self): - """Tests that self._is_mtls is reset to False if an exception is raised during configuration.""" - with ( - mock.patch.dict( - os.environ, {"GOOGLE_API_USE_CLIENT_CERTIFICATE": "true"} - ), - mock.patch("os.path.exists") as mock_exists, - mock.patch( - "builtins.open", - mock.mock_open(read_data=json.dumps(VALID_WORKLOAD_CONFIG)), - ), - mock.patch( - "google.auth.aio.transport.mtls.get_client_cert_and_key" - ) as mock_helper, - mock.patch( - "google.auth.aio.transport.mtls.make_client_cert_ssl_context" - ) as mock_make_context, - ): - mock_exists.return_value = True - mock_helper.return_value = (True, b"fake_cert_data", b"fake_key_data") - mock_make_context.side_effect = exceptions.ClientCertError("Mock error") - - mock_creds = mock.AsyncMock(spec=credentials.Credentials) - session = sessions.AsyncAuthorizedSession(mock_creds) - - with pytest.raises(exceptions.MutualTLSChannelError): - await session.configure_mtls_channel() - - assert session._is_mtls is False - await session.close() - - @pytest.mark.asyncio - async def test_configure_mtls_channel_transport_error_resets_flag(self): - """Tests that self._is_mtls is reset to False if a TransportError is raised.""" - with ( - mock.patch.dict( - os.environ, {"GOOGLE_API_USE_CLIENT_CERTIFICATE": "true"} - ), - mock.patch("os.path.exists") as mock_exists, - mock.patch( - "builtins.open", - mock.mock_open(read_data=json.dumps(VALID_WORKLOAD_CONFIG)), - ), - mock.patch( - "google.auth.aio.transport.mtls.get_client_cert_and_key" - ) as mock_helper, - mock.patch( - "google.auth.aio.transport.mtls.make_client_cert_ssl_context" - ) as mock_make_context, - ): - mock_exists.return_value = True - mock_helper.return_value = (True, b"fake_cert_data", b"fake_key_data") - mock_make_context.side_effect = exceptions.TransportError("Mock error") - - mock_creds = mock.AsyncMock(spec=credentials.Credentials) - session = sessions.AsyncAuthorizedSession(mock_creds) - - with pytest.raises(exceptions.MutualTLSChannelError): - await session.configure_mtls_channel() - - assert session._is_mtls is False - await session.close() - - @pytest.mark.asyncio - async def test_configure_mtls_channel_atomic_on_exception(self): - """Tests that if configure_mtls_channel already succeeded, a subsequent failure preserves state.""" - # Step 1: Successful configuration - with ( - mock.patch.dict( - os.environ, {"GOOGLE_API_USE_CLIENT_CERTIFICATE": "true"} - ), - mock.patch("os.path.exists") as mock_exists, - mock.patch( - "builtins.open", - mock.mock_open(read_data=json.dumps(VALID_WORKLOAD_CONFIG)), - ), - mock.patch( - "google.auth.aio.transport.mtls.get_client_cert_and_key" - ) as mock_helper, - mock.patch( - "google.auth.aio.transport.mtls.make_client_cert_ssl_context" - ) as mock_make_context, - mock.patch("aiohttp.TCPConnector"), - mock.patch("aiohttp.ClientSession") as mock_session, - ): - mock_session.return_value.close = mock.AsyncMock() - mock_exists.return_value = True - mock_helper.return_value = ( - True, - b"fake_cert_data_1", - b"fake_key_data_1", - ) - - mock_context = mock.Mock(spec=ssl.SSLContext) - mock_make_context.return_value = mock_context - - mock_creds = mock.AsyncMock(spec=credentials.Credentials) - session = sessions.AsyncAuthorizedSession(mock_creds) - - await session.configure_mtls_channel() - assert session._is_mtls is True - assert session._cached_cert == b"fake_cert_data_1" - first_auth_request = session._auth_request - - # Step 2: Failed subsequent configuration attempt - session._mtls_init_task = None - mock_make_context.side_effect = exceptions.ClientCertError("Mock error") - - with pytest.raises(exceptions.MutualTLSChannelError): - await session.configure_mtls_channel() - - assert session._is_mtls is True - assert session._cached_cert == b"fake_cert_data_1" - assert session._auth_request is first_auth_request - await session.close() - - @pytest.mark.asyncio - async def test_configure_mtls_channel_close_exception_does_not_abort(self): - """Tests that an exception in old_auth_request.close() does not abort configuration.""" - with ( - mock.patch.dict( - os.environ, {"GOOGLE_API_USE_CLIENT_CERTIFICATE": "true"} - ), - mock.patch("os.path.exists") as mock_exists, - mock.patch( - "builtins.open", - mock.mock_open(read_data=json.dumps(VALID_WORKLOAD_CONFIG)), - ), - mock.patch( - "google.auth.aio.transport.mtls.get_client_cert_and_key" - ) as mock_helper, - mock.patch( - "google.auth.aio.transport.mtls.make_client_cert_ssl_context" - ) as mock_make_context, - mock.patch("aiohttp.TCPConnector"), - mock.patch("aiohttp.ClientSession") as mock_session, - ): - mock_session.return_value.close = mock.AsyncMock() - mock_exists.return_value = True - mock_helper.return_value = (True, b"fake_cert_data", b"fake_key_data") - - mock_context = mock.Mock(spec=ssl.SSLContext) - mock_make_context.return_value = mock_context - - mock_creds = mock.AsyncMock(spec=credentials.Credentials) - session = sessions.AsyncAuthorizedSession(mock_creds) - - session._auth_request.close = mock.AsyncMock( - side_effect=Exception("Mock close error") - ) - - await session.configure_mtls_channel() - - assert session._is_mtls is True - assert session._cached_cert == b"fake_cert_data" - await session.close() - - @pytest.mark.asyncio - async def test_cert_rotation_failure_raises_error(self): - mock_creds = mock.AsyncMock(spec=credentials.Credentials) - mock_creds.before_request = mock.AsyncMock(return_value=None) - - mock_resp = mock.Mock() - mock_resp.status_code = http_client.UNAUTHORIZED - mock_auth_req = mock.AsyncMock(return_value=mock_resp) - - session = sessions.AsyncAuthorizedSession( - mock_creds, auth_request=mock_auth_req - ) - session._is_mtls = True - session._cached_cert = b"old_cert" - - new_cert = b"new_cert" - new_key = b"new_key" - - with ( - mock.patch( - "google.auth.aio.transport.mtls.check_parameters_for_unauthorized_response", - new_callable=mock.AsyncMock, - ) as mock_check, - mock.patch.object( - session, "configure_mtls_channel", new_callable=mock.AsyncMock - ) as mock_conf, - ): - mock_check.return_value = (new_cert, new_key, b"old_fp", b"new_fp") - mock_conf.side_effect = Exception("Failed to reconfigure") - - with pytest.raises(exceptions.MutualTLSChannelError): - await session.request("GET", "https://pubsub.mtls.googleapis.com/test") - - mock_check.assert_called_once() - mock_conf.assert_called_once() - - await session.close() - - @pytest.mark.asyncio - async def test_cert_rotation_check_params_fails(self): - mock_creds = mock.AsyncMock(spec=credentials.Credentials) - mock_creds.before_request = mock.AsyncMock(return_value=None) - mock_creds.refresh = mock.AsyncMock() - - mock_resp = mock.Mock() - mock_resp.status_code = http_client.UNAUTHORIZED - mock_auth_req = mock.AsyncMock(return_value=mock_resp) - - session = sessions.AsyncAuthorizedSession( - mock_creds, auth_request=mock_auth_req - ) - session._is_mtls = True - session._cached_cert = b"old_cert" - - with ( - mock.patch( - "google.auth.aio.transport.mtls.check_parameters_for_unauthorized_response", - new_callable=mock.AsyncMock, - ) as mock_check, - mock.patch.object( - session, "configure_mtls_channel", new_callable=mock.AsyncMock - ) as mock_conf, - ): - mock_check.side_effect = exceptions.MutualTLSChannelError( - "Failed to check params" - ) - - resp = await session.request( - "GET", "https://pubsub.mtls.googleapis.com/test" - ) - - assert resp == mock_resp - mock_check.assert_called_once() - mock_creds.refresh.assert_not_called() - mock_conf.assert_not_called() - - await session.close() - - @pytest.mark.asyncio - async def test_no_cert_rotation_when_cert_matches_and_mtls_enabled(self): - mock_creds = mock.AsyncMock(spec=credentials.Credentials) - mock_creds.before_request = mock.AsyncMock(return_value=None) - mock_creds.refresh = mock.AsyncMock(return_value=None) - - mock_resp_401 = mock.Mock() - mock_resp_401.status_code = http_client.UNAUTHORIZED - mock_resp_401.close = mock.AsyncMock() - - mock_resp_200 = mock.Mock() - mock_resp_200.status_code = http_client.OK - - # 401 on initial request, 200 on retry after refresh - mock_auth_req = mock.AsyncMock(side_effect=[mock_resp_401, mock_resp_200]) - - session = sessions.AsyncAuthorizedSession( - mock_creds, auth_request=mock_auth_req - ) - session._is_mtls = True - session._cached_cert = b"old_cert" - - new_cert = b"new_cert" - new_key = b"new_key" - - with ( - mock.patch( - "google.auth.aio.transport.mtls.check_parameters_for_unauthorized_response", - new_callable=mock.AsyncMock, - ) as mock_check, - mock.patch.object( - session, "configure_mtls_channel", new_callable=mock.AsyncMock - ) as mock_conf, - ): - # Matching fingerprints mean no mTLS rotation is needed - mock_check.return_value = (new_cert, new_key, b"old_fp", b"old_fp") - - resp = await session.request( - "GET", "https://pubsub.mtls.googleapis.com/test" - ) - - assert resp == mock_resp_200 - mock_check.assert_called_once() - mock_conf.assert_not_called() - mock_creds.refresh.assert_called_once() - assert mock_auth_req.call_count == 2 - mock_resp_401.close.assert_awaited_once() - - await session.close() - - @pytest.mark.asyncio - async def test_cert_rotation_success_and_retry(self): - mock_creds = mock.AsyncMock(spec=credentials.Credentials) - mock_creds.before_request = mock.AsyncMock(return_value=None) - mock_creds.refresh = mock.AsyncMock(return_value=None) - - mock_resp_401 = mock.Mock() - mock_resp_401.status_code = http_client.UNAUTHORIZED - mock_resp_200 = mock.Mock() - mock_resp_200.status_code = http_client.OK - - mock_auth_req = mock.AsyncMock(side_effect=[mock_resp_401, mock_resp_200]) - - session = sessions.AsyncAuthorizedSession( - mock_creds, auth_request=mock_auth_req - ) - session._is_mtls = True - session._cached_cert = b"old_cert" - - new_cert = b"new_cert" - new_key = b"new_key" - - with ( - mock.patch( - "google.auth.aio.transport.mtls.check_parameters_for_unauthorized_response", - new_callable=mock.AsyncMock, - ) as mock_check, - mock.patch.object( - session, "configure_mtls_channel", new_callable=mock.AsyncMock - ) as mock_conf, - ): - mock_check.return_value = (new_cert, new_key, b"old_fp", b"new_fp") - - resp = await session.request( - "GET", "https://pubsub.mtls.googleapis.com/test" - ) - - assert resp == mock_resp_200 - mock_check.assert_called_once() - mock_conf.assert_called_once_with(mock.ANY) - mock_creds.refresh.assert_called_once() - assert mock_creds.before_request.call_count == 2 - - await session.close() - - @pytest.mark.asyncio - async def test_cert_rotation_lock_contention(self): - mock_creds = mock.AsyncMock(spec=credentials.Credentials) - mock_creds.before_request = mock.AsyncMock(return_value=None) - mock_creds.refresh = mock.AsyncMock(return_value=None) - - mock_resp_401 = mock.Mock() - mock_resp_401.status_code = http_client.UNAUTHORIZED - mock_auth_req = mock.AsyncMock(return_value=mock_resp_401) - - session = sessions.AsyncAuthorizedSession( - mock_creds, auth_request=mock_auth_req - ) - session._is_mtls = True - session._cached_cert = b"old_cert" - - new_cert = b"new_cert" - new_key = b"new_key" - - async def mock_configure_mtls_channel(*args, **kwargs): - await asyncio.sleep(0.01) - session._cached_cert = new_cert - - async def mock_check_side_effect(cached_cert, callback=None): - if cached_cert == b"old_cert": - return (new_cert, new_key, b"old_fp", b"new_fp") - return (new_cert, new_key, b"new_fp", b"new_fp") - - with ( - mock.patch( - "google.auth.aio.transport.mtls.check_parameters_for_unauthorized_response", - new_callable=mock.AsyncMock, - ) as mock_check, - mock.patch.object( - session, "configure_mtls_channel", new_callable=mock.AsyncMock - ) as mock_conf, - ): - mock_check.side_effect = mock_check_side_effect - mock_conf.side_effect = mock_configure_mtls_channel - - tasks = [ - session.request("GET", "https://pubsub.mtls.googleapis.com/test") - for _ in range(3) - ] - responses = await asyncio.gather(*tasks) - - for resp in responses: - assert resp == mock_resp_401 - - mock_conf.assert_called_once() - - await session.close() - - @pytest.mark.asyncio - async def test_cert_rotation_lock_contention_no_cert_change(self): - mock_creds = mock.AsyncMock(spec=credentials.Credentials) - mock_creds.before_request = mock.AsyncMock(return_value=None) - mock_creds.refresh = mock.AsyncMock(return_value=None) - - mock_resp_401 = mock.Mock() - mock_resp_401.status_code = http_client.UNAUTHORIZED - mock_resp_200 = mock.Mock() - mock_resp_200.status_code = http_client.OK - - mock_auth_req = mock.AsyncMock( - side_effect=[mock_resp_401] * 3 + [mock_resp_200] * 3 - ) - - session = sessions.AsyncAuthorizedSession( - mock_creds, auth_request=mock_auth_req - ) - session._is_mtls = True - session._cached_cert = b"old_cert" - - async def mock_check_side_effect(cached_cert, callback=None): - await asyncio.sleep(0.01) - return (b"old_cert", b"old_key", b"old_fp", b"old_fp") - - with ( - mock.patch( - "google.auth.aio.transport.mtls.check_parameters_for_unauthorized_response", - new_callable=mock.AsyncMock, - ) as mock_check, - mock.patch.object( - session, "configure_mtls_channel", new_callable=mock.AsyncMock - ) as mock_conf, - ): - mock_check.side_effect = mock_check_side_effect - - tasks = [ - session.request("GET", "https://pubsub.mtls.googleapis.com/test") - for _ in range(3) - ] - responses = await asyncio.gather(*tasks) - - for resp in responses: - assert resp == mock_resp_200 - - mock_check.assert_called_once() - mock_conf.assert_not_called() - # Concurrent 401s properly deduplicate to 1 refresh - assert mock_creds.refresh.call_count == 1 - - await session.close() - - @pytest.mark.asyncio - async def test_psc_endpoint_triggers_cert_rotation(self): - """Verifies that PSC endpoints (*.p.googleapis.com) are recognized as mTLS endpoints.""" - mock_creds = mock.AsyncMock(spec=credentials.Credentials) - mock_creds.before_request = mock.AsyncMock(return_value=None) - mock_creds.refresh = mock.AsyncMock(return_value=None) - - mock_resp_401 = mock.Mock() - mock_resp_401.status_code = http_client.UNAUTHORIZED - mock_resp_200 = mock.Mock() - mock_resp_200.status_code = http_client.OK - - mock_auth_req = mock.AsyncMock(side_effect=[mock_resp_401, mock_resp_200]) - - session = sessions.AsyncAuthorizedSession( - mock_creds, auth_request=mock_auth_req - ) - session._is_mtls = True - session._cached_cert = b"old_cert" - - new_cert = b"new_cert" - new_key = b"new_key" - - with ( - mock.patch( - "google.auth.aio.transport.mtls.check_parameters_for_unauthorized_response", - new_callable=mock.AsyncMock, - ) as mock_check, - mock.patch.object( - session, "configure_mtls_channel", new_callable=mock.AsyncMock - ) as mock_conf, - ): - mock_check.return_value = (new_cert, new_key, b"old_fp", b"new_fp") - - resp = await session.request( - "GET", "https://pubsub.p.googleapis.com/test" - ) - - assert resp == mock_resp_200 - mock_check.assert_called_once() - mock_conf.assert_called_once_with(mock.ANY) - - await session.close() - - @pytest.mark.asyncio - async def test_non_mtls_url_bypasses_rotation(self): - mock_creds = mock.AsyncMock(spec=credentials.Credentials) - mock_creds.before_request = mock.AsyncMock(return_value=None) - mock_creds.refresh = mock.AsyncMock(return_value=None) - - mock_resp_401 = mock.Mock() - mock_resp_401.status_code = http_client.UNAUTHORIZED - mock_auth_req = mock.AsyncMock(return_value=mock_resp_401) - - session = sessions.AsyncAuthorizedSession( - mock_creds, auth_request=mock_auth_req - ) - - session._is_mtls = True - session._cached_cert = b"old_cert" - - with ( - mock.patch( - "google.auth.aio.transport.mtls.check_parameters_for_unauthorized_response", - new_callable=mock.AsyncMock, - ) as mock_check, - mock.patch.object( - session, "configure_mtls_channel", new_callable=mock.AsyncMock - ) as mock_conf, - ): - resp = await session.request("GET", "https://example.com/test") - - assert resp == mock_resp_401 - mock_check.assert_not_called() - mock_conf.assert_not_called() - assert mock_creds.refresh.call_count == 2 - assert mock_auth_req.call_count == 3 - - await session.close() - - @pytest.mark.asyncio - async def test_cert_rotation_skips_retry_for_streaming(self): - mock_creds = mock.AsyncMock(spec=credentials.Credentials) - mock_creds.before_request = mock.AsyncMock(return_value=None) - mock_creds.refresh = mock.AsyncMock() - - mock_resp = mock.Mock() - mock_resp.status_code = http_client.UNAUTHORIZED - mock_auth_req = mock.AsyncMock(return_value=mock_resp) - - session = sessions.AsyncAuthorizedSession( - mock_creds, auth_request=mock_auth_req - ) - session._is_mtls = True - session._cached_cert = b"old_cert" - - class MockStream: - - def read(self): - pass - - with ( - mock.patch( - "google.auth.aio.transport.mtls.check_parameters_for_unauthorized_response", - new_callable=mock.AsyncMock, - ) as mock_check, - mock.patch.object( - session, "configure_mtls_channel", new_callable=mock.AsyncMock - ) as mock_conf, - ): - mock_check.return_value = (b"new", b"new", b"old_fp", b"new_fp") - - resp = await session.request( - "GET", "https://pubsub.mtls.googleapis.com/test", data=MockStream() - ) - - assert resp == mock_resp - mock_conf.assert_called_once() - - mock_creds.refresh.assert_called_once() - await session.close() - - @pytest.mark.asyncio - async def test_cert_rotation_credential_refresh_fails(self): - """Covers the except block for RefreshError when credentials fail to refresh.""" - mock_creds = mock.AsyncMock(spec=credentials.Credentials) - mock_creds.before_request = mock.AsyncMock(return_value=None) - mock_creds.refresh = mock.AsyncMock( - side_effect=exceptions.RefreshError("Refresh failed") - ) - - mock_resp = mock.Mock() - mock_resp.status_code = http_client.UNAUTHORIZED - mock_auth_req = mock.AsyncMock(return_value=mock_resp) - - session = sessions.AsyncAuthorizedSession( - mock_creds, auth_request=mock_auth_req - ) - session._is_mtls = True - session._cached_cert = b"old_cert" - - with ( - mock.patch( - "google.auth.aio.transport.mtls.check_parameters_for_unauthorized_response", - new_callable=mock.AsyncMock, - ) as mock_check, - mock.patch.object( - session, "configure_mtls_channel", new_callable=mock.AsyncMock - ), - ): - mock_check.return_value = (b"new", b"new", b"old_fp", b"new_fp") - - resp = await session.request( - "GET", "https://pubsub.mtls.googleapis.com/test" - ) - - assert resp == mock_resp - mock_creds.refresh.assert_called_once() - mock_auth_req.assert_called_once() - - await session.close() - - @pytest.mark.asyncio - async def test_cert_rotation_max_retries_exceeded(self): - """Covers the `if _auth_retry_count < 2:` max retry limit.""" - mock_creds = mock.AsyncMock(spec=credentials.Credentials) - mock_creds.before_request = mock.AsyncMock(return_value=None) - mock_creds.refresh = mock.AsyncMock(return_value=None) - - mock_resp = mock.Mock() - mock_resp.status_code = http_client.UNAUTHORIZED - mock_resp.close = mock.AsyncMock() - mock_auth_req = mock.AsyncMock(return_value=mock_resp) - - session = sessions.AsyncAuthorizedSession( - mock_creds, auth_request=mock_auth_req - ) - session._is_mtls = True - session._cached_cert = b"old_cert" - - with ( - mock.patch( - "google.auth.aio.transport.mtls.check_parameters_for_unauthorized_response", - new_callable=mock.AsyncMock, - ) as mock_check, - mock.patch.object( - session, "configure_mtls_channel", new_callable=mock.AsyncMock - ), + @pytest.mark.asyncio + async def test_configure_mtls_channel(self): + """Tests that the mTLS channel configures correctly when a valid workload config is mocked.""" + with ( + mock.patch.dict(os.environ, {"GOOGLE_API_USE_CLIENT_CERTIFICATE": "true"}), + mock.patch("os.path.exists") as mock_exists, + mock.patch( + "builtins.open", + mock.mock_open(read_data=json.dumps(VALID_WORKLOAD_CONFIG)), + ), + mock.patch( + "google.auth.aio.transport.mtls.get_client_cert_and_key" + ) as mock_helper, + mock.patch( + "google.auth.aio.transport.mtls.make_client_cert_ssl_context" + ) as mock_make_context, + mock.patch("aiohttp.TCPConnector") as mock_connector, + mock.patch("aiohttp.ClientSession") as mock_session, + ): + mock_session.return_value.close = mock.AsyncMock() + mock_exists.return_value = True + mock_helper.return_value = (True, b"fake_cert_data", b"fake_key_data") + + mock_context = mock.Mock(spec=ssl.SSLContext) + mock_make_context.return_value = mock_context + + mock_creds = mock.AsyncMock(spec=credentials.Credentials) + session = sessions.AsyncAuthorizedSession(mock_creds) + + await session.configure_mtls_channel() + + assert session._is_mtls is True + mock_make_context.assert_called_once_with( + b"fake_cert_data", b"fake_key_data" + ) + mock_connector.assert_called_once_with(ssl=mock_context) + mock_session.assert_called_once_with(connector=mock_connector.return_value) + await session.close() + + @pytest.mark.asyncio + async def test_configure_mtls_channel_disabled(self): + """Tests behavior when the config file does not exist.""" + with ( + mock.patch.dict(os.environ, {"GOOGLE_API_USE_CLIENT_CERTIFICATE": "true"}), + mock.patch("os.path.exists") as mock_exists, + ): + mock_exists.return_value = False + mock_creds = mock.AsyncMock(spec=credentials.Credentials) + session = sessions.AsyncAuthorizedSession(mock_creds) + await session.configure_mtls_channel() + assert session._is_mtls is False + await session.close() + + @pytest.mark.asyncio + async def test_configure_mtls_channel_invalid_format(self): + """Verifies that the MutualTLSChannelError is raised for bad formats.""" + with ( + mock.patch.dict(os.environ, {"GOOGLE_API_USE_CLIENT_CERTIFICATE": "true"}), + mock.patch("os.path.exists") as mock_exists, + mock.patch( + "builtins.open", mock.mock_open(read_data='{"invalid": "format"}') + ), + ): + mock_exists.return_value = True + mock_creds = mock.AsyncMock(spec=credentials.Credentials) + session = sessions.AsyncAuthorizedSession(mock_creds) + + with pytest.raises(exceptions.MutualTLSChannelError): + await session.configure_mtls_channel() + await session.close() + + @pytest.mark.asyncio + async def test_configure_mtls_channel_invalid_fields(self): + """If cert is missing expected keys, it should fail gracefully.""" + with ( + mock.patch.dict(os.environ, {"GOOGLE_API_USE_CLIENT_CERTIFICATE": "true"}), + mock.patch("os.path.exists") as mock_exists, + mock.patch( + "builtins.open", mock.mock_open(read_data='{"cert_configs": {}}') + ), + ): + mock_exists.return_value = True + mock_creds = mock.AsyncMock(spec=credentials.Credentials) + session = sessions.AsyncAuthorizedSession(mock_creds) + await session.configure_mtls_channel() + assert session._is_mtls is False + await session.close() + + @pytest.mark.asyncio + async def test_configure_mtls_channel_mock_callback(self): + """Tests mTLS configuration using bytes-returning callback.""" + + def mock_callback(): + return (b"fake_cert_bytes", b"fake_key_bytes") + + with ( + mock.patch.dict(os.environ, {"GOOGLE_API_USE_CLIENT_CERTIFICATE": "true"}), + mock.patch( + "google.auth.transport.mtls.has_default_client_cert_source", + return_value=True, + ), + mock.patch( + "google.auth.aio.transport.mtls.make_client_cert_ssl_context" + ) as mock_make_context, + mock.patch("aiohttp.TCPConnector") as mock_connector, + mock.patch("aiohttp.ClientSession") as mock_session, + ): + mock_session.return_value.close = mock.AsyncMock() + mock_context = mock.Mock(spec=ssl.SSLContext) + mock_make_context.return_value = mock_context + + mock_creds = mock.AsyncMock(spec=credentials.Credentials) + session = sessions.AsyncAuthorizedSession(mock_creds) + + await session.configure_mtls_channel(client_cert_callback=mock_callback) + + assert session._is_mtls is True + mock_make_context.assert_called_once_with( + b"fake_cert_bytes", b"fake_key_bytes" + ) + mock_connector.assert_called_once_with(ssl=mock_context) + mock_session.assert_called_once_with(connector=mock_connector.return_value) + await session.close() + + @pytest.mark.asyncio + async def test_configure_mtls_channel_custom_request(self): + """Tests that if _auth_request is not an AiohttpRequest, _is_mtls is set to False.""" + with ( + mock.patch.dict(os.environ, {"GOOGLE_API_USE_CLIENT_CERTIFICATE": "true"}), + mock.patch("os.path.exists") as mock_exists, + mock.patch( + "builtins.open", + mock.mock_open(read_data=json.dumps(VALID_WORKLOAD_CONFIG)), + ), + mock.patch( + "google.auth.aio.transport.mtls.get_client_cert_and_key" + ) as mock_helper, + mock.patch( + "google.auth.aio.transport.mtls.make_client_cert_ssl_context" + ) as mock_make_context, + ): + mock_exists.return_value = True + mock_helper.return_value = (True, b"fake_cert_data", b"fake_key_data") + + mock_context = mock.Mock(spec=ssl.SSLContext) + mock_make_context.return_value = mock_context + + mock_creds = mock.AsyncMock(spec=credentials.Credentials) + mock_auth_request = mock.AsyncMock(spec=transport.Request) + session = sessions.AsyncAuthorizedSession( + mock_creds, auth_request=mock_auth_request + ) + + with pytest.warns(UserWarning, match="Attempted to establish mTLS"): + await session.configure_mtls_channel() + + assert session._is_mtls is False + mock_make_context.assert_not_called() + await session.close() + + @pytest.mark.asyncio + async def test_configure_mtls_channel_exception_resets_flag(self): + """Tests that self._is_mtls is reset to False if an exception is raised during configuration.""" + with ( + mock.patch.dict(os.environ, {"GOOGLE_API_USE_CLIENT_CERTIFICATE": "true"}), + mock.patch("os.path.exists") as mock_exists, + mock.patch( + "builtins.open", + mock.mock_open(read_data=json.dumps(VALID_WORKLOAD_CONFIG)), + ), + mock.patch( + "google.auth.aio.transport.mtls.get_client_cert_and_key" + ) as mock_helper, + mock.patch( + "google.auth.aio.transport.mtls.make_client_cert_ssl_context" + ) as mock_make_context, + ): + mock_exists.return_value = True + mock_helper.return_value = (True, b"fake_cert_data", b"fake_key_data") + mock_make_context.side_effect = exceptions.ClientCertError("Mock error") + + mock_creds = mock.AsyncMock(spec=credentials.Credentials) + session = sessions.AsyncAuthorizedSession(mock_creds) + + with pytest.raises(exceptions.MutualTLSChannelError): + await session.configure_mtls_channel() + + assert session._is_mtls is False + await session.close() + + @pytest.mark.asyncio + async def test_configure_mtls_channel_transport_error_resets_flag(self): + """Tests that self._is_mtls is reset to False if a TransportError is raised.""" + with ( + mock.patch.dict(os.environ, {"GOOGLE_API_USE_CLIENT_CERTIFICATE": "true"}), + mock.patch("os.path.exists") as mock_exists, + mock.patch( + "builtins.open", + mock.mock_open(read_data=json.dumps(VALID_WORKLOAD_CONFIG)), + ), + mock.patch( + "google.auth.aio.transport.mtls.get_client_cert_and_key" + ) as mock_helper, + mock.patch( + "google.auth.aio.transport.mtls.make_client_cert_ssl_context" + ) as mock_make_context, + ): + mock_exists.return_value = True + mock_helper.return_value = (True, b"fake_cert_data", b"fake_key_data") + mock_make_context.side_effect = exceptions.TransportError("Mock error") + + mock_creds = mock.AsyncMock(spec=credentials.Credentials) + session = sessions.AsyncAuthorizedSession(mock_creds) + + with pytest.raises(exceptions.MutualTLSChannelError): + await session.configure_mtls_channel() + + assert session._is_mtls is False + await session.close() + + @pytest.mark.asyncio + async def test_configure_mtls_channel_atomic_on_exception(self): + """Tests that if configure_mtls_channel already succeeded, a subsequent failure preserves state.""" + # Step 1: Successful configuration + with ( + mock.patch.dict(os.environ, {"GOOGLE_API_USE_CLIENT_CERTIFICATE": "true"}), + mock.patch("os.path.exists") as mock_exists, + mock.patch( + "builtins.open", + mock.mock_open(read_data=json.dumps(VALID_WORKLOAD_CONFIG)), + ), + mock.patch( + "google.auth.aio.transport.mtls.get_client_cert_and_key" + ) as mock_helper, + mock.patch( + "google.auth.aio.transport.mtls.make_client_cert_ssl_context" + ) as mock_make_context, + mock.patch("aiohttp.TCPConnector"), + mock.patch("aiohttp.ClientSession") as mock_session, + ): + mock_session.return_value.close = mock.AsyncMock() + mock_exists.return_value = True + mock_helper.return_value = ( + True, + b"fake_cert_data_1", + b"fake_key_data_1", + ) + + mock_context = mock.Mock(spec=ssl.SSLContext) + mock_make_context.return_value = mock_context + + mock_creds = mock.AsyncMock(spec=credentials.Credentials) + session = sessions.AsyncAuthorizedSession(mock_creds) + + await session.configure_mtls_channel() + assert session._is_mtls is True + assert session._cached_cert == b"fake_cert_data_1" + first_auth_request = session._auth_request + + # Step 2: Failed subsequent configuration attempt + session._mtls_init_task = None + mock_make_context.side_effect = exceptions.ClientCertError("Mock error") + + with pytest.raises(exceptions.MutualTLSChannelError): + await session.configure_mtls_channel() + + assert session._is_mtls is True + assert session._cached_cert == b"fake_cert_data_1" + assert session._auth_request is first_auth_request + await session.close() + + @pytest.mark.asyncio + async def test_configure_mtls_channel_close_exception_does_not_abort(self): + """Tests that an exception in old_auth_request.close() does not abort configuration.""" + with ( + mock.patch.dict(os.environ, {"GOOGLE_API_USE_CLIENT_CERTIFICATE": "true"}), + mock.patch("os.path.exists") as mock_exists, + mock.patch( + "builtins.open", + mock.mock_open(read_data=json.dumps(VALID_WORKLOAD_CONFIG)), + ), + mock.patch( + "google.auth.aio.transport.mtls.get_client_cert_and_key" + ) as mock_helper, + mock.patch( + "google.auth.aio.transport.mtls.make_client_cert_ssl_context" + ) as mock_make_context, + mock.patch("aiohttp.TCPConnector"), + mock.patch("aiohttp.ClientSession") as mock_session, + ): + mock_session.return_value.close = mock.AsyncMock() + mock_exists.return_value = True + mock_helper.return_value = (True, b"fake_cert_data", b"fake_key_data") + + mock_context = mock.Mock(spec=ssl.SSLContext) + mock_make_context.return_value = mock_context + + mock_creds = mock.AsyncMock(spec=credentials.Credentials) + session = sessions.AsyncAuthorizedSession(mock_creds) + + session._auth_request.close = mock.AsyncMock( + side_effect=Exception("Mock close error") + ) + + await session.configure_mtls_channel() + + assert session._is_mtls is True + assert session._cached_cert == b"fake_cert_data" + await session.close() + + @pytest.mark.asyncio + async def test_cert_rotation_failure_raises_error(self): + mock_creds = mock.AsyncMock(spec=credentials.Credentials) + mock_creds.before_request = mock.AsyncMock(return_value=None) + + mock_resp = mock.Mock() + mock_resp.status_code = http_client.UNAUTHORIZED + mock_auth_req = mock.AsyncMock(return_value=mock_resp) + + session = sessions.AsyncAuthorizedSession( + mock_creds, auth_request=mock_auth_req + ) + session._is_mtls = True + session._cached_cert = b"old_cert" + + new_cert = b"new_cert" + new_key = b"new_key" + + with ( + mock.patch( + "google.auth.aio.transport.mtls.check_parameters_for_unauthorized_response", + new_callable=mock.AsyncMock, + ) as mock_check, + mock.patch.object( + session, "configure_mtls_channel", new_callable=mock.AsyncMock + ) as mock_conf, + ): + mock_check.return_value = (new_cert, new_key, b"old_fp", b"new_fp") + mock_conf.side_effect = Exception("Failed to reconfigure") + + with pytest.raises(exceptions.MutualTLSChannelError): + await session.request("GET", "https://pubsub.mtls.googleapis.com/test") + + mock_check.assert_called_once() + mock_conf.assert_called_once() + + await session.close() + + @pytest.mark.asyncio + async def test_cert_rotation_check_params_fails(self): + mock_creds = mock.AsyncMock(spec=credentials.Credentials) + mock_creds.before_request = mock.AsyncMock(return_value=None) + mock_creds.refresh = mock.AsyncMock() + + mock_resp = mock.Mock() + mock_resp.status_code = http_client.UNAUTHORIZED + mock_auth_req = mock.AsyncMock(return_value=mock_resp) + + session = sessions.AsyncAuthorizedSession( + mock_creds, auth_request=mock_auth_req + ) + session._is_mtls = True + session._cached_cert = b"old_cert" + + with ( + mock.patch( + "google.auth.aio.transport.mtls.check_parameters_for_unauthorized_response", + new_callable=mock.AsyncMock, + ) as mock_check, + mock.patch.object( + session, "configure_mtls_channel", new_callable=mock.AsyncMock + ) as mock_conf, + ): + mock_check.side_effect = exceptions.MutualTLSChannelError( + "Failed to check params" + ) + + resp = await session.request( + "GET", "https://pubsub.mtls.googleapis.com/test" + ) + + assert resp == mock_resp + mock_check.assert_called_once() + mock_creds.refresh.assert_not_called() + mock_conf.assert_not_called() + + await session.close() + + @pytest.mark.asyncio + async def test_no_cert_rotation_when_cert_matches_and_mtls_enabled(self): + mock_creds = mock.AsyncMock(spec=credentials.Credentials) + mock_creds.before_request = mock.AsyncMock(return_value=None) + mock_creds.refresh = mock.AsyncMock(return_value=None) + + mock_resp_401 = mock.Mock() + mock_resp_401.status_code = http_client.UNAUTHORIZED + mock_resp_401.close = mock.AsyncMock() + + mock_resp_200 = mock.Mock() + mock_resp_200.status_code = http_client.OK + + # 401 on initial request, 200 on retry after refresh + mock_auth_req = mock.AsyncMock(side_effect=[mock_resp_401, mock_resp_200]) + + session = sessions.AsyncAuthorizedSession( + mock_creds, auth_request=mock_auth_req + ) + session._is_mtls = True + session._cached_cert = b"old_cert" + + new_cert = b"new_cert" + new_key = b"new_key" + + with ( + mock.patch( + "google.auth.aio.transport.mtls.check_parameters_for_unauthorized_response", + new_callable=mock.AsyncMock, + ) as mock_check, + mock.patch.object( + session, "configure_mtls_channel", new_callable=mock.AsyncMock + ) as mock_conf, + ): + # Matching fingerprints mean no mTLS rotation is needed + mock_check.return_value = (new_cert, new_key, b"old_fp", b"old_fp") + + resp = await session.request( + "GET", "https://pubsub.mtls.googleapis.com/test" + ) + + assert resp == mock_resp_200 + mock_check.assert_called_once() + mock_conf.assert_not_called() + mock_creds.refresh.assert_called_once() + assert mock_auth_req.call_count == 2 + mock_resp_401.close.assert_awaited_once() + + await session.close() + + @pytest.mark.asyncio + async def test_cert_rotation_success_and_retry(self): + mock_creds = mock.AsyncMock(spec=credentials.Credentials) + mock_creds.before_request = mock.AsyncMock(return_value=None) + mock_creds.refresh = mock.AsyncMock(return_value=None) + + mock_resp_401 = mock.Mock() + mock_resp_401.status_code = http_client.UNAUTHORIZED + mock_resp_200 = mock.Mock() + mock_resp_200.status_code = http_client.OK + + mock_auth_req = mock.AsyncMock(side_effect=[mock_resp_401, mock_resp_200]) + + session = sessions.AsyncAuthorizedSession( + mock_creds, auth_request=mock_auth_req + ) + session._is_mtls = True + session._cached_cert = b"old_cert" + + new_cert = b"new_cert" + new_key = b"new_key" + + with ( + mock.patch( + "google.auth.aio.transport.mtls.check_parameters_for_unauthorized_response", + new_callable=mock.AsyncMock, + ) as mock_check, + mock.patch.object( + session, "configure_mtls_channel", new_callable=mock.AsyncMock + ) as mock_conf, + ): + mock_check.return_value = (new_cert, new_key, b"old_fp", b"new_fp") + + resp = await session.request( + "GET", "https://pubsub.mtls.googleapis.com/test" + ) + + assert resp == mock_resp_200 + mock_check.assert_called_once() + mock_conf.assert_called_once_with(mock.ANY) + mock_creds.refresh.assert_called_once() + assert mock_creds.before_request.call_count == 2 + + await session.close() + + @pytest.mark.asyncio + async def test_cert_rotation_lock_contention(self): + mock_creds = mock.AsyncMock(spec=credentials.Credentials) + mock_creds.before_request = mock.AsyncMock(return_value=None) + mock_creds.refresh = mock.AsyncMock(return_value=None) + + mock_resp_401 = mock.Mock() + mock_resp_401.status_code = http_client.UNAUTHORIZED + mock_auth_req = mock.AsyncMock(return_value=mock_resp_401) + + session = sessions.AsyncAuthorizedSession( + mock_creds, auth_request=mock_auth_req + ) + session._is_mtls = True + session._cached_cert = b"old_cert" + + new_cert = b"new_cert" + new_key = b"new_key" + + async def mock_configure_mtls_channel(*args, **kwargs): + await asyncio.sleep(0.01) + session._cached_cert = new_cert + + async def mock_check_side_effect(cached_cert, callback=None): + if cached_cert == b"old_cert": + return (new_cert, new_key, b"old_fp", b"new_fp") + return (new_cert, new_key, b"new_fp", b"new_fp") + + with ( + mock.patch( + "google.auth.aio.transport.mtls.check_parameters_for_unauthorized_response", + new_callable=mock.AsyncMock, + ) as mock_check, + mock.patch.object( + session, "configure_mtls_channel", new_callable=mock.AsyncMock + ) as mock_conf, + ): + mock_check.side_effect = mock_check_side_effect + mock_conf.side_effect = mock_configure_mtls_channel + + tasks = [ + session.request("GET", "https://pubsub.mtls.googleapis.com/test") + for _ in range(3) + ] + responses = await asyncio.gather(*tasks) + + for resp in responses: + assert resp == mock_resp_401 + + mock_conf.assert_called_once() + + await session.close() + + @pytest.mark.asyncio + async def test_cert_rotation_lock_contention_no_cert_change(self): + mock_creds = mock.AsyncMock(spec=credentials.Credentials) + mock_creds.before_request = mock.AsyncMock(return_value=None) + mock_creds.refresh = mock.AsyncMock(return_value=None) + + mock_resp_401 = mock.Mock() + mock_resp_401.status_code = http_client.UNAUTHORIZED + mock_resp_200 = mock.Mock() + mock_resp_200.status_code = http_client.OK + + mock_auth_req = mock.AsyncMock( + side_effect=[mock_resp_401] * 3 + [mock_resp_200] * 3 + ) + + session = sessions.AsyncAuthorizedSession( + mock_creds, auth_request=mock_auth_req + ) + session._is_mtls = True + session._cached_cert = b"old_cert" + + async def mock_check_side_effect(cached_cert, callback=None): + await asyncio.sleep(0.01) + return (b"old_cert", b"old_key", b"old_fp", b"old_fp") + + with ( + mock.patch( + "google.auth.aio.transport.mtls.check_parameters_for_unauthorized_response", + new_callable=mock.AsyncMock, + ) as mock_check, + mock.patch.object( + session, "configure_mtls_channel", new_callable=mock.AsyncMock + ) as mock_conf, + ): + mock_check.side_effect = mock_check_side_effect + + tasks = [ + session.request("GET", "https://pubsub.mtls.googleapis.com/test") + for _ in range(3) + ] + responses = await asyncio.gather(*tasks) + + for resp in responses: + assert resp == mock_resp_200 + + mock_check.assert_called_once() + mock_conf.assert_not_called() + # Concurrent 401s properly deduplicate to 1 refresh + assert mock_creds.refresh.call_count == 1 + + await session.close() + + @pytest.mark.asyncio + async def test_psc_endpoint_triggers_cert_rotation(self): + """Verifies that PSC endpoints (*.p.googleapis.com) are recognized as mTLS endpoints.""" + mock_creds = mock.AsyncMock(spec=credentials.Credentials) + mock_creds.before_request = mock.AsyncMock(return_value=None) + mock_creds.refresh = mock.AsyncMock(return_value=None) + + mock_resp_401 = mock.Mock() + mock_resp_401.status_code = http_client.UNAUTHORIZED + mock_resp_200 = mock.Mock() + mock_resp_200.status_code = http_client.OK + + mock_auth_req = mock.AsyncMock(side_effect=[mock_resp_401, mock_resp_200]) + + session = sessions.AsyncAuthorizedSession( + mock_creds, auth_request=mock_auth_req + ) + session._is_mtls = True + session._cached_cert = b"old_cert" + + new_cert = b"new_cert" + new_key = b"new_key" + + with ( + mock.patch( + "google.auth.aio.transport.mtls.check_parameters_for_unauthorized_response", + new_callable=mock.AsyncMock, + ) as mock_check, + mock.patch.object( + session, "configure_mtls_channel", new_callable=mock.AsyncMock + ) as mock_conf, + ): + mock_check.return_value = (new_cert, new_key, b"old_fp", b"new_fp") + + resp = await session.request("GET", "https://pubsub.p.googleapis.com/test") + + assert resp == mock_resp_200 + mock_check.assert_called_once() + mock_conf.assert_called_once_with(mock.ANY) + + await session.close() + + @pytest.mark.asyncio + async def test_non_mtls_url_bypasses_rotation(self): + mock_creds = mock.AsyncMock(spec=credentials.Credentials) + mock_creds.before_request = mock.AsyncMock(return_value=None) + mock_creds.refresh = mock.AsyncMock(return_value=None) + + mock_resp_401 = mock.Mock() + mock_resp_401.status_code = http_client.UNAUTHORIZED + mock_auth_req = mock.AsyncMock(return_value=mock_resp_401) + + session = sessions.AsyncAuthorizedSession( + mock_creds, auth_request=mock_auth_req + ) + + session._is_mtls = True + session._cached_cert = b"old_cert" + + with ( + mock.patch( + "google.auth.aio.transport.mtls.check_parameters_for_unauthorized_response", + new_callable=mock.AsyncMock, + ) as mock_check, + mock.patch.object( + session, "configure_mtls_channel", new_callable=mock.AsyncMock + ) as mock_conf, + ): + resp = await session.request("GET", "https://example.com/test") + + assert resp == mock_resp_401 + mock_check.assert_not_called() + mock_conf.assert_not_called() + assert mock_creds.refresh.call_count == 2 + assert mock_auth_req.call_count == 3 + + await session.close() + + @pytest.mark.asyncio + async def test_cert_rotation_skips_retry_for_streaming(self): + mock_creds = mock.AsyncMock(spec=credentials.Credentials) + mock_creds.before_request = mock.AsyncMock(return_value=None) + mock_creds.refresh = mock.AsyncMock() + + mock_resp = mock.Mock() + mock_resp.status_code = http_client.UNAUTHORIZED + mock_auth_req = mock.AsyncMock(return_value=mock_resp) + + session = sessions.AsyncAuthorizedSession( + mock_creds, auth_request=mock_auth_req + ) + session._is_mtls = True + session._cached_cert = b"old_cert" + + class MockStream: + def read(self): + pass + + with ( + mock.patch( + "google.auth.aio.transport.mtls.check_parameters_for_unauthorized_response", + new_callable=mock.AsyncMock, + ) as mock_check, + mock.patch.object( + session, "configure_mtls_channel", new_callable=mock.AsyncMock + ) as mock_conf, + ): + mock_check.return_value = (b"new", b"new", b"old_fp", b"new_fp") + + resp = await session.request( + "GET", "https://pubsub.mtls.googleapis.com/test", data=MockStream() + ) + + assert resp == mock_resp + mock_conf.assert_called_once() + + mock_creds.refresh.assert_called_once() + await session.close() + + @pytest.mark.asyncio + async def test_cert_rotation_credential_refresh_fails(self): + """Covers the except block for RefreshError when credentials fail to refresh.""" + mock_creds = mock.AsyncMock(spec=credentials.Credentials) + mock_creds.before_request = mock.AsyncMock(return_value=None) + mock_creds.refresh = mock.AsyncMock( + side_effect=exceptions.RefreshError("Refresh failed") + ) + + mock_resp = mock.Mock() + mock_resp.status_code = http_client.UNAUTHORIZED + mock_auth_req = mock.AsyncMock(return_value=mock_resp) + + session = sessions.AsyncAuthorizedSession( + mock_creds, auth_request=mock_auth_req + ) + session._is_mtls = True + session._cached_cert = b"old_cert" + + with ( + mock.patch( + "google.auth.aio.transport.mtls.check_parameters_for_unauthorized_response", + new_callable=mock.AsyncMock, + ) as mock_check, + mock.patch.object( + session, "configure_mtls_channel", new_callable=mock.AsyncMock + ), + ): + mock_check.return_value = (b"new", b"new", b"old_fp", b"new_fp") + + resp = await session.request( + "GET", "https://pubsub.mtls.googleapis.com/test" + ) + + assert resp == mock_resp + mock_creds.refresh.assert_called_once() + mock_auth_req.assert_called_once() + + await session.close() + + @pytest.mark.asyncio + async def test_cert_rotation_max_retries_exceeded(self): + """Covers the `if _auth_retry_count < 2:` max retry limit.""" + mock_creds = mock.AsyncMock(spec=credentials.Credentials) + mock_creds.before_request = mock.AsyncMock(return_value=None) + mock_creds.refresh = mock.AsyncMock(return_value=None) + + mock_resp = mock.Mock() + mock_resp.status_code = http_client.UNAUTHORIZED + mock_resp.close = mock.AsyncMock() + mock_auth_req = mock.AsyncMock(return_value=mock_resp) + + session = sessions.AsyncAuthorizedSession( + mock_creds, auth_request=mock_auth_req + ) + session._is_mtls = True + session._cached_cert = b"old_cert" + + with ( + mock.patch( + "google.auth.aio.transport.mtls.check_parameters_for_unauthorized_response", + new_callable=mock.AsyncMock, + ) as mock_check, + mock.patch.object( + session, "configure_mtls_channel", new_callable=mock.AsyncMock + ), + ): + mock_check.return_value = (b"new", b"new", b"old_fp", b"new_fp") + + resp = await session.request( + "GET", "https://pubsub.mtls.googleapis.com/test" + ) + + assert resp == mock_resp + assert mock_auth_req.call_count == 3 + assert mock_check.call_count == 2 + assert mock_resp.close.call_count == 2 + + await session.close() + + @pytest.mark.asyncio + async def test_session_close_cleans_old_auth_requests(self): + """Covers the loop in the `close()` method that drains `_old_auth_requests`.""" + mock_creds = mock.AsyncMock(spec=credentials.Credentials) + session = sessions.AsyncAuthorizedSession( + mock_creds, auth_request=mock.AsyncMock() + ) + + mock_old_req_1 = mock.AsyncMock() + mock_old_req_2 = mock.AsyncMock() + mock_old_req_3_fails = mock.AsyncMock() + mock_old_req_3_fails.close.side_effect = Exception("Close error") + + session._old_auth_requests.extend( + [mock_old_req_1, mock_old_req_2, mock_old_req_3_fails] + ) + + await session.close() + + mock_old_req_1.close.assert_called_once() + mock_old_req_2.close.assert_called_once() + mock_old_req_3_fails.close.assert_called_once() + assert len(session._old_auth_requests) == 0 + + @pytest.mark.asyncio + async def test_request_401_streaming_refreshes_creds_and_returns_open_response( + self, ): - mock_check.return_value = (b"new", b"new", b"old_fp", b"new_fp") - - resp = await session.request( - "GET", "https://pubsub.mtls.googleapis.com/test" - ) - - assert resp == mock_resp - assert mock_auth_req.call_count == 3 - assert mock_check.call_count == 2 - assert mock_resp.close.call_count == 2 - - await session.close() - - @pytest.mark.asyncio - async def test_session_close_cleans_old_auth_requests(self): - """Covers the loop in the `close()` method that drains `_old_auth_requests`.""" - mock_creds = mock.AsyncMock(spec=credentials.Credentials) - session = sessions.AsyncAuthorizedSession( - mock_creds, auth_request=mock.AsyncMock() - ) - - mock_old_req_1 = mock.AsyncMock() - mock_old_req_2 = mock.AsyncMock() - mock_old_req_3_fails = mock.AsyncMock() - mock_old_req_3_fails.close.side_effect = Exception("Close error") - - session._old_auth_requests.extend( - [mock_old_req_1, mock_old_req_2, mock_old_req_3_fails] - ) - - await session.close() - - mock_old_req_1.close.assert_called_once() - mock_old_req_2.close.assert_called_once() - mock_old_req_3_fails.close.assert_called_once() - assert len(session._old_auth_requests) == 0 - - @pytest.mark.asyncio - async def test_request_401_streaming_refreshes_creds_and_returns_open_response( - self, - ): - """Verifies that streaming requests refresh credentials but return the unclosed response.""" - mock_creds = mock.AsyncMock(spec=credentials.Credentials) - mock_creds.before_request = mock.AsyncMock(return_value=None) - mock_creds.refresh = mock.AsyncMock(return_value=None) - - mock_resp_401 = mock.Mock() - mock_resp_401.status_code = http_client.UNAUTHORIZED - mock_resp_401.close = mock.AsyncMock() - - mock_auth_req = mock.AsyncMock(return_value=mock_resp_401) - session = sessions.AsyncAuthorizedSession( - mock_creds, auth_request=mock_auth_req - ) - - streaming_data = (chunk for chunk in [b"chunk1", b"chunk2"]) - response = await session.request( - "POST", "https://example.com", data=streaming_data - ) - - assert response == mock_resp_401 - mock_creds.refresh.assert_awaited_once() - mock_resp_401.close.assert_not_called() - await session.close() - - @pytest.mark.asyncio - async def test_request_401_closes_response_on_timeout_during_recovery(self): - """Verifies that response is closed when auth_with_timeout times out during _recover_auth_state.""" - mock_creds = mock.AsyncMock(spec=credentials.Credentials) - mock_creds.before_request = mock.AsyncMock(return_value=None) - - async def slow_refresh(*args, **kwargs): - await asyncio.sleep(10) - - mock_creds.refresh = mock.AsyncMock(side_effect=slow_refresh) - - mock_resp_401 = mock.Mock() - mock_resp_401.status_code = http_client.UNAUTHORIZED - mock_resp_401.close = mock.AsyncMock() - - mock_auth_req = mock.AsyncMock(return_value=mock_resp_401) - session = sessions.AsyncAuthorizedSession( - mock_creds, auth_request=mock_auth_req - ) - - with pytest.raises(TimeoutError): - await session.request("GET", "https://example.com", max_allowed_time=0.01) - - mock_resp_401.close.assert_awaited_once() - await session.close() - - @pytest.mark.asyncio - async def test_request_401_closes_response_on_cancellation(self): - """Verifies that response is closed and CancelledError propagated if task is cancelled.""" - mock_creds = mock.AsyncMock(spec=credentials.Credentials) - mock_creds.before_request = mock.AsyncMock(return_value=None) - - refresh_started = asyncio.Event() - - async def cancel_on_refresh(*args, **kwargs): - refresh_started.set() - await asyncio.sleep(10) - - mock_creds.refresh = mock.AsyncMock(side_effect=cancel_on_refresh) - - mock_resp_401 = mock.Mock() - mock_resp_401.status_code = http_client.UNAUTHORIZED - mock_resp_401.close = mock.AsyncMock() - - mock_auth_req = mock.AsyncMock(return_value=mock_resp_401) - session = sessions.AsyncAuthorizedSession( - mock_creds, auth_request=mock_auth_req - ) - - task = asyncio.create_task(session.request("GET", "https://example.com")) - await refresh_started.wait() - task.cancel() - - with pytest.raises(asyncio.CancelledError): - await task - - mock_resp_401.close.assert_awaited_once() - await session.close() - - @pytest.mark.asyncio - async def test_request_401_concurrent_refreshes_are_deduplicated(self): - """Verifies that concurrent 401s execute only one credentials.refresh call.""" - mock_creds = mock.AsyncMock(spec=credentials.Credentials) - mock_creds.before_request = mock.AsyncMock(return_value=None) - - mock_resp_401 = mock.Mock() - mock_resp_401.status_code = http_client.UNAUTHORIZED - mock_resp_401.close = mock.AsyncMock() - - mock_resp_200 = mock.Mock() - mock_resp_200.status_code = http_client.OK - mock_resp_200.close = mock.AsyncMock() - - mock_auth_req = mock.AsyncMock( - side_effect=[ - mock_resp_401, - mock_resp_401, - mock_resp_200, - mock_resp_200, - ] - ) - session = sessions.AsyncAuthorizedSession( - mock_creds, auth_request=mock_auth_req - ) - - refresh_count = 0 - - async def slow_refresh(*args, **kwargs): - nonlocal refresh_count - refresh_count += 1 - await asyncio.sleep(0.05) - - mock_creds.refresh = mock.AsyncMock(side_effect=slow_refresh) - - results = await asyncio.gather( - session.request("GET", "https://example.com/1"), - session.request("GET", "https://example.com/2"), - ) - - assert all(r.status_code == 200 for r in results) - assert refresh_count == 1 - await session.close() + """Verifies that streaming requests refresh credentials but return the unclosed response.""" + mock_creds = mock.AsyncMock(spec=credentials.Credentials) + mock_creds.before_request = mock.AsyncMock(return_value=None) + mock_creds.refresh = mock.AsyncMock(return_value=None) + + mock_resp_401 = mock.Mock() + mock_resp_401.status_code = http_client.UNAUTHORIZED + mock_resp_401.close = mock.AsyncMock() + + mock_auth_req = mock.AsyncMock(return_value=mock_resp_401) + session = sessions.AsyncAuthorizedSession( + mock_creds, auth_request=mock_auth_req + ) + + streaming_data = (chunk for chunk in [b"chunk1", b"chunk2"]) + response = await session.request( + "POST", "https://example.com", data=streaming_data + ) + + assert response == mock_resp_401 + mock_creds.refresh.assert_awaited_once() + mock_resp_401.close.assert_not_called() + await session.close() + + @pytest.mark.asyncio + async def test_request_401_closes_response_on_timeout_during_recovery(self): + """Verifies that response is closed when auth_with_timeout times out during _recover_auth_state.""" + mock_creds = mock.AsyncMock(spec=credentials.Credentials) + mock_creds.before_request = mock.AsyncMock(return_value=None) + + async def slow_refresh(*args, **kwargs): + await asyncio.sleep(10) + + mock_creds.refresh = mock.AsyncMock(side_effect=slow_refresh) + + mock_resp_401 = mock.Mock() + mock_resp_401.status_code = http_client.UNAUTHORIZED + mock_resp_401.close = mock.AsyncMock() + + mock_auth_req = mock.AsyncMock(return_value=mock_resp_401) + session = sessions.AsyncAuthorizedSession( + mock_creds, auth_request=mock_auth_req + ) + + with pytest.raises(TimeoutError): + await session.request("GET", "https://example.com", max_allowed_time=0.01) + + mock_resp_401.close.assert_awaited_once() + await session.close() + + @pytest.mark.asyncio + async def test_request_401_closes_response_on_cancellation(self): + """Verifies that response is closed and CancelledError propagated if task is cancelled.""" + mock_creds = mock.AsyncMock(spec=credentials.Credentials) + mock_creds.before_request = mock.AsyncMock(return_value=None) + + refresh_started = asyncio.Event() + + async def cancel_on_refresh(*args, **kwargs): + refresh_started.set() + await asyncio.sleep(10) + + mock_creds.refresh = mock.AsyncMock(side_effect=cancel_on_refresh) + + mock_resp_401 = mock.Mock() + mock_resp_401.status_code = http_client.UNAUTHORIZED + mock_resp_401.close = mock.AsyncMock() + + mock_auth_req = mock.AsyncMock(return_value=mock_resp_401) + session = sessions.AsyncAuthorizedSession( + mock_creds, auth_request=mock_auth_req + ) + + task = asyncio.create_task(session.request("GET", "https://example.com")) + await refresh_started.wait() + task.cancel() + + with pytest.raises(asyncio.CancelledError): + await task + + mock_resp_401.close.assert_awaited_once() + await session.close() + + @pytest.mark.asyncio + async def test_request_401_concurrent_refreshes_are_deduplicated(self): + """Verifies that concurrent 401s execute only one credentials.refresh call.""" + mock_creds = mock.AsyncMock(spec=credentials.Credentials) + mock_creds.before_request = mock.AsyncMock(return_value=None) + + mock_resp_401 = mock.Mock() + mock_resp_401.status_code = http_client.UNAUTHORIZED + mock_resp_401.close = mock.AsyncMock() + + mock_resp_200 = mock.Mock() + mock_resp_200.status_code = http_client.OK + mock_resp_200.close = mock.AsyncMock() + + mock_auth_req = mock.AsyncMock( + side_effect=[ + mock_resp_401, + mock_resp_401, + mock_resp_200, + mock_resp_200, + ] + ) + session = sessions.AsyncAuthorizedSession( + mock_creds, auth_request=mock_auth_req + ) + + refresh_count = 0 + + async def slow_refresh(*args, **kwargs): + nonlocal refresh_count + refresh_count += 1 + await asyncio.sleep(0.05) + + mock_creds.refresh = mock.AsyncMock(side_effect=slow_refresh) + + results = await asyncio.gather( + session.request("GET", "https://example.com/1"), + session.request("GET", "https://example.com/2"), + ) + + assert all(r.status_code == 200 for r in results) + assert refresh_count == 1 + await session.close() From f1b18adb83161f085e922085aae126e52505f3c5 Mon Sep 17 00:00:00 2001 From: Radhika Agrawal Date: Thu, 3 Sep 2026 18:01:08 +0000 Subject: [PATCH 74/79] Fix: Fix the lint errors Signed-off-by: Radhika Agrawal --- .../google/auth/aio/transport/sessions.py | 2 +- .../tests/transport/aio/test_mtls.py | 274 +++++++++--------- .../tests/transport/aio/test_sessions_mtls.py | 2 + 3 files changed, 148 insertions(+), 130 deletions(-) diff --git a/packages/google-auth/google/auth/aio/transport/sessions.py b/packages/google-auth/google/auth/aio/transport/sessions.py index cb7fdf9361b9..ae5bfb0ff3e8 100644 --- a/packages/google-auth/google/auth/aio/transport/sessions.py +++ b/packages/google-auth/google/auth/aio/transport/sessions.py @@ -46,7 +46,7 @@ _MTLS_URL_PREFIXES = [ "mtls.googleapis.com", "mtls.sandbox.googleapis.com", - ".p.googleapis.com", + "p.googleapis.com", ] # Tracks the internal aiohttp installation and usage diff --git a/packages/google-auth/tests/transport/aio/test_mtls.py b/packages/google-auth/tests/transport/aio/test_mtls.py index 20a66380b2bd..5343e1473dd7 100644 --- a/packages/google-auth/tests/transport/aio/test_mtls.py +++ b/packages/google-auth/tests/transport/aio/test_mtls.py @@ -14,15 +14,20 @@ from unittest import mock +import pytest + from google.auth import exceptions from google.auth.aio.transport import mtls -import pytest -CERT_BYTES = b"-----BEGIN CERTIFICATE-----\nMIID...CERT1...=\n-----END CERTIFICATE-----\n" +CERT_BYTES = ( + b"-----BEGIN CERTIFICATE-----\nMIID...CERT1...=\n-----END CERTIFICATE-----\n" +) KEY_BYTES = ( b"-----BEGIN PRIVATE KEY-----\nMIIE...KEY1...==\n-----END PRIVATE KEY-----\n" ) -NEW_CERT_BYTES = b"-----BEGIN CERTIFICATE-----\nMIID...CERT2...=\n-----END CERTIFICATE-----\n" +NEW_CERT_BYTES = ( + b"-----BEGIN CERTIFICATE-----\nMIID...CERT2...=\n-----END CERTIFICATE-----\n" +) NEW_KEY_BYTES = ( b"-----BEGIN PRIVATE KEY-----\nMIIE...KEY2...==\n-----END PRIVATE KEY-----\n" ) @@ -30,172 +35,183 @@ @pytest.mark.asyncio async def test_check_parameters_no_client_cert(): - """Test when no certificate is discovered (has_cert is False).""" - with mock.patch.object( - mtls, "get_client_cert_and_key", return_value=(False, None, None) - ): - cert, key, cached_fp, current_fp = ( - await mtls.check_parameters_for_unauthorized_response( + """Test when no certificate is discovered (has_cert is False).""" + with mock.patch.object( + mtls, "get_client_cert_and_key", return_value=(False, None, None) + ): + ( + cert, + key, + cached_fp, + current_fp, + ) = await mtls.check_parameters_for_unauthorized_response( cached_cert=b"stale_cert", client_cert_callback=None ) - ) - assert cert is None - assert key is None - assert cached_fp is None - assert current_fp is None + assert cert is None + assert key is None + assert cached_fp is None + assert current_fp is None @pytest.mark.asyncio async def test_check_parameters_cert_matched(): - """Test when newly retrieved certificate matches the cached certificate.""" - - def callback(): - return CERT_BYTES, KEY_BYTES - - with ( - mock.patch("google.auth._agent_identity_utils.parse_certificate"), - mock.patch( - "google.auth._agent_identity_utils.calculate_certificate_fingerprint", - return_value="FINGERPRINT_A", - ), - mock.patch( - "google.auth._agent_identity_utils.get_cached_cert_fingerprint", - return_value="FINGERPRINT_A", - ), - ): - cert, key, cached_fp, current_fp = ( - await mtls.check_parameters_for_unauthorized_response( + """Test when newly retrieved certificate matches the cached certificate.""" + + def callback(): + return CERT_BYTES, KEY_BYTES + + with ( + mock.patch("google.auth._agent_identity_utils.parse_certificate"), + mock.patch( + "google.auth._agent_identity_utils.calculate_certificate_fingerprint", + return_value="FINGERPRINT_A", + ), + mock.patch( + "google.auth._agent_identity_utils.get_cached_cert_fingerprint", + return_value="FINGERPRINT_A", + ), + ): + ( + cert, + key, + cached_fp, + current_fp, + ) = await mtls.check_parameters_for_unauthorized_response( cached_cert=CERT_BYTES, client_cert_callback=callback ) - ) - assert cert == CERT_BYTES - assert key == KEY_BYTES - assert cached_fp == "FINGERPRINT_A" - assert current_fp == "FINGERPRINT_A" - assert cached_fp == current_fp + assert cert == CERT_BYTES + assert key == KEY_BYTES + assert cached_fp == "FINGERPRINT_A" + assert current_fp == "FINGERPRINT_A" + assert cached_fp == current_fp @pytest.mark.asyncio async def test_check_parameters_cert_mismatch_rotation(): - """Test when newly retrieved certificate differs from the cached certificate (rotation occurred).""" - - def callback(): - return NEW_CERT_BYTES, NEW_KEY_BYTES - - with ( - mock.patch("google.auth._agent_identity_utils.parse_certificate"), - mock.patch( - "google.auth._agent_identity_utils.calculate_certificate_fingerprint", - return_value="FINGERPRINT_NEW", - ), - mock.patch( - "google.auth._agent_identity_utils.get_cached_cert_fingerprint", - return_value="FINGERPRINT_OLD", - ), - ): - cert, key, cached_fp, current_fp = ( - await mtls.check_parameters_for_unauthorized_response( + """Test when newly retrieved certificate differs from the cached certificate (rotation occurred).""" + + def callback(): + return NEW_CERT_BYTES, NEW_KEY_BYTES + + with ( + mock.patch("google.auth._agent_identity_utils.parse_certificate"), + mock.patch( + "google.auth._agent_identity_utils.calculate_certificate_fingerprint", + return_value="FINGERPRINT_NEW", + ), + mock.patch( + "google.auth._agent_identity_utils.get_cached_cert_fingerprint", + return_value="FINGERPRINT_OLD", + ), + ): + ( + cert, + key, + cached_fp, + current_fp, + ) = await mtls.check_parameters_for_unauthorized_response( cached_cert=CERT_BYTES, client_cert_callback=callback ) - ) - assert cert == NEW_CERT_BYTES - assert key == NEW_KEY_BYTES - assert cached_fp == "FINGERPRINT_OLD" - assert current_fp == "FINGERPRINT_NEW" - assert cached_fp != current_fp + assert cert == NEW_CERT_BYTES + assert key == NEW_KEY_BYTES + assert cached_fp == "FINGERPRINT_OLD" + assert current_fp == "FINGERPRINT_NEW" + assert cached_fp != current_fp @pytest.mark.asyncio async def test_check_parameters_without_cached_cert(): - """Test when cached_cert is None.""" - - def callback(): - return CERT_BYTES, KEY_BYTES - - with ( - mock.patch("google.auth._agent_identity_utils.parse_certificate"), - mock.patch( - "google.auth._agent_identity_utils.calculate_certificate_fingerprint", - return_value="FINGERPRINT_CURRENT", - ), - mock.patch( - "google.auth._agent_identity_utils.get_cached_cert_fingerprint" - ) as mock_get_cached, - ): - cert, key, cached_fp, current_fp = ( - await mtls.check_parameters_for_unauthorized_response( + """Test when cached_cert is None.""" + + def callback(): + return CERT_BYTES, KEY_BYTES + + with ( + mock.patch("google.auth._agent_identity_utils.parse_certificate"), + mock.patch( + "google.auth._agent_identity_utils.calculate_certificate_fingerprint", + return_value="FINGERPRINT_CURRENT", + ), + mock.patch( + "google.auth._agent_identity_utils.get_cached_cert_fingerprint" + ) as mock_get_cached, + ): + ( + cert, + key, + cached_fp, + current_fp, + ) = await mtls.check_parameters_for_unauthorized_response( cached_cert=None, client_cert_callback=callback ) - ) - assert cert == CERT_BYTES - assert key == KEY_BYTES - assert cached_fp == "FINGERPRINT_CURRENT" - assert current_fp == "FINGERPRINT_CURRENT" - mock_get_cached.assert_not_called() + assert cert == CERT_BYTES + assert key == KEY_BYTES + assert cached_fp == "FINGERPRINT_CURRENT" + assert current_fp == "FINGERPRINT_CURRENT" + mock_get_cached.assert_not_called() @pytest.mark.asyncio async def test_check_parameters_executor_fingerprint_computation(): - """Test that fingerprint computation is properly offloaded to the executor.""" - - def callback(): - return CERT_BYTES, KEY_BYTES - - with ( - mock.patch.object( - mtls, "_run_in_executor", wraps=mtls._run_in_executor - ) as mock_run_in_executor, - mock.patch("google.auth._agent_identity_utils.parse_certificate"), - mock.patch( - "google.auth._agent_identity_utils.calculate_certificate_fingerprint", - return_value="FP_CURRENT", - ), - mock.patch( - "google.auth._agent_identity_utils.get_cached_cert_fingerprint", - return_value="FP_CACHED", - ), - ): - cert, key, cached_fp, current_fp = ( - await mtls.check_parameters_for_unauthorized_response( + """Test that fingerprint computation is properly offloaded to the executor.""" + + def callback(): + return CERT_BYTES, KEY_BYTES + + with ( + mock.patch.object( + mtls, "_run_in_executor", wraps=mtls._run_in_executor + ) as mock_run_in_executor, + mock.patch("google.auth._agent_identity_utils.parse_certificate"), + mock.patch( + "google.auth._agent_identity_utils.calculate_certificate_fingerprint", + return_value="FP_CURRENT", + ), + mock.patch( + "google.auth._agent_identity_utils.get_cached_cert_fingerprint", + return_value="FP_CACHED", + ), + ): + ( + cert, + key, + cached_fp, + current_fp, + ) = await mtls.check_parameters_for_unauthorized_response( cached_cert=CERT_BYTES, client_cert_callback=callback ) - ) - assert mock_run_in_executor.called - assert cert == CERT_BYTES - assert cached_fp == "FP_CACHED" - assert current_fp == "FP_CURRENT" + assert mock_run_in_executor.called + assert cert == CERT_BYTES + assert cached_fp == "FP_CACHED" + assert current_fp == "FP_CURRENT" @pytest.mark.asyncio async def test_check_parameters_callback_exception_propagation(): - """Test that exceptions raised by client_cert_callback propagate cleanly.""" + """Test that exceptions raised by client_cert_callback propagate cleanly.""" - def failing_callback(): - raise exceptions.ClientCertError("Client cert provider failed") + def failing_callback(): + raise exceptions.ClientCertError("Client cert provider failed") - with pytest.raises( - exceptions.ClientCertError, match="Client cert provider failed" - ): - await mtls.check_parameters_for_unauthorized_response( - cached_cert=CERT_BYTES, client_cert_callback=failing_callback - ) + with pytest.raises(exceptions.ClientCertError, match="Client cert provider failed"): + await mtls.check_parameters_for_unauthorized_response( + cached_cert=CERT_BYTES, client_cert_callback=failing_callback + ) @pytest.mark.asyncio async def test_check_parameters_async_callback_exception_propagation(): - """Test that exceptions raised in an async client_cert_callback propagate cleanly.""" + """Test that exceptions raised in an async client_cert_callback propagate cleanly.""" - async def failing_async_callback(): - raise OSError("Disk read error while loading certificates") + async def failing_async_callback(): + raise OSError("Disk read error while loading certificates") - with pytest.raises( - OSError, match="Disk read error while loading certificates" - ): - await mtls.check_parameters_for_unauthorized_response( - cached_cert=CERT_BYTES, client_cert_callback=failing_async_callback - ) + with pytest.raises(OSError, match="Disk read error while loading certificates"): + await mtls.check_parameters_for_unauthorized_response( + cached_cert=CERT_BYTES, client_cert_callback=failing_async_callback + ) diff --git a/packages/google-auth/tests/transport/aio/test_sessions_mtls.py b/packages/google-auth/tests/transport/aio/test_sessions_mtls.py index 0e7f3bff1787..f9aafdd8cff8 100644 --- a/packages/google-auth/tests/transport/aio/test_sessions_mtls.py +++ b/packages/google-auth/tests/transport/aio/test_sessions_mtls.py @@ -18,7 +18,9 @@ import os import ssl from unittest import mock + import pytest + from google.auth import exceptions from google.auth.aio import credentials from google.auth.aio import transport From 4cb7e057ab503a7cd95960227c2262502d12f05c Mon Sep 17 00:00:00 2001 From: agrawalradhika-cell Date: Thu, 3 Sep 2026 13:38:09 -0700 Subject: [PATCH 75/79] chore: Modify check_parameters_for_unauthorized_response function Updated the 'check_parameters_for_unauthorized_response' function to include an optional client_cert_callback parameter and added detailed docstring for better understanding. --- .../google/auth/aio/transport/mtls.py | 18 ++++++++++++++++-- 1 file changed, 16 insertions(+), 2 deletions(-) diff --git a/packages/google-auth/google/auth/aio/transport/mtls.py b/packages/google-auth/google/auth/aio/transport/mtls.py index 1d76fb744d3d..45e471ba099c 100644 --- a/packages/google-auth/google/auth/aio/transport/mtls.py +++ b/packages/google-auth/google/auth/aio/transport/mtls.py @@ -179,8 +179,22 @@ async def get_client_cert_and_key(client_cert_callback=None): return has_cert, cert, key -async def check_parameters_for_unauthorized_response(cached_cert, client_cert_callback): - """Async helper to retrieve certs and compute fingerprints for mTLS rotation.""" +async def check_parameters_for_unauthorized_response(cached_cert, client_cert_callback=None): + """Async helper to retrieve certs and compute fingerprints for mTLS rotation. + + Args: + cached_cert (bytes): The cached client certificate. + client_cert_callback (Optional[Callable[[], (bytes, bytes)]]): An + optional callback which returns client certificate bytes and private + key bytes both in PEM format. + + Returns: + bytes: The client callback cert bytes. + bytes: The client callback key bytes. + str: The base64-encoded SHA256 cached fingerprint. + str: The base64-encoded SHA256 current cert fingerprint. + + """ is_mtls, call_cert_bytes, call_key_bytes = await get_client_cert_and_key( client_cert_callback ) From 215857d6c540b8cb77a6c361fb6306eb579ca50e Mon Sep 17 00:00:00 2001 From: agrawalradhika-cell Date: Thu, 3 Sep 2026 13:45:43 -0700 Subject: [PATCH 76/79] chore: Store refresh counter during request handling Store refresh counter at error for better tracking. --- packages/google-auth/google/auth/aio/transport/sessions.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/google-auth/google/auth/aio/transport/sessions.py b/packages/google-auth/google/auth/aio/transport/sessions.py index ae5bfb0ff3e8..e95938bea88f 100644 --- a/packages/google-auth/google/auth/aio/transport/sessions.py +++ b/packages/google-auth/google/auth/aio/transport/sessions.py @@ -319,6 +319,7 @@ async def request( ) request_headers = dict(headers) if headers is not None else {} start_time = time.monotonic() + refresh_counter_at_error = self._refresh_counter async with timeout_guard(max_allowed_time) as with_timeout: await with_timeout( # Note: before_request will attempt to refresh credentials if expired. @@ -365,7 +366,6 @@ async def request( async def _recover_auth_state(): is_mtls_endpoint = False - refresh_counter_at_error = self._refresh_counter if getattr(self, "is_mtls", False): hostname = urllib.parse.urlsplit(url).hostname if hostname: From f5984327c62cb1de3f10d92fcb056c064fc50c1d Mon Sep 17 00:00:00 2001 From: agrawalradhika-cell Date: Thu, 3 Sep 2026 13:52:52 -0700 Subject: [PATCH 77/79] fix: Refactor test_mtls.py for improved mock handling Updated mock patches to ensure correct assertions and added bound mocks for certificate parsing and caching. --- .../tests/transport/aio/test_mtls.py | 40 ++++++++++++------- 1 file changed, 25 insertions(+), 15 deletions(-) diff --git a/packages/google-auth/tests/transport/aio/test_mtls.py b/packages/google-auth/tests/transport/aio/test_mtls.py index 5343e1473dd7..3006aae36dbe 100644 --- a/packages/google-auth/tests/transport/aio/test_mtls.py +++ b/packages/google-auth/tests/transport/aio/test_mtls.py @@ -62,7 +62,7 @@ def callback(): return CERT_BYTES, KEY_BYTES with ( - mock.patch("google.auth._agent_identity_utils.parse_certificate"), + mock.patch("google.auth._agent_identity_utils.parse_certificate") as mock_parse, mock.patch( "google.auth._agent_identity_utils.calculate_certificate_fingerprint", return_value="FINGERPRINT_A", @@ -70,7 +70,7 @@ def callback(): mock.patch( "google.auth._agent_identity_utils.get_cached_cert_fingerprint", return_value="FINGERPRINT_A", - ), + ) as mock_get_cached, ): ( cert, @@ -81,30 +81,36 @@ def callback(): cached_cert=CERT_BYTES, client_cert_callback=callback ) - assert cert == CERT_BYTES - assert key == KEY_BYTES - assert cached_fp == "FINGERPRINT_A" - assert current_fp == "FINGERPRINT_A" - assert cached_fp == current_fp + assert cert == CERT_BYTES + assert key == KEY_BYTES + assert cached_fp == "FINGERPRINT_A" + assert current_fp == "FINGERPRINT_A" + assert cached_fp == current_fp + + # These assertions will now work correctly using the bound mocks + mock_parse.assert_called_once_with(CERT_BYTES) + mock_get_cached.assert_called_once_with(CERT_BYTES) + @pytest.mark.asyncio async def test_check_parameters_cert_mismatch_rotation(): """Test when newly retrieved certificate differs from the cached certificate (rotation occurred).""" - def callback(): return NEW_CERT_BYTES, NEW_KEY_BYTES with ( - mock.patch("google.auth._agent_identity_utils.parse_certificate"), + # Add `as mock_parse` here + mock.patch("google.auth._agent_identity_utils.parse_certificate") as mock_parse, mock.patch( "google.auth._agent_identity_utils.calculate_certificate_fingerprint", return_value="FINGERPRINT_NEW", ), + # Add `as mock_get_cached` here mock.patch( "google.auth._agent_identity_utils.get_cached_cert_fingerprint", return_value="FINGERPRINT_OLD", - ), + ) as mock_get_cached, ): ( cert, @@ -115,11 +121,15 @@ def callback(): cached_cert=CERT_BYTES, client_cert_callback=callback ) - assert cert == NEW_CERT_BYTES - assert key == NEW_KEY_BYTES - assert cached_fp == "FINGERPRINT_OLD" - assert current_fp == "FINGERPRINT_NEW" - assert cached_fp != current_fp + assert cert == NEW_CERT_BYTES + assert key == NEW_KEY_BYTES + assert cached_fp == "FINGERPRINT_OLD" + assert current_fp == "FINGERPRINT_NEW" + assert cached_fp != current_fp + + # Now you can add the assertions requested by the reviewer at the end of the test: + mock_parse.assert_called_once_with(CERT_BYTES) + mock_get_cached.assert_called_once_with(CERT_BYTES) @pytest.mark.asyncio From aa0ee3780b3f73173ce6733ba5f51001770a700b Mon Sep 17 00:00:00 2001 From: agrawalradhika-cell Date: Thu, 3 Sep 2026 13:56:01 -0700 Subject: [PATCH 78/79] fix: Change mock responses to include 200 status code Update mock authentication response handling in tests. --- .../tests/transport/aio/test_sessions_mtls.py | 16 ++++++++++++---- 1 file changed, 12 insertions(+), 4 deletions(-) diff --git a/packages/google-auth/tests/transport/aio/test_sessions_mtls.py b/packages/google-auth/tests/transport/aio/test_sessions_mtls.py index f9aafdd8cff8..5d6e5c59a944 100644 --- a/packages/google-auth/tests/transport/aio/test_sessions_mtls.py +++ b/packages/google-auth/tests/transport/aio/test_sessions_mtls.py @@ -531,7 +531,13 @@ async def test_cert_rotation_lock_contention(self): mock_resp_401 = mock.Mock() mock_resp_401.status_code = http_client.UNAUTHORIZED - mock_auth_req = mock.AsyncMock(return_value=mock_resp_401) + + mock_resp_200 = mock.Mock() + mock_resp_200.status_code = http_client.OK + + mock_auth_req = mock.AsyncMock( + side_effect=[mock_resp_401] * 3 + [mock_resp_200] * 3 + ) session = sessions.AsyncAuthorizedSession( mock_creds, auth_request=mock_auth_req @@ -569,10 +575,12 @@ async def mock_check_side_effect(cached_cert, callback=None): ] responses = await asyncio.gather(*tasks) - for resp in responses: - assert resp == mock_resp_401 + for resp in responses: + assert resp == mock_resp_200 - mock_conf.assert_called_once() + mock_check.assert_called_once() + mock_conf.assert_called_once() + assert mock_creds.refresh.call_count == 1 await session.close() From 97e2b1a4ee66cc424ce2d0722c30b269de21b816 Mon Sep 17 00:00:00 2001 From: Radhika Agrawal Date: Thu, 3 Sep 2026 21:16:28 +0000 Subject: [PATCH 79/79] Fix: fix the unit tests in test_mtls Signed-off-by: Radhika Agrawal --- .../google-auth/google/auth/aio/transport/mtls.py | 8 +++++--- .../google-auth/tests/transport/aio/test_mtls.py | 13 ++++++------- .../tests/transport/aio/test_sessions_mtls.py | 4 ++-- 3 files changed, 13 insertions(+), 12 deletions(-) diff --git a/packages/google-auth/google/auth/aio/transport/mtls.py b/packages/google-auth/google/auth/aio/transport/mtls.py index 45e471ba099c..46310b856779 100644 --- a/packages/google-auth/google/auth/aio/transport/mtls.py +++ b/packages/google-auth/google/auth/aio/transport/mtls.py @@ -179,7 +179,9 @@ async def get_client_cert_and_key(client_cert_callback=None): return has_cert, cert, key -async def check_parameters_for_unauthorized_response(cached_cert, client_cert_callback=None): +async def check_parameters_for_unauthorized_response( + cached_cert, client_cert_callback=None +): """Async helper to retrieve certs and compute fingerprints for mTLS rotation. Args: @@ -187,13 +189,13 @@ async def check_parameters_for_unauthorized_response(cached_cert, client_cert_ca client_cert_callback (Optional[Callable[[], (bytes, bytes)]]): An optional callback which returns client certificate bytes and private key bytes both in PEM format. - + Returns: bytes: The client callback cert bytes. bytes: The client callback key bytes. str: The base64-encoded SHA256 cached fingerprint. str: The base64-encoded SHA256 current cert fingerprint. - + """ is_mtls, call_cert_bytes, call_key_bytes = await get_client_cert_and_key( client_cert_callback diff --git a/packages/google-auth/tests/transport/aio/test_mtls.py b/packages/google-auth/tests/transport/aio/test_mtls.py index 3006aae36dbe..7d583dcfb108 100644 --- a/packages/google-auth/tests/transport/aio/test_mtls.py +++ b/packages/google-auth/tests/transport/aio/test_mtls.py @@ -86,27 +86,25 @@ def callback(): assert cached_fp == "FINGERPRINT_A" assert current_fp == "FINGERPRINT_A" assert cached_fp == current_fp - + # These assertions will now work correctly using the bound mocks mock_parse.assert_called_once_with(CERT_BYTES) mock_get_cached.assert_called_once_with(CERT_BYTES) - @pytest.mark.asyncio async def test_check_parameters_cert_mismatch_rotation(): """Test when newly retrieved certificate differs from the cached certificate (rotation occurred).""" + def callback(): return NEW_CERT_BYTES, NEW_KEY_BYTES with ( - # Add `as mock_parse` here - mock.patch("google.auth._agent_identity_utils.parse_certificate") as mock_parse, + mock.patch("google.auth._agent_identity_utils.parse_certificate") as mock_parse, mock.patch( "google.auth._agent_identity_utils.calculate_certificate_fingerprint", return_value="FINGERPRINT_NEW", ), - # Add `as mock_get_cached` here mock.patch( "google.auth._agent_identity_utils.get_cached_cert_fingerprint", return_value="FINGERPRINT_OLD", @@ -127,8 +125,9 @@ def callback(): assert current_fp == "FINGERPRINT_NEW" assert cached_fp != current_fp - # Now you can add the assertions requested by the reviewer at the end of the test: - mock_parse.assert_called_once_with(CERT_BYTES) + # The fix: mock_parse is called with the NEW cert from the callback + mock_parse.assert_called_once_with(NEW_CERT_BYTES) + # mock_get_cached is called with the old cached cert mock_get_cached.assert_called_once_with(CERT_BYTES) diff --git a/packages/google-auth/tests/transport/aio/test_sessions_mtls.py b/packages/google-auth/tests/transport/aio/test_sessions_mtls.py index 5d6e5c59a944..0972a6180299 100644 --- a/packages/google-auth/tests/transport/aio/test_sessions_mtls.py +++ b/packages/google-auth/tests/transport/aio/test_sessions_mtls.py @@ -531,10 +531,10 @@ async def test_cert_rotation_lock_contention(self): mock_resp_401 = mock.Mock() mock_resp_401.status_code = http_client.UNAUTHORIZED - + mock_resp_200 = mock.Mock() mock_resp_200.status_code = http_client.OK - + mock_auth_req = mock.AsyncMock( side_effect=[mock_resp_401] * 3 + [mock_resp_200] * 3 )