scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-09 22:31:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 86

fix(vertex-ai): update tests to match new get_vertex_region model_cost lookup

Browse source 
- Remove redundant get_vertex_region() call in partner models main.py
  (already called inside get_complete_vertex_url)
- Rewrite test mocks to use patch.dict(litellm.model_cost) instead of
  patching the removed is_global_only_vertex_model symbol
- Align test assertions with new behavior: user-specified region is
  preserved (not overridden) for global-only models
This commit is contained in:
Chesars 2026-03-11 14:25:08 -03:00
parent 8d4e98faa1
commit 689cbaa6c1
3 changed files with 54 additions and 59 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

5
litellm/llms/vertex_ai/vertex_ai_partner_models/main.py
View file

@ -161,11 +161,6 @@ class VertexAIPartnerModels(VertexBase):
else:
raise ValueError(f"Unknown partner model: {model}")
# Resolve vertex_location based on model's supported_regions
vertex_location = self.get_vertex_region(
vertex_region=vertex_location, model=model
)
api_base = self.get_complete_vertex_url(
custom_api_base=api_base,
vertex_location=vertex_location,

42
tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py
View file

@ -588,39 +588,43 @@ def test_is_global_only_vertex_model(supported_regions, expected_result):
@pytest.mark.parametrize(
"model_is_global_only, vertex_region, expected_region",
"model_cost_entry, vertex_region, expected_region",
[
(True, None, "global"), # Global-only model with no region specified
(True, "us-central1", "global"), # Global-only model overrides specified region
(True, "europe-west1", "global"), # Global-only model overrides any region
(False, None, "us-central1"), # Non-global model defaults to us-central1
(
False,
"europe-west1",
"europe-west1",
), # Non-global model uses specified region
(False, "us-east1", "us-east1"), # Non-global model uses specified region
# Model with supported_regions=["global"], no user region -> use "global"
({"supported_regions": ["global"]}, None, "global"),
# Model with supported_regions=["global"], user specifies region -> trust user
({"supported_regions": ["global"]}, "us-central1", "us-central1"),
# Model with supported_regions=["global"], user specifies region -> trust user
({"supported_regions": ["global"]}, "europe-west1", "europe-west1"),
# Model with supported_regions=["us-west2"], no user region -> use "us-west2"
({"supported_regions": ["us-west2"]}, None, "us-west2"),
# No model_cost entry, no user region -> default us-central1
({}, None, "us-central1"),
# No model_cost entry, user specifies region -> use specified region
({}, "europe-west1", "europe-west1"),
# No model_cost entry, user specifies region -> use specified region
({}, "us-east1", "us-east1"),
],
)
def test_get_vertex_region_global_only_model(
model_is_global_only, vertex_region, expected_region
model_cost_entry, vertex_region, expected_region
):
"""Test get_vertex_region ensures global-only models default to 'global' region"""
"""Test get_vertex_region resolves region from model_cost supported_regions"""
import litellm
from litellm.llms.vertex_ai.vertex_llm_base import VertexBase
vertex_base = VertexBase()
with patch(
"litellm.llms.vertex_ai.vertex_llm_base.is_global_only_vertex_model"
) as mock_is_global_only:
mock_is_global_only.return_value = model_is_global_only
with patch.dict(
litellm.model_cost,
{"vertex_ai/test-model": model_cost_entry},
clear=False,
):
result = vertex_base.get_vertex_region(
vertex_region=vertex_region, model="test-model"
)
assert result == expected_region
mock_is_global_only.assert_called_once_with("test-model")
def test_vertex_filter_format_uri():

66
tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/qwen/test_vertex_ai_qwen_global_endpoint.py
View file

@ -81,47 +81,46 @@ class TestQwenGlobalOnlyDetection:
class TestVertexBaseGetVertexRegion:
"""Test the get_vertex_region method."""
"""Test the get_vertex_region method using model_cost lookup."""
def test_global_only_model_returns_global(self):
"""Test that global-only models return 'global' regardless of input."""
def test_global_model_no_user_region_returns_global(self):
"""Test that global-only models return 'global' when user doesn't specify region."""
vertex_base = VertexBase()
with patch(
"litellm.llms.vertex_ai.vertex_llm_base.is_global_only_vertex_model",
return_value=True,
):
result = vertex_base.get_vertex_region(
vertex_region="us-central1",
model="vertex_ai/qwen/qwen3-next-80b-a3b-instruct-maas",
)
assert result == "global"
def test_global_only_model_with_none_returns_global(self):
"""Test that global-only models return 'global' even with None input."""
vertex_base = VertexBase()
with patch(
"litellm.llms.vertex_ai.vertex_llm_base.is_global_only_vertex_model",
return_value=True,
with patch.dict(
litellm.model_cost,
{"vertex_ai/qwen/qwen3-next-80b-a3b-instruct-maas": {"supported_regions": ["global"]}},
clear=False,
):
result = vertex_base.get_vertex_region(
vertex_region=None,
model="vertex_ai/qwen/qwen3-next-80b-a3b-instruct-maas",
model="qwen/qwen3-next-80b-a3b-instruct-maas",
)
assert result == "global"
def test_global_model_with_user_region_trusts_user(self):
"""Test that user-specified region is preserved even for global-only models."""
vertex_base = VertexBase()
with patch.dict(
litellm.model_cost,
{"vertex_ai/qwen/qwen3-next-80b-a3b-instruct-maas": {"supported_regions": ["global"]}},
clear=False,
):
result = vertex_base.get_vertex_region(
vertex_region="us-central1",
model="qwen/qwen3-next-80b-a3b-instruct-maas",
)
assert result == "us-central1"
def test_non_global_model_uses_provided_region(self):
"""Test that non-global models use the provided region."""
vertex_base = VertexBase()
with patch(
"litellm.llms.vertex_ai.vertex_llm_base.is_global_only_vertex_model",
return_value=False,
):
with patch.dict(litellm.model_cost, {}, clear=False):
result = vertex_base.get_vertex_region(
vertex_region="europe-west1",
model="vertex_ai/gemini-1.5-pro",
model="gemini-1.5-pro",
)
assert result == "europe-west1"
@ -129,13 +128,10 @@ class TestVertexBaseGetVertexRegion:
"""Test that non-global models with None region fallback to us-central1."""
vertex_base = VertexBase()
with patch(
"litellm.llms.vertex_ai.vertex_llm_base.is_global_only_vertex_model",
return_value=False,
):
with patch.dict(litellm.model_cost, {}, clear=False):
result = vertex_base.get_vertex_region(
vertex_region=None,
model="vertex_ai/gemini-1.5-pro",
model="unknown-model-xyz",
)
assert result == "us-central1"
@ -217,15 +213,15 @@ async def test_vertex_ai_qwen_global_endpoint_url():
client, "post", side_effect=mock_post_func
) as mock_post, patch.object(
VertexLLM, "_ensure_access_token", return_value=("fake-token", "test-project")
), patch(
"litellm.llms.vertex_ai.vertex_llm_base.is_global_only_vertex_model",
return_value=True,
), patch.dict(
litellm.model_cost,
{"vertex_ai/qwen/qwen3-next-80b-a3b-instruct-maas": {"supported_regions": ["global"]}},
clear=False,
):
response = await litellm.acompletion(
model="vertex_ai/qwen/qwen3-next-80b-a3b-instruct-maas",
messages=[{"role": "user", "content": "Hello"}],
vertex_ai_project="test-project",
vertex_ai_location="us-central1",
client=client,
)