diff --git a/tests/unit/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py b/tests/unit/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py index a438b62c94c..7913700c8a7 100644 --- a/tests/unit/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py +++ b/tests/unit/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py @@ -1911,6 +1911,87 @@ class TestCheckCachePagination: assert result is None assert self.mock_async_client.get.call_count == 1 + @pytest.mark.parametrize( + "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] + ) + @patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching") + def test_check_cache_pagination_max_pages_limit( + self, mock_get_token_url, custom_llm_provider + ): + mock_get_token_url.return_value = ("token", "https://test-url.com") + cache_key_to_find = "nonexistent_cache_key" + + def create_page_response(page_num): + response = MagicMock() + response.json.return_value = { + "cachedContents": [ + {"name": f"cache_{page_num}", "displayName": f"key_{page_num}"} + ], + "nextPageToken": f"token_page_{page_num + 1}", + } + return response + + self.mock_client.get.side_effect = [ + create_page_response(i) for i in range(100) + ] + + result = self.context_caching.check_cache( + cache_key=cache_key_to_find, + client=self.mock_client, + headers={"Authorization": "Bearer token"}, + api_key="test_key", + api_base=None, + logging_obj=self.mock_logging, + custom_llm_provider=custom_llm_provider, + vertex_project="test_project", + vertex_location="us-central1", + vertex_auth_header="Bearer test-token", + ) + + assert result is None + assert self.mock_client.get.call_count == 100 + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] + ) + @patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching") + async def test_async_check_cache_pagination_max_pages_limit( + self, mock_get_token_url, custom_llm_provider + ): + mock_get_token_url.return_value = ("token", "https://test-url.com") + cache_key_to_find = "nonexistent_cache_key" + + def create_page_response(page_num): + response = MagicMock() + response.json.return_value = { + "cachedContents": [ + {"name": f"cache_{page_num}", "displayName": f"key_{page_num}"} + ], + "nextPageToken": f"token_page_{page_num + 1}", + } + return response + + self.mock_async_client.get = AsyncMock( + side_effect=[create_page_response(i) for i in range(100)] + ) + + result = await self.context_caching.async_check_cache( + cache_key=cache_key_to_find, + client=self.mock_async_client, + headers={"Authorization": "Bearer token"}, + api_key="test_key", + api_base=None, + logging_obj=self.mock_logging, + custom_llm_provider=custom_llm_provider, + vertex_project="test_project", + vertex_location="us-central1", + vertex_auth_header="Bearer test-token", + ) + + assert result is None + assert self.mock_async_client.get.call_count == 100 + diff --git a/tests/unit/llms/vertex_ai/realtime/test_vertex_ai_realtime_transformation.py b/tests/unit/llms/vertex_ai/realtime/test_vertex_ai_realtime_transformation.py index a71ef21d997..14b3bdb48a1 100644 --- a/tests/unit/llms/vertex_ai/realtime/test_vertex_ai_realtime_transformation.py +++ b/tests/unit/llms/vertex_ai/realtime/test_vertex_ai_realtime_transformation.py @@ -387,6 +387,43 @@ def test_vertex_does_not_warn_when_dropping_non_guardrail_session_update(caplog) ) +@pytest.mark.asyncio +async def test_async_realtime_does_not_forward_client_query_params_to_vertex_backend( + monkeypatch, +): + import websockets + + from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler + + cfg = VertexAIRealtimeConfig( + access_token="tok", project="my-proj", location="us-central1" + ) + + captured = {} + + def fake_connect(url, *args, **kwargs): + captured["url"] = url + raise RuntimeError("stop before establishing the backend connection") + + monkeypatch.setattr(websockets, "connect", fake_connect) + + await BaseLLMHTTPHandler().async_realtime( + model="gemini-live-2.5-flash-preview-native-audio-09-2025", + websocket=AsyncMock(), + logging_obj=MagicMock(), + provider_config=cfg, + headers={}, + query_params={ + "model": "gemini-live-2.5-flash-preview-native-audio-09-2025", + "intent": "chat", + }, + ) + + assert "?" not in captured["url"] + assert "model=" not in captured["url"] + assert "intent=" not in captured["url"] + + def test_vertex_function_call_output_omits_id():