|
13 | 13 | # limitations under the License. |
14 | 14 |
|
15 | 15 | import asyncio |
| 16 | +import time |
16 | 17 | from unittest.mock import ANY |
17 | 18 | from unittest.mock import AsyncMock |
18 | 19 | from unittest.mock import Mock |
@@ -1180,6 +1181,175 @@ async def test_refresh_credential_no_refresher(self): |
1180 | 1181 | assert result == mock_credential |
1181 | 1182 | assert was_refreshed is False |
1182 | 1183 |
|
| 1184 | + @pytest.mark.asyncio |
| 1185 | + @pytest.mark.parametrize( |
| 1186 | + "auth_scheme, auth_type", |
| 1187 | + [ |
| 1188 | + ( |
| 1189 | + OAuth2( |
| 1190 | + flows=OAuthFlows( |
| 1191 | + clientCredentials=OAuthFlowClientCredentials( |
| 1192 | + tokenUrl="https://example.com/token", |
| 1193 | + scopes={}, |
| 1194 | + ) |
| 1195 | + ) |
| 1196 | + ), |
| 1197 | + AuthCredentialTypes.OAUTH2, |
| 1198 | + ), |
| 1199 | + ( |
| 1200 | + OpenIdConnectWithConfig( |
| 1201 | + authorization_endpoint="https://example.com/auth", |
| 1202 | + token_endpoint="https://example.com/token", |
| 1203 | + grant_types_supported=["client_credentials"], |
| 1204 | + ), |
| 1205 | + AuthCredentialTypes.OPEN_ID_CONNECT, |
| 1206 | + ), |
| 1207 | + ], |
| 1208 | + ids=["oauth2", "oidc"], |
| 1209 | + ) |
| 1210 | + async def test_get_auth_credential_refetches_expired_client_credentials_token( |
| 1211 | + self, auth_scheme, auth_type |
| 1212 | + ): |
| 1213 | + """Client credentials tokens have no refresh token, so fetch a new one.""" |
| 1214 | + raw_credential = AuthCredential( |
| 1215 | + auth_type=auth_type, |
| 1216 | + oauth2=OAuth2Auth( |
| 1217 | + client_id="mock_client_id", |
| 1218 | + client_secret="mock_client_secret", |
| 1219 | + ), |
| 1220 | + ) |
| 1221 | + stored_credential = raw_credential.model_copy(deep=True) |
| 1222 | + stored_credential.oauth2.access_token = "expired_token" |
| 1223 | + stored_credential.oauth2.expires_at = int(time.time()) - 60 |
| 1224 | + manager = CredentialManager( |
| 1225 | + AuthConfig(auth_scheme=auth_scheme, raw_auth_credential=raw_credential) |
| 1226 | + ) |
| 1227 | + manager._load_existing_credential = AsyncMock( |
| 1228 | + return_value=stored_credential |
| 1229 | + ) |
| 1230 | + manager._save_credential = AsyncMock() |
| 1231 | + context = Mock() |
| 1232 | + |
| 1233 | + with patch.object( |
| 1234 | + oauth2_credential_util, "OAuth2Session", autospec=True |
| 1235 | + ) as mock_oauth2_session: |
| 1236 | + mock_client = mock_oauth2_session.return_value |
| 1237 | + mock_client.fetch_token.return_value = { |
| 1238 | + "access_token": "new_token", |
| 1239 | + "expires_at": int(time.time()) + 3600, |
| 1240 | + } |
| 1241 | + # There is no refresh token to send. |
| 1242 | + mock_client.refresh_token.side_effect = AssertionError( |
| 1243 | + "refresh_token() called without a refresh token" |
| 1244 | + ) |
| 1245 | + result = await manager.get_auth_credential(context) |
| 1246 | + |
| 1247 | + assert result.oauth2.access_token == "new_token" |
| 1248 | + mock_client.fetch_token.assert_called_once() |
| 1249 | + manager._save_credential.assert_awaited_once_with(context, result) |
| 1250 | + # The shared tool config is not mutated. |
| 1251 | + assert raw_credential.oauth2.access_token is None |
| 1252 | + |
| 1253 | + @pytest.mark.asyncio |
| 1254 | + async def test_get_auth_credential_refreshes_token_that_has_refresh_token( |
| 1255 | + self, |
| 1256 | + ): |
| 1257 | + """A client credentials token that has a refresh token is refreshed.""" |
| 1258 | + auth_scheme = OAuth2( |
| 1259 | + flows=OAuthFlows( |
| 1260 | + clientCredentials=OAuthFlowClientCredentials( |
| 1261 | + tokenUrl="https://example.com/token", |
| 1262 | + scopes={}, |
| 1263 | + ) |
| 1264 | + ) |
| 1265 | + ) |
| 1266 | + raw_credential = AuthCredential( |
| 1267 | + auth_type=AuthCredentialTypes.OAUTH2, |
| 1268 | + oauth2=OAuth2Auth( |
| 1269 | + client_id="mock_client_id", |
| 1270 | + client_secret="mock_client_secret", |
| 1271 | + ), |
| 1272 | + ) |
| 1273 | + stored_credential = raw_credential.model_copy(deep=True) |
| 1274 | + stored_credential.oauth2.access_token = "expired_token" |
| 1275 | + stored_credential.oauth2.refresh_token = "refresh_token" |
| 1276 | + stored_credential.oauth2.expires_at = int(time.time()) - 60 |
| 1277 | + manager = CredentialManager( |
| 1278 | + AuthConfig(auth_scheme=auth_scheme, raw_auth_credential=raw_credential) |
| 1279 | + ) |
| 1280 | + manager._load_existing_credential = AsyncMock( |
| 1281 | + return_value=stored_credential |
| 1282 | + ) |
| 1283 | + manager._save_credential = AsyncMock() |
| 1284 | + context = Mock() |
| 1285 | + |
| 1286 | + with patch.object( |
| 1287 | + oauth2_credential_util, "OAuth2Session", autospec=True |
| 1288 | + ) as mock_oauth2_session: |
| 1289 | + mock_client = mock_oauth2_session.return_value |
| 1290 | + mock_client.refresh_token.return_value = { |
| 1291 | + "access_token": "refreshed_token", |
| 1292 | + "refresh_token": "refresh_token", |
| 1293 | + "expires_at": int(time.time()) + 3600, |
| 1294 | + } |
| 1295 | + result = await manager.get_auth_credential(context) |
| 1296 | + |
| 1297 | + assert result.oauth2.access_token == "refreshed_token" |
| 1298 | + mock_client.refresh_token.assert_called_once() |
| 1299 | + assert ( |
| 1300 | + mock_client.refresh_token.call_args.kwargs["refresh_token"] |
| 1301 | + == "refresh_token" |
| 1302 | + ) |
| 1303 | + mock_client.fetch_token.assert_not_called() |
| 1304 | + |
| 1305 | + @pytest.mark.asyncio |
| 1306 | + async def test_get_auth_credential_keeps_stored_token_when_refetch_fails( |
| 1307 | + self, |
| 1308 | + ): |
| 1309 | + """A failed client credentials refetch returns the stored token unsaved.""" |
| 1310 | + auth_scheme = OAuth2( |
| 1311 | + flows=OAuthFlows( |
| 1312 | + clientCredentials=OAuthFlowClientCredentials( |
| 1313 | + tokenUrl="https://example.com/token", |
| 1314 | + scopes={}, |
| 1315 | + ) |
| 1316 | + ) |
| 1317 | + ) |
| 1318 | + raw_credential = AuthCredential( |
| 1319 | + auth_type=AuthCredentialTypes.OAUTH2, |
| 1320 | + oauth2=OAuth2Auth( |
| 1321 | + client_id="mock_client_id", |
| 1322 | + client_secret="mock_client_secret", |
| 1323 | + ), |
| 1324 | + ) |
| 1325 | + stored_credential = raw_credential.model_copy(deep=True) |
| 1326 | + stored_credential.oauth2.access_token = "expired_token" |
| 1327 | + stored_credential.oauth2.expires_at = int(time.time()) - 60 |
| 1328 | + manager = CredentialManager( |
| 1329 | + AuthConfig(auth_scheme=auth_scheme, raw_auth_credential=raw_credential) |
| 1330 | + ) |
| 1331 | + manager._load_existing_credential = AsyncMock( |
| 1332 | + return_value=stored_credential |
| 1333 | + ) |
| 1334 | + manager._save_credential = AsyncMock() |
| 1335 | + context = Mock() |
| 1336 | + |
| 1337 | + with patch.object( |
| 1338 | + oauth2_credential_util, "OAuth2Session", autospec=True |
| 1339 | + ) as mock_oauth2_session: |
| 1340 | + mock_client = mock_oauth2_session.return_value |
| 1341 | + mock_client.fetch_token.side_effect = OAuth2Error(error="invalid_client") |
| 1342 | + mock_client.refresh_token.side_effect = AssertionError( |
| 1343 | + "refresh_token() called without a refresh token" |
| 1344 | + ) |
| 1345 | + result = await manager.get_auth_credential(context) |
| 1346 | + |
| 1347 | + assert result is stored_credential |
| 1348 | + assert result.oauth2.access_token == "expired_token" |
| 1349 | + mock_client.fetch_token.assert_called_once() |
| 1350 | + manager._save_credential.assert_not_awaited() |
| 1351 | + assert raw_credential.oauth2.access_token is None |
| 1352 | + |
1183 | 1353 | @pytest.mark.asyncio |
1184 | 1354 | async def test_is_credential_ready_api_key(self): |
1185 | 1355 | """Test _is_credential_ready with API key credential.""" |
|
0 commit comments