Skip to content

Commit b992294

Browse files
authored
feat(storage): add async credential wrapper (#1659)
Adds async credential wrapper for async client. Context: - The standard google auth credentials are currently synchronous, and it's asynchronous credentials classes are either not available or marked private. - Credential retrieval and refreshing will remain synchronous under the hood. Rationale: As authentication tokens typically possess an expiration lifetime of one hour, the blocking time required for token fetching occurs infrequently. The performance impact of blocking (or utilizing a separate thread for offloading) once per hour is deemed negligible when weighed against the considerable engineering cost of developing and maintaining a asynchronous authentication.
1 parent 096642f commit b992294

File tree

2 files changed

+165
-0
lines changed

2 files changed

+165
-0
lines changed
Lines changed: 57 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,57 @@
1+
"""Async Wrapper around Google Auth Credentials"""
2+
3+
import asyncio
4+
from google.auth.transport.requests import Request
5+
6+
try:
7+
from google.auth.aio import credentials as aio_creds_module
8+
BaseCredentials = aio_creds_module.Credentials
9+
_AIO_AVAILABLE = True
10+
except ImportError:
11+
BaseCredentials = object
12+
_AIO_AVAILABLE = False
13+
14+
class AsyncCredsWrapper(BaseCredentials):
15+
"""Wraps synchronous Google Auth credentials to provide an asynchronous interface.
16+
17+
Args:
18+
sync_creds (google.auth.credentials.Credentials): The synchronous credentials
19+
instance to wrap.
20+
21+
Raises:
22+
ImportError: If instantiated in an environment where 'google.auth.aio'
23+
is not available.
24+
"""
25+
26+
def __init__(self, sync_creds):
27+
if not _AIO_AVAILABLE:
28+
raise ImportError(
29+
"Failed to import 'google.auth.aio'. This module requires a newer version "
30+
"of 'google-auth' which supports asyncio."
31+
)
32+
33+
super().__init__()
34+
self.creds = sync_creds
35+
36+
async def refresh(self, request):
37+
"""Refreshes the access token."""
38+
loop = asyncio.get_running_loop()
39+
await loop.run_in_executor(
40+
None, self.creds.refresh, Request()
41+
)
42+
43+
@property
44+
def valid(self):
45+
"""Checks the validity of the credentials."""
46+
return self.creds.valid
47+
48+
async def before_request(self, request, method, url, headers):
49+
"""Performs credential-specific before request logic."""
50+
if self.valid:
51+
self.creds.apply(headers)
52+
return
53+
54+
loop = asyncio.get_running_loop()
55+
await loop.run_in_executor(
56+
None, self.creds.before_request, Request(), method, url, headers
57+
)
Lines changed: 108 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,108 @@
1+
import sys
2+
import unittest.mock
3+
import pytest
4+
from google.auth import credentials as google_creds
5+
from google.cloud.storage._experimental.asyncio import async_creds
6+
7+
@pytest.fixture
8+
def mock_aio_modules():
9+
"""Patches sys.modules to simulate google.auth.aio existence."""
10+
mock_creds_module = unittest.mock.MagicMock()
11+
# We must set the base class to object so our wrapper can inherit safely in tests
12+
mock_creds_module.Credentials = object
13+
14+
modules = {
15+
'google.auth.aio': unittest.mock.MagicMock(),
16+
'google.auth.aio.credentials': mock_creds_module,
17+
}
18+
19+
with unittest.mock.patch.dict(sys.modules, modules):
20+
# We also need to manually flip the flag in the module to True for the test context
21+
# because the module was likely already imported with the flag set to False/True
22+
# depending on the real environment.
23+
with unittest.mock.patch.object(async_creds, '_AIO_AVAILABLE', True):
24+
# We also need to ensure BaseCredentials in the module points to our mock
25+
# if we want strictly correct inheritance, though duck typing usually suffices.
26+
with unittest.mock.patch.object(async_creds, 'BaseCredentials', object):
27+
yield
28+
29+
@pytest.fixture
30+
def mock_sync_creds():
31+
"""Creates a mock of the synchronous Google Credentials object."""
32+
creds = unittest.mock.create_autospec(google_creds.Credentials, instance=True)
33+
type(creds).valid = unittest.mock.PropertyMock(return_value=True)
34+
return creds
35+
36+
@pytest.fixture
37+
def async_wrapper(mock_aio_modules, mock_sync_creds):
38+
"""Instantiates the wrapper with the mock credentials."""
39+
# This instantiation would raise ImportError if mock_aio_modules didn't set _AIO_AVAILABLE=True
40+
return async_creds.AsyncCredsWrapper(mock_sync_creds)
41+
42+
class TestAsyncCredsWrapper:
43+
44+
@pytest.mark.asyncio
45+
async def test_init_sets_attributes(self, async_wrapper, mock_sync_creds):
46+
"""Test that the wrapper initializes correctly."""
47+
assert async_wrapper.creds == mock_sync_creds
48+
49+
@pytest.mark.asyncio
50+
async def test_valid_property_delegates(self, async_wrapper, mock_sync_creds):
51+
"""Test that the .valid property maps to the sync creds .valid property."""
52+
type(mock_sync_creds).valid = unittest.mock.PropertyMock(return_value=True)
53+
assert async_wrapper.valid is True
54+
55+
type(mock_sync_creds).valid = unittest.mock.PropertyMock(return_value=False)
56+
assert async_wrapper.valid is False
57+
58+
@pytest.mark.asyncio
59+
async def test_refresh_offloads_to_executor(self, async_wrapper, mock_sync_creds):
60+
"""Test that refresh() gets the running loop and calls sync refresh in executor."""
61+
with unittest.mock.patch('asyncio.get_running_loop') as mock_get_loop:
62+
mock_loop = unittest.mock.AsyncMock()
63+
mock_get_loop.return_value = mock_loop
64+
65+
await async_wrapper.refresh(None)
66+
67+
mock_loop.run_in_executor.assert_called_once()
68+
args, _ = mock_loop.run_in_executor.call_args
69+
assert args[1] == mock_sync_creds.refresh
70+
71+
@pytest.mark.asyncio
72+
async def test_before_request_valid_creds(self, async_wrapper, mock_sync_creds):
73+
"""Test before_request when credentials are ALREADY valid."""
74+
type(mock_sync_creds).valid = unittest.mock.PropertyMock(return_value=True)
75+
76+
headers = {}
77+
await async_wrapper.before_request(None, "GET", "http://example.com", headers)
78+
79+
mock_sync_creds.apply.assert_called_once_with(headers)
80+
mock_sync_creds.before_request.assert_not_called()
81+
82+
@pytest.mark.asyncio
83+
async def test_before_request_invalid_creds(self, async_wrapper, mock_sync_creds):
84+
"""Test before_request when credentials are INVALID (refresh path)."""
85+
type(mock_sync_creds).valid = unittest.mock.PropertyMock(return_value=False)
86+
87+
headers = {}
88+
method = "GET"
89+
url = "http://example.com"
90+
91+
with unittest.mock.patch('asyncio.get_running_loop') as mock_get_loop:
92+
mock_loop = unittest.mock.AsyncMock()
93+
mock_get_loop.return_value = mock_loop
94+
95+
await async_wrapper.before_request(None, method, url, headers)
96+
97+
mock_loop.run_in_executor.assert_called_once()
98+
args, _ = mock_loop.run_in_executor.call_args
99+
assert args[1] == mock_sync_creds.before_request
100+
101+
def test_missing_aio_raises_error(self, mock_sync_creds):
102+
"""Ensure ImportError is raised if _AIO_AVAILABLE is False."""
103+
# We manually simulate the environment where AIO is missing
104+
with unittest.mock.patch.object(async_creds, '_AIO_AVAILABLE', False):
105+
with pytest.raises(ImportError) as excinfo:
106+
async_creds.AsyncCredsWrapper(mock_sync_creds)
107+
108+
assert "Failed to import 'google.auth.aio'" in str(excinfo.value)

0 commit comments

Comments
 (0)