From bd242d0896f80289e06a0292e55247fc37c324c8 Mon Sep 17 00:00:00 2001 From: Adnaan Ali Date: Thu, 12 Feb 2026 04:33:30 +0000 Subject: [PATCH] test: add unit test for vertex_ai embedding extra_headers --- .../test_vertex_ai_embedding_headers.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) create mode 100644 tests/test_litellm/test_vertex_ai_embedding_headers.py diff --git a/tests/test_litellm/test_vertex_ai_embedding_headers.py b/tests/test_litellm/test_vertex_ai_embedding_headers.py new file mode 100644 index 00000000000..a46e1e76bbb --- /dev/null +++ b/tests/test_litellm/test_vertex_ai_embedding_headers.py @@ -0,0 +1,16 @@ +import pytest +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")