Skip to content

Commit ee0aa61

Browse files
twishabansaltwishabansal
committed
add tests
1 parent 6626f69 commit ee0aa61

File tree

3 files changed

+200
-2
lines changed

3 files changed

+200
-2
lines changed

‎packages/toolbox-core/pyproject.toml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,7 @@ test = [
4848
"pytest-aioresponses==0.3.0",
4949
"pytest-asyncio==0.26.0",
5050
"pytest-cov==6.1.1",
51+
"pytest-mock==3.14.0",
5152
"google-cloud-secret-manager==2.23.3",
5253
"google-cloud-storage==3.1.0",
5354
]

‎packages/toolbox-core/src/toolbox_core/toolbox_auth_methods.py renamed to ‎packages/toolbox-core/src/toolbox_core/auth_methods.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -16,9 +16,9 @@
1616
# and are intended to be passed in the "Authorization" header of HTTP requests.
1717
#
1818
# Example User Experience:
19-
# from toolbox_core import toolbox_auth_methods
19+
# from toolbox_core import auth_methods
2020
#
21-
# auth_token_provider = toolbox_auth_methods.aget_google_id_token
21+
# auth_token_provider = auth_methods.aget_google_id_token
2222
# toolbox = ToolboxClient(
2323
# URL,
2424
# client_headers={"Authorization": auth_token_provider},

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

