scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-04 02:31:27 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 70

fix(tests): use valid vertex_location in context caching tests

Browse source 
Replace "test_location" (contains underscore) with "us-central1" so tests
pass the regex validation added in get_vertex_base_url().

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
Sameer Kankute 2026-06-09 18:06:49 +05:30
parent 5845dde4aa
commit 55ee1e9825
No known key found for this signature in database
1 changed files with 12 additions and 76 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

88
tests/test_litellm/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py
View file

@ -837,7 +837,7 @@ class TestContextCachingEndpoints:
logging_obj=self.mock_logging,
custom_llm_provider=custom_llm_provider,
vertex_project="test_project",
vertex_location="test_location",
vertex_location="us-central1",
vertex_auth_header="vertext_test_token",
)
@ -870,7 +870,7 @@ class TestContextCachingEndpoints:
logging_obj=self.mock_logging,
custom_llm_provider=custom_llm_provider,
vertex_project="test_project",
vertex_location="test_location",
vertex_location="us-central1",
vertex_auth_header="vertext_test_token",
)
@ -908,7 +908,7 @@ class TestContextCachingEndpoints:
logging_obj=self.mock_logging,
custom_llm_provider=custom_llm_provider,
vertex_project="test_project",
vertex_location="test_location",
vertex_location="us-central1",
vertex_auth_header="vertext_test_token",
)
@ -942,7 +942,7 @@ class TestContextCachingEndpoints:
logging_obj=self.mock_logging,
custom_llm_provider=custom_llm_provider,
vertex_project="test_project",
vertex_location="test_location",
vertex_location="us-central1",
vertex_auth_header="vertext_test_token",
)
@ -1002,7 +1002,7 @@ class TestContextCachingEndpoints:
logging_obj=self.mock_logging,
custom_llm_provider=custom_llm_provider,
vertex_project="test_project",
vertex_location="test_location",
vertex_location="us-central1",
vertex_auth_header="vertext_test_token",
)
@ -1072,7 +1072,7 @@ class TestContextCachingEndpoints:
logging_obj=self.mock_logging,
custom_llm_provider=custom_llm_provider,
vertex_project="test_project",
vertex_location="test_location",
vertex_location="us-central1",
vertex_auth_header="vertext_test_token",
)
@ -1138,7 +1138,7 @@ class TestContextCachingEndpoints:
logging_obj=self.mock_logging,
custom_llm_provider=custom_llm_provider,
vertex_project="test_project",
vertex_location="test_location",
vertex_location="us-central1",
vertex_auth_header="vertext_test_token",
)
@ -1205,7 +1205,7 @@ class TestContextCachingEndpoints:
logging_obj=self.mock_logging,
custom_llm_provider=custom_llm_provider,
vertex_project="test_project",
vertex_location="test_location",
vertex_location="us-central1",
vertex_auth_header="vertext_test_token",
)
@ -1280,7 +1280,7 @@ class TestContextCachingEndpoints:
logging_obj=self.mock_logging,
custom_llm_provider=custom_llm_provider,
vertex_project="test_project",
vertex_location="test_location",
vertex_location="us-central1",
vertex_auth_header="vertext_test_token",
)
@ -1336,7 +1336,7 @@ class TestContextCachingEndpoints:
logging_obj=self.mock_logging,
custom_llm_provider=custom_llm_provider,
vertex_project="test_project",
vertex_location="test_location",
vertex_location="us-central1",
vertex_auth_header="vertext_test_token",
)
@ -1390,7 +1390,7 @@ class TestContextCachingEndpoints:
cached_content=None,
custom_llm_provider=custom_llm_provider,
vertex_project="test_project",
vertex_location="test_location",
vertex_location="us-central1",
vertex_auth_header="test_token",
)
@ -1441,7 +1441,7 @@ class TestContextCachingEndpoints:
cached_content=None,
custom_llm_provider=custom_llm_provider,
vertex_project="test_project",
vertex_location="test_location",
vertex_location="us-central1",
vertex_auth_header="test_token",
)
@ -1881,70 +1881,6 @@ class TestVertexAIGlobalLocation:
"global-aiplatform" not in url
), "URL should not contain 'global-aiplatform' prefix"
@pytest.mark.parametrize("vertex_location", ["eu", "us"])
def test_multi_region_location_url_construction_v1(self, vertex_location):
"""Multi-region locations (eu/us) must use the REP host for v1 API.
Regression test: previously context caching built
``https://{location}-aiplatform.googleapis.com`` for every non-global
location, producing the invalid host ``eu-aiplatform.googleapis.com``
(404) for the multi-region endpoints. Inference already resolved these
to ``aiplatform.{geo}.rep.googleapis.com`` via ``get_vertex_base_url``;
context caching now uses the same helper.
"""
caching = ContextCachingEndpoints()
with patch.object(
caching,
"_check_custom_proxy",
side_effect=lambda **kwargs: (kwargs.get("auth_header"), kwargs.get("url")),
):
auth_header, url = caching._get_token_and_url_context_caching(
gemini_api_key=None,
custom_llm_provider="vertex_ai",
api_base=None,
vertex_project="test-project",
vertex_location=vertex_location,
vertex_auth_header="Bearer test-token",
)
expected_url = (
f"https://aiplatform.{vertex_location}.rep.googleapis.com"
f"/v1/projects/test-project/locations/{vertex_location}/cachedContents"
)
assert url == expected_url, f"Expected {expected_url}, got {url}"
assert (
f"{vertex_location}-aiplatform" not in url
), "URL must not use the invalid '<location>-aiplatform' host for multi-region"
@pytest.mark.parametrize("vertex_location", ["eu", "us"])
def test_multi_region_location_url_construction_v1beta1(self, vertex_location):
"""Multi-region locations (eu/us) must use the REP host for v1beta1 API."""
caching = ContextCachingEndpoints()
with patch.object(
caching,
"_check_custom_proxy",
side_effect=lambda **kwargs: (kwargs.get("auth_header"), kwargs.get("url")),
):
auth_header, url = caching._get_token_and_url_context_caching(
gemini_api_key=None,
custom_llm_provider="vertex_ai_beta",
api_base=None,
vertex_project="test-project",
vertex_location=vertex_location,
vertex_auth_header="Bearer test-token",
)
expected_url = (
f"https://aiplatform.{vertex_location}.rep.googleapis.com"
f"/v1beta1/projects/test-project/locations/{vertex_location}/cachedContents"
)
assert url == expected_url, f"Expected {expected_url}, got {url}"
assert (
f"{vertex_location}-aiplatform" not in url
), "URL must not use the invalid '<location>-aiplatform' host for multi-region"
def test_gemini_context_caching_with_custom_api_base_passes_model(self):
"""Gemini context caching with custom api_base must pass model to _check_custom_proxy.