From 2f0acfbea18a59d8b6d11d6b0bf51901f9719d44 Mon Sep 17 00:00:00 2001 From: Adnaan Ali Date: Fri, 13 Feb 2026 04:48:09 +0000 Subject: [PATCH] test: refactor vertex_ai embedding test to use mocks --- .../test_vertex_ai_embedding_headers.py | 40 +++++++++++++------ 1 file changed, 28 insertions(+), 12 deletions(-) diff --git a/tests/test_litellm/test_vertex_ai_embedding_headers.py b/tests/test_litellm/test_vertex_ai_embedding_headers.py index a46e1e76bbb..5300df1b029 100644 --- a/tests/test_litellm/test_vertex_ai_embedding_headers.py +++ b/tests/test_litellm/test_vertex_ai_embedding_headers.py @@ -1,16 +1,32 @@ import pytest +from unittest.mock import patch, MagicMock from litellm import embedding def test_vertex_ai_embedding_extra_headers(): - # Test that extra_headers are passed without crashing - try: - response = embedding( - model="vertex_ai/text-embedding-004", - input=["hello"], - extra_headers={"X-Custom-Header": "test-value"} - ) - except Exception as e: - # We expect a 401/404 if no real creds, - # but we are checking that it doesn't fail with a TypeError - if "extra_headers" in str(e): - pytest.fail("extra_headers not accepted by vertex_ai embedding") + """ + Test that extra_headers are correctly forwarded to the + vertex_embedding.embedding function. + """ + # We patch the exact location where main.py calls the vertex provider + with patch("litellm.main.vertex_embedding.embedding") as mock_vertex_embedding: + # Mock a successful return so the call doesn't fail + mock_vertex_embedding.return_value = MagicMock() + + # Trigger the embedding call + try: + embedding( + model="vertex_ai/text-embedding-004", + input=["hello"], + extra_headers={"X-Custom-Header": "test-value"}, + ) + except Exception: + # We don't care about subsequent errors, only the forwarding + pass + + # VERIFICATION: This is the important part + mock_vertex_embedding.assert_called_once() + call_kwargs = mock_vertex_embedding.call_args.kwargs + + # Check that the headers we passed actually reached the provider + assert call_kwargs.get("extra_headers") == {"X-Custom-Header": "test-value"} + print("\n✅ Success: extra_headers correctly forwarded to Vertex AI provider!")