Lines changed: 197 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,197 @@
1+
# Copyright 2025 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
import google.auth.exceptions
16+
import pytest
17+
from unittest.mock import patch, AsyncMock, MagicMock
18+
19+
from toolbox_core import auth_methods
20+
21+
@pytest.mark.asyncio
22+
@patch("toolbox_core.auth_methods.partial")
23+
@patch("toolbox_core.auth_methods._aiohttp_requests.Request")
24+
@patch("toolbox_core.auth_methods.Credentials")
25+
@patch("toolbox_core.auth_methods.default_async")
26+
async def test_aget_google_id_token_success(
27+
mock_default_async,
28+
mock_credentials_class,
29+
mock_aiohttp_request_class,
30+
mock_partial,
31+
):
32+
"""
33+
Test aget_google_id_token successfully retrieves and formats a token using pytest.
34+
"""
35+
# Setup mock for default_async() -> (creds, project_id)
36+
mock_creds_instance = AsyncMock()
37+
mock_creds_instance.id_token = "test_async_id_token_123"
38+
mock_default_async.return_value = (mock_creds_instance, "test-project")
39+
40+
# Setup mock for _aiohttp_requests.Request()
41+
mock_aio_request_instance = MagicMock()
42+
mock_aiohttp_request_class.return_value = mock_aio_request_instance
43+
44+
# Setup mock for Credentials.before_request (class attribute used in partial)
45+
mock_unbound_before_request = MagicMock()
46+
mock_credentials_class.before_request = mock_unbound_before_request
47+
48+
# Setup mock for partial()
49+
mock_partial_object = MagicMock()
50+
mock_partial.return_value = mock_partial_object
51+
52+
token = await auth_methods.aget_google_id_token()
53+
54+
mock_default_async.assert_called_once_with()
55+
mock_aiohttp_request_class.assert_called_once_with()
56+
mock_creds_instance.refresh.assert_called_once_with(mock_aio_request_instance)
57+
58+
mock_partial.assert_called_once_with(
59+
mock_unbound_before_request, mock_creds_instance
60+
)
61+
assert mock_creds_instance.before_request == mock_partial_object
62+
63+
assert token == "Bearer test_async_id_token_123"
64+
65+
66+
@pytest.mark.asyncio
67+
@patch("toolbox_core.auth_methods.default_async")
68+
async def test_aget_google_id_token_default_credentials_error(mock_default_async):
69+
"""
70+
Test aget_google_id_token when default_async raises DefaultCredentialsError.
71+
"""
72+
mock_default_async.side_effect = google.auth.exceptions.DefaultCredentialsError(
73+
"ADC not found"
74+
)
75+
76+
with pytest.raises(
77+
google.auth.exceptions.DefaultCredentialsError, match="ADC not found"
78+
):
79+
await auth_methods.aget_google_id_token()
80+
81+
mock_default_async.assert_called_once_with()
82+
83+
84+
@pytest.mark.asyncio
85+
@patch("toolbox_core.auth_methods._aiohttp_requests.Request")
86+
@patch("toolbox_core.auth_methods.default_async")
87+
async def test_aget_google_id_token_refresh_error(
88+
mock_default_async,
89+
mock_aiohttp_request_class,
90+
):
91+
"""
92+
Test aget_google_id_token when creds.refresh raises RefreshError.
93+
The `partial` call should not happen if refresh fails.
94+
"""
95+
mock_creds_instance = AsyncMock()
96+
mock_creds_instance.refresh.side_effect = google.auth.exceptions.RefreshError(
97+
"Token refresh failed"
98+
)
99+
mock_default_async.return_value = (mock_creds_instance, "test-project")
100+
101+
mock_aio_request_instance = MagicMock()
102+
mock_aiohttp_request_class.return_value = mock_aio_request_instance
103+
104+
with pytest.raises(
105+
google.auth.exceptions.RefreshError, match="Token refresh failed"
106+
):
107+
await auth_methods.aget_google_id_token()
108+
109+
mock_default_async.assert_called_once_with()
110+
mock_aiohttp_request_class.assert_called_once_with()
111+
mock_creds_instance.refresh.assert_called_once_with(mock_aio_request_instance)
112+
113+
114+
# --- Synchronous Tests ---
115+
116+
@patch("toolbox_core.auth_methods.Request")
117+
@patch("toolbox_core.auth_methods.AuthorizedSession")
118+
@patch("toolbox_core.auth_methods.google.auth.default")
119+
def test_get_google_id_token_success(
120+
mock_google_auth_default,
121+
mock_authorized_session_class,
122+
mock_request_class,
123+
):
124+
"""
125+
Test get_google_id_token successfully retrieves and formats a token using pytest.
126+
"""
127+
# Setup mock for google.auth.default() -> (credentials, project_id)
128+
mock_creds_instance = MagicMock()
129+
mock_creds_instance.id_token = "test_sync_id_token_456"
130+
mock_google_auth_default.return_value = (mock_creds_instance, "test-project")
131+
132+
# Setup mock for AuthorizedSession()
133+
mock_session_instance = MagicMock()
134+
mock_authorized_session_class.return_value = mock_session_instance
135+
136+
# Setup mock for Request()
137+
mock_request_instance = MagicMock()
138+
mock_request_class.return_value = mock_request_instance
139+
140+
token = auth_methods.get_google_id_token()
141+
142+
mock_google_auth_default.assert_called_once_with()
143+
mock_authorized_session_class.assert_called_once_with(mock_creds_instance)
144+
mock_request_class.assert_called_once_with(mock_session_instance)
145+
mock_creds_instance.refresh.assert_called_once_with(mock_request_instance)
146+
assert token == "Bearer test_sync_id_token_456"
147+
148+
149+
@patch("toolbox_core.auth_methods.google.auth.default")
150+
def test_get_google_id_token_default_credentials_error(mock_google_auth_default):
151+
"""
152+
Test get_google_id_token when google.auth.default raises DefaultCredentialsError.
153+
"""
154+
mock_google_auth_default.side_effect = (
155+
google.auth.exceptions.DefaultCredentialsError("Sync ADC not found")
156+
)
157+
158+
with pytest.raises(
159+
google.auth.exceptions.DefaultCredentialsError, match="Sync ADC not found"
160+
):
161+
auth_methods.get_google_id_token()
162+
163+
mock_google_auth_default.assert_called_once_with()
164+
165+
166+
@patch("toolbox_core.auth_methods.Request")
167+
@patch("toolbox_core.auth_methods.AuthorizedSession")
168+
@patch("toolbox_core.auth_methods.google.auth.default")
169+
def test_get_google_id_token_refresh_error(
170+
mock_google_auth_default,
171+
mock_authorized_session_class,
172+
mock_request_class,
173+
):
174+
"""
175+
Test get_google_id_token when credentials.refresh raises RefreshError.
176+
"""
177+
mock_creds_instance = MagicMock()
178+
mock_creds_instance.refresh.side_effect = google.auth.exceptions.RefreshError(
179+
"Sync token refresh failed"
180+
)
181+
mock_google_auth_default.return_value = (mock_creds_instance, "test-project")
182+
183+
mock_session_instance = MagicMock()
184+
mock_authorized_session_class.return_value = mock_session_instance
185+
186+
mock_request_instance = MagicMock()
187+
mock_request_class.return_value = mock_request_instance
188+
189+
with pytest.raises(
190+
google.auth.exceptions.RefreshError, match="Sync token refresh failed"
191+
):
192+
auth_methods.get_google_id_token()
193+
194+
mock_google_auth_default.assert_called_once_with()
195+
mock_authorized_session_class.assert_called_once_with(mock_creds_instance)
196+
mock_request_class.assert_called_once_with(mock_session_instance)
197+
mock_creds_instance.refresh.assert_called_once_with(mock_request_instance)

0 commit comments

Comments
 (0)