From 689cbaa6c1d768d2ef18633d01b422fbc969dab6 Mon Sep 17 00:00:00 2001 From: Chesars Date: Wed, 11 Mar 2026 14:25:08 -0300 Subject: [PATCH] fix(vertex-ai): update tests to match new get_vertex_region model_cost lookup - 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 --- .../vertex_ai_partner_models/main.py | 5 -- .../vertex_ai/test_vertex_ai_common_utils.py | 42 ++++++------ .../test_vertex_ai_qwen_global_endpoint.py | 66 +++++++++---------- 3 files changed, 54 insertions(+), 59 deletions(-) diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/main.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/main.py index 67fc6adc127..123d925f7c1 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/main.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/main.py @@ -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, diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py b/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py index 94323e06901..4dfde61509a 100644 --- a/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py +++ b/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py @@ -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(): diff --git a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/qwen/test_vertex_ai_qwen_global_endpoint.py b/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/qwen/test_vertex_ai_qwen_global_endpoint.py index dafd58b06d8..46f3f1e7817 100644 --- a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/qwen/test_vertex_ai_qwen_global_endpoint.py +++ b/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/qwen/test_vertex_ai_qwen_global_endpoint.py @@ -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, )