From eddae66697adfa322a281939cc14d418278cfefb Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Fri, 24 Apr 2026 22:00:53 -0700 Subject: [PATCH] [Test] mock vertex schema tests so they don't make real network calls test_vertex_schema_test and test_gemini_nullable_object_tool_schema_httpx were calling vertex_ai/gemini-2.5-flash for real. When Vertex was slow, the calls hung past pytest's 300s timeout, pytest-timeout fired os._exit(1), and the xdist worker disappeared mid-run as "Not properly terminated". The actual purpose of these tests is to exercise the schema-conversion code path (anyOf with null types, nullable object tool params). Mocking the HTTP layer and the auth token retrieval makes them hermetic and removes the worker-crash failure mode entirely. --- .../test_amazing_vertex_completion.py | 142 +++++++++++++++--- 1 file changed, 123 insertions(+), 19 deletions(-) diff --git a/tests/local_testing/test_amazing_vertex_completion.py b/tests/local_testing/test_amazing_vertex_completion.py index 5cb1e2c78e8..0cd8bb63581 100644 --- a/tests/local_testing/test_amazing_vertex_completion.py +++ b/tests/local_testing/test_amazing_vertex_completion.py @@ -3687,10 +3687,9 @@ def test_vertex_ai_llama_tool_calling(): def test_vertex_schema_test(): - load_vertex_ai_credentials() - - def tool_call(text: str | None) -> str: - return text or "No text provided" + """Verify Vertex schema conversion accepts anyOf with null types in tool params.""" + from litellm.llms.custom_httpx.http_handler import HTTPHandler + from litellm.llms.vertex_ai.vertex_llm_base import VertexBase tool = { "type": "function", @@ -3714,22 +3713,65 @@ def test_vertex_schema_test(): }, } - response = litellm.completion( - model="vertex_ai/gemini-2.5-flash", - messages=[{"role": "user", "content": "call the tool"}], - tools=[tool], - tool_choice="required", - ) + client = HTTPHandler() + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.headers = {"Content-Type": "application/json"} + mock_response.json.return_value = { + "candidates": [ + { + "content": { + "role": "model", + "parts": [ + { + "functionCall": { + "name": "git_create_branch", + "args": { + "repo_path": "/tmp/repo", + "branch_name": "feature", + }, + } + } + ], + }, + "finishReason": "STOP", + } + ], + "usageMetadata": { + "promptTokenCount": 10, + "candidatesTokenCount": 5, + "totalTokenCount": 15, + }, + } - print(response) + with ( + patch.object(client, "post", return_value=mock_response), + patch.object( + VertexBase, + "_ensure_access_token", + return_value=("fake-token", "fake-project"), + ), + ): + response = litellm.completion( + model="vertex_ai/gemini-2.5-flash", + messages=[{"role": "user", "content": "call the tool"}], + tools=[tool], + tool_choice="required", + client=client, + ) + + assert response.choices[0].message.tool_calls is not None + assert ( + response.choices[0].message.tool_calls[0].function.name == "git_create_branch" + ) def test_gemini_nullable_object_tool_schema_httpx(): """ Ensure nullable object tool params preserve nested properties in Vertex schema conversion. """ - load_vertex_ai_credentials() - litellm._turn_on_debug() + from litellm.llms.custom_httpx.http_handler import HTTPHandler + from litellm.llms.vertex_ai.vertex_llm_base import VertexBase tools = [ { @@ -3769,14 +3811,76 @@ def test_gemini_nullable_object_tool_schema_httpx(): } ] - response = litellm.completion( - model="vertex_ai/gemini-2.5-flash", - messages=[{"role": "user", "content": "call the tool"}], - tools=tools, - tool_choice="required", + client = HTTPHandler() + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.headers = {"Content-Type": "application/json"} + mock_response.json.return_value = { + "candidates": [ + { + "content": { + "role": "model", + "parts": [ + { + "functionCall": { + "name": "create_support_ticket", + "args": { + "ticket_id": "T-1", + "customer_context": { + "user_id": "u-1", + "plan": "pro", + }, + }, + } + } + ], + }, + "finishReason": "STOP", + } + ], + "usageMetadata": { + "promptTokenCount": 10, + "candidatesTokenCount": 5, + "totalTokenCount": 15, + }, + } + + with ( + patch.object(client, "post", return_value=mock_response) as mock_post, + patch.object( + VertexBase, + "_ensure_access_token", + return_value=("fake-token", "fake-project"), + ), + ): + response = litellm.completion( + model="vertex_ai/gemini-2.5-flash", + messages=[{"role": "user", "content": "call the tool"}], + tools=tools, + tool_choice="required", + client=client, + ) + + assert response.choices[0].message.tool_calls is not None + assert ( + response.choices[0].message.tool_calls[0].function.name + == "create_support_ticket" ) - print(response) + # Verify schema conversion preserved the nested customer_context fields + # somewhere in the converted tool definition (the converter may rewrap + # nullable objects under anyOf, so check the serialized form rather than + # a fixed shape). + sent_body = mock_post.call_args.kwargs.get( + "json" + ) or mock_post.call_args.kwargs.get("data") + assert sent_body is not None, "expected request body to be sent" + if isinstance(sent_body, str): + sent_body = json.loads(sent_body) + sent_tool_str = json.dumps(sent_body["tools"][0]["function_declarations"][0]) + assert "customer_context" in sent_tool_str + assert "user_id" in sent_tool_str + assert "plan" in sent_tool_str def test_vertex_ai_response_id():