From 85d4000af6d3932c43e52d5542d3f18e4e45bee0 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Wed, 26 Nov 2025 11:37:55 -0800 Subject: [PATCH] test_vertex_ai_partner_models_token_counting_endpoint --- tests/proxy_unit_tests/test_proxy_token_counter.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/tests/proxy_unit_tests/test_proxy_token_counter.py b/tests/proxy_unit_tests/test_proxy_token_counter.py index eb364ebf5b..a8486358f7 100644 --- a/tests/proxy_unit_tests/test_proxy_token_counter.py +++ b/tests/proxy_unit_tests/test_proxy_token_counter.py @@ -821,18 +821,22 @@ async def test_vertex_ai_anthropic_token_counting(): assert response.original_response["input_tokens"] == 15 @pytest.mark.parametrize("vertex_location", ["global", "us-central1"]) -def test_vertex_ai_gemini_token_counting_endpoint(vertex_location): +def test_vertex_ai_partner_models_token_counting_endpoint(vertex_location): + """ + Test that the VertexAIPartnerModelsTokenCounter builds the correct endpoint URL + for different vertex locations, including the special 'global' location. + """ from litellm.llms.vertex_ai.vertex_ai_partner_models.count_tokens.handler import ( VertexAIPartnerModelsTokenCounter, ) endpoint = VertexAIPartnerModelsTokenCounter()._build_count_tokens_endpoint( - model="gemini-2.5-pro", + model="claude-3-5-sonnet-20241022", project_id="test-project", vertex_location=vertex_location, api_base=None, ) if vertex_location == "global": - assert endpoint == "https://aiplatform.googleapis.com" + assert endpoint.startswith("https://aiplatform.googleapis.com") else: - assert endpoint == f"https://{vertex_location}-aiplatform.googleapis.com" \ No newline at end of file + assert endpoint.startswith(f"https://{vertex_location}-aiplatform.googleapis.com") \ No newline at end of file