From 988ae7ca80b66dd9930c1a32ca399d18086561fa Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Thu, 3 Sep 2026 20:04:05 -0400 Subject: [PATCH] test(vertex_ai): assert the publisher-model batch payload instead of only mock calls --- tests/test_litellm/llms/vertex_ai/batches/test_handler.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/tests/test_litellm/llms/vertex_ai/batches/test_handler.py b/tests/test_litellm/llms/vertex_ai/batches/test_handler.py index 8db4775a2bc..d65ba92afd7 100644 --- a/tests/test_litellm/llms/vertex_ai/batches/test_handler.py +++ b/tests/test_litellm/llms/vertex_ai/batches/test_handler.py @@ -179,13 +179,14 @@ def test_create_batch_async_returns_coroutine_and_uses_async_client(): def test_create_batch_sync_does_not_resolve_publisher_models(): - """Publisher-model jobs must not incur the endpoint-resolution GET.""" + """Publisher-model jobs must not incur the endpoint-resolution GET, and the job model must + stay the publisher path untouched.""" h = _make_handler() client = MagicMock() client.post.return_value = _http_response() with patch(f"{HMOD}._get_httpx_client", return_value=client): - h.create_batch( + out = h.create_batch( _is_async=False, create_batch_data=CREATE_DATA, api_base=None, @@ -196,6 +197,9 @@ def test_create_batch_sync_does_not_resolve_publisher_models(): max_retries=None, ) + assert isinstance(out, LiteLLMBatch) + sent = json.loads(client.post.call_args.kwargs["data"]) + assert sent["model"] == "publishers/google/models/gemini-1.5-flash-001" client.get.assert_not_called()