scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-07 08:26:10 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 89

fix: update test to reflect change

Browse source
This commit is contained in:
Krrish Dholakia 2025-05-19 20:26:48 -07:00
parent defd602513
commit edfd8a138d
1 changed files with 32 additions and 24 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

56
tests/litellm/llms/vertex_ai/test_vertex_llm_base.py
View file

@ -53,24 +53,23 @@ class TestVertexBase:
assert project == "project-1"
assert token == "fake-token-1"
# Test case 2: Prevent using credentials from different project
# Test case 2: Allow using credentials from different project
with patch.object(
vertex_base, "load_auth", return_value=(mock_creds, "project-1")
):
with pytest.raises(ValueError, match="Could not resolve project_id"):
if is_async:
result = await vertex_base._ensure_access_token_async(
credentials={"type": "service_account"},
project_id="different-project",
custom_llm_provider="vertex_ai",
)
else:
result = vertex_base._ensure_access_token(
credentials={"type": "service_account"},
project_id="different-project",
custom_llm_provider="vertex_ai",
)
print(f"result: {result}")
if is_async:
result = await vertex_base._ensure_access_token_async(
credentials={"type": "service_account"},
project_id="different-project",
custom_llm_provider="vertex_ai",
)
else:
result = vertex_base._ensure_access_token(
credentials={"type": "service_account"},
project_id="different-project",
custom_llm_provider="vertex_ai",
)
print(f"result: {result}")
@pytest.mark.parametrize("is_async", [True, False], ids=["async", "sync"])
@pytest.mark.asyncio
@ -189,7 +188,7 @@ class TestVertexBase:
"quota_project_id": "test-project",
"refresh_token": "fake-refresh-token",
"type": "authorized_user",
"universe_domain": "googleapis.com"
"universe_domain": "googleapis.com",
}
mock_creds = MagicMock()
@ -197,10 +196,12 @@ class TestVertexBase:
mock_creds.expired = False
mock_creds.quota_project_id = quota_project_id
with patch.object(
vertex_base, "_credentials_from_authorized_user", return_value=mock_creds
) as mock_credentials_from_authorized_user, patch.object(vertex_base, "refresh_auth") as mock_refresh:
) as mock_credentials_from_authorized_user, patch.object(
vertex_base, "refresh_auth"
) as mock_refresh:
def mock_refresh_impl(creds):
creds.token = "refreshed-token"
@ -209,11 +210,15 @@ class TestVertexBase:
# 1. Test that authorized_user-style credentials are correctly handled and uses quota_project_id
if is_async:
token, project = await vertex_base._ensure_access_token_async(
credentials=credentials, project_id=None, custom_llm_provider="vertex_ai"
credentials=credentials,
project_id=None,
custom_llm_provider="vertex_ai",
)
else:
token, project = vertex_base._ensure_access_token(
credentials=credentials, project_id=None, custom_llm_provider="vertex_ai"
credentials=credentials,
project_id=None,
custom_llm_provider="vertex_ai",
)
assert mock_credentials_from_authorized_user.called
@ -224,13 +229,16 @@ class TestVertexBase:
not_quota_project_id = "new-project"
if is_async:
token, project = await vertex_base._ensure_access_token_async(
credentials=credentials, project_id=not_quota_project_id, custom_llm_provider="vertex_ai"
credentials=credentials,
project_id=not_quota_project_id,
custom_llm_provider="vertex_ai",
)
else:
token, project = vertex_base._ensure_access_token(
credentials=credentials, project_id=not_quota_project_id, custom_llm_provider="vertex_ai"
credentials=credentials,
project_id=not_quota_project_id,
custom_llm_provider="vertex_ai",
)
assert token == "refreshed-token"
assert project == not_quota_project_id
assert project == not_quota_project_id