Skip to content

Commit e0a789b

Browse files
twishabansaltwishabansal
committed
lint
1 parent 71fac64 commit e0a789b

File tree

2 files changed

+55
-25
lines changed

2 files changed

+55
-25
lines changed

‎packages/toolbox-core/src/toolbox_core/auth_methods.py

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -71,7 +71,10 @@ def _is_cached_token_valid(
7171
return False
7272

7373
# Ensure expires_at_value is timezone-aware (UTC).
74-
if expires_at_value.tzinfo is None or expires_at_value.tzinfo.utcoffset(expires_at_value) is None:
74+
if (
75+
expires_at_value.tzinfo is None
76+
or expires_at_value.tzinfo.utcoffset(expires_at_value) is None
77+
):
7578
expires_at_value = expires_at_value.replace(tzinfo=timezone.utc)
7679

7780
current_time_utc = datetime.now(timezone.utc)

‎packages/toolbox-core/tests/test_auth_methods.py

Lines changed: 51 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@
1212
# See the License for the specific language governing permissions and
1313
# limitations under the License.
1414

15-
from unittest.mock import AsyncMock, MagicMock, patch, PropertyMock
15+
from unittest.mock import AsyncMock, MagicMock, PropertyMock, patch
1616

1717
import pytest
1818

@@ -22,7 +22,9 @@
2222
MOCK_GOOGLE_ID_TOKEN = "test_id_token_123"
2323
MOCK_PROJECT_ID = "test-project"
2424
# A realistic expiry timestamp (e.g., 1 hour from now)
25-
MOCK_EXPIRY_DATETIME = auth_methods.datetime.now(auth_methods.timezone.utc) + auth_methods.timedelta(hours=1)
25+
MOCK_EXPIRY_DATETIME = auth_methods.datetime.now(
26+
auth_methods.timezone.utc
27+
) + auth_methods.timedelta(hours=1)
2628

2729

2830
# Expected exception messages from auth_methods.py
@@ -56,7 +58,9 @@ async def test_aget_google_id_token_success_first_call(
5658
"""Tests successful fetching of an async token on the first call."""
5759
mock_creds_instance = AsyncMock()
5860
mock_creds_instance.id_token = MOCK_GOOGLE_ID_TOKEN
59-
type(mock_creds_instance).expiry = PropertyMock(return_value=MOCK_EXPIRY_DATETIME)
61+
type(mock_creds_instance).expiry = PropertyMock(
62+
return_value=MOCK_EXPIRY_DATETIME
63+
)
6064
mock_default_async.return_value = (mock_creds_instance, MOCK_PROJECT_ID)
6165

6266
mock_async_req_instance = MagicMock()
@@ -85,9 +89,10 @@ async def test_aget_google_id_token_success_uses_cache(
8589
):
8690
"""Tests that subsequent calls use the cached token if valid."""
8791
auth_methods._cached_google_id_token["token"] = MOCK_GOOGLE_ID_TOKEN
88-
auth_methods._cached_google_id_token["expires_at"] = (
89-
auth_methods.datetime.now(auth_methods.timezone.utc) +
90-
auth_methods.timedelta(seconds=auth_methods.CACHE_REFRESH_MARGIN_SECONDS + 100)
92+
auth_methods._cached_google_id_token["expires_at"] = auth_methods.datetime.now(
93+
auth_methods.timezone.utc
94+
) + auth_methods.timedelta(
95+
seconds=auth_methods.CACHE_REFRESH_MARGIN_SECONDS + 100
9196
) # Ensure it's valid
9297

9398
token = await auth_methods.aget_google_id_token()
@@ -106,13 +111,17 @@ async def test_aget_google_id_token_refreshes_expired_cache(
106111
):
107112
"""Tests that an expired cached token triggers a refresh."""
108113
auth_methods._cached_google_id_token["token"] = "expired_token"
109-
auth_methods._cached_google_id_token["expires_at"] = (
110-
auth_methods.datetime.now(auth_methods.timezone.utc) - auth_methods.timedelta(seconds=100)
114+
auth_methods._cached_google_id_token["expires_at"] = auth_methods.datetime.now(
115+
auth_methods.timezone.utc
116+
) - auth_methods.timedelta(
117+
seconds=100
111118
) # Expired
112119

113120
mock_creds_instance = AsyncMock()
114121
mock_creds_instance.id_token = MOCK_GOOGLE_ID_TOKEN # New token after refresh
115-
type(mock_creds_instance).expiry = PropertyMock(return_value=MOCK_EXPIRY_DATETIME)
122+
type(mock_creds_instance).expiry = PropertyMock(
123+
return_value=MOCK_EXPIRY_DATETIME
124+
)
116125
mock_default_async.return_value = (mock_creds_instance, MOCK_PROJECT_ID)
117126

118127
mock_async_req_instance = MagicMock()
@@ -125,8 +134,9 @@ async def test_aget_google_id_token_refreshes_expired_cache(
125134
mock_creds_instance.refresh.assert_called_once_with(mock_async_req_instance)
126135
assert token == f"{auth_methods.BEARER_TOKEN_PREFIX}{MOCK_GOOGLE_ID_TOKEN}"
127136
assert auth_methods._cached_google_id_token["token"] == MOCK_GOOGLE_ID_TOKEN
128-
assert auth_methods._cached_google_id_token["expires_at"] == MOCK_EXPIRY_DATETIME
129-
137+
assert (
138+
auth_methods._cached_google_id_token["expires_at"] == MOCK_EXPIRY_DATETIME
139+
)
130140

131141
@pytest.mark.asyncio
132142
@patch("toolbox_core.auth_methods._aiohttp_requests.Request")
@@ -137,7 +147,9 @@ async def test_aget_google_id_token_fetch_failure(
137147
"""Tests error handling when fetching the token fails (no id_token returned)."""
138148
mock_creds_instance = AsyncMock()
139149
mock_creds_instance.id_token = None # Simulate no ID token after refresh
140-
type(mock_creds_instance).expiry = PropertyMock(return_value=MOCK_EXPIRY_DATETIME) # Still need expiry for update_cache
150+
type(mock_creds_instance).expiry = PropertyMock(
151+
return_value=MOCK_EXPIRY_DATETIME
152+
) # Still need expiry for update_cache
141153
mock_default_async.return_value = (mock_creds_instance, MOCK_PROJECT_ID)
142154
mock_async_req_class.return_value = MagicMock()
143155

@@ -178,11 +190,13 @@ async def test_aget_google_id_token_no_expiry_info(
178190
"""Tests that a token without expiry info is still cached but effectively expired."""
179191
mock_creds_instance = AsyncMock()
180192
mock_creds_instance.id_token = MOCK_GOOGLE_ID_TOKEN
181-
type(mock_creds_instance).expiry = PropertyMock(return_value=None) # Simulate no expiry info
193+
type(mock_creds_instance).expiry = PropertyMock(
194+
return_value=None
195+
) # Simulate no expiry info
182196
mock_default_async.return_value = (mock_creds_instance, MOCK_PROJECT_ID)
183197

184198
mock_async_req_class.return_value = MagicMock()
185-
199+
186200
token = await auth_methods.aget_google_id_token()
187201

188202
assert token == f"{auth_methods.BEARER_TOKEN_PREFIX}{MOCK_GOOGLE_ID_TOKEN}"
@@ -208,7 +222,9 @@ def test_get_google_id_token_success_first_call(
208222
"""Tests successful fetching of a sync token on the first call."""
209223
mock_creds_instance = MagicMock()
210224
mock_creds_instance.id_token = MOCK_GOOGLE_ID_TOKEN
211-
type(mock_creds_instance).expiry = PropertyMock(return_value=MOCK_EXPIRY_DATETIME)
225+
type(mock_creds_instance).expiry = PropertyMock(
226+
return_value=MOCK_EXPIRY_DATETIME
227+
)
212228
mock_sync_default.return_value = (mock_creds_instance, MOCK_PROJECT_ID)
213229

214230
mock_session_instance = MagicMock()
@@ -241,9 +257,10 @@ def test_get_google_id_token_success_uses_cache(
241257
):
242258
"""Tests that subsequent calls use the cached token if valid."""
243259
auth_methods._cached_google_id_token["token"] = MOCK_GOOGLE_ID_TOKEN
244-
auth_methods._cached_google_id_token["expires_at"] = (
245-
auth_methods.datetime.now(auth_methods.timezone.utc) +
246-
auth_methods.timedelta(seconds=auth_methods.CACHE_REFRESH_MARGIN_SECONDS + 100)
260+
auth_methods._cached_google_id_token["expires_at"] = auth_methods.datetime.now(
261+
auth_methods.timezone.utc
262+
) + auth_methods.timedelta(
263+
seconds=auth_methods.CACHE_REFRESH_MARGIN_SECONDS + 100
247264
) # Ensure it's valid
248265

249266
token = auth_methods.get_google_id_token()
@@ -267,13 +284,17 @@ def test_get_google_id_token_refreshes_expired_cache(
267284
"""Tests that an expired cached token triggers a refresh."""
268285
# Prime the cache with an expired token
269286
auth_methods._cached_google_id_token["token"] = "expired_token_sync"
270-
auth_methods._cached_google_id_token["expires_at"] = (
271-
auth_methods.datetime.now(auth_methods.timezone.utc) - auth_methods.timedelta(seconds=100)
287+
auth_methods._cached_google_id_token["expires_at"] = auth_methods.datetime.now(
288+
auth_methods.timezone.utc
289+
) - auth_methods.timedelta(
290+
seconds=100
272291
) # Expired
273292

274293
mock_creds_instance = MagicMock()
275294
mock_creds_instance.id_token = MOCK_GOOGLE_ID_TOKEN # New token after refresh
276-
type(mock_creds_instance).expiry = PropertyMock(return_value=MOCK_EXPIRY_DATETIME)
295+
type(mock_creds_instance).expiry = PropertyMock(
296+
return_value=MOCK_EXPIRY_DATETIME
297+
)
277298
mock_sync_default.return_value = (mock_creds_instance, MOCK_PROJECT_ID)
278299

279300
mock_session_instance = MagicMock()
@@ -290,7 +311,9 @@ def test_get_google_id_token_refreshes_expired_cache(
290311
mock_creds_instance.refresh.assert_called_once_with(mock_sync_request_instance)
291312
assert token == f"{auth_methods.BEARER_TOKEN_PREFIX}{MOCK_GOOGLE_ID_TOKEN}"
292313
assert auth_methods._cached_google_id_token["token"] == MOCK_GOOGLE_ID_TOKEN
293-
assert auth_methods._cached_google_id_token["expires_at"] == MOCK_EXPIRY_DATETIME
314+
assert (
315+
auth_methods._cached_google_id_token["expires_at"] == MOCK_EXPIRY_DATETIME
316+
)
294317

295318
@patch("toolbox_core.auth_methods.Request")
296319
@patch("toolbox_core.auth_methods.AuthorizedSession")
@@ -301,7 +324,9 @@ def test_get_google_id_token_fetch_failure(
301324
"""Tests error handling when fetching the token fails (no id_token returned)."""
302325
mock_creds_instance = MagicMock()
303326
mock_creds_instance.id_token = None # Simulate no ID token after refresh
304-
type(mock_creds_instance).expiry = PropertyMock(return_value=MOCK_EXPIRY_DATETIME) # Still need expiry for update_cache
327+
type(mock_creds_instance).expiry = PropertyMock(
328+
return_value=MOCK_EXPIRY_DATETIME
329+
) # Still need expiry for update_cache
305330
mock_sync_default.return_value = (mock_creds_instance, MOCK_PROJECT_ID)
306331

307332
mock_session_instance = MagicMock()
@@ -357,7 +382,9 @@ def test_get_google_id_token_no_expiry_info(
357382
"""Tests that a token without expiry info is still cached but effectively expired."""
358383
mock_creds_instance = MagicMock()
359384
mock_creds_instance.id_token = MOCK_GOOGLE_ID_TOKEN
360-
type(mock_creds_instance).expiry = PropertyMock(return_value=None) # Simulate no expiry info
385+
type(mock_creds_instance).expiry = PropertyMock(
386+
return_value=None
387+
) # Simulate no expiry info
361388
mock_sync_default.return_value = (mock_creds_instance, MOCK_PROJECT_ID)
362389

363390
mock_session_instance = MagicMock()

0 commit comments

Comments
 (0)