From f90ca3ff17fc579c44355e034a5070b93da2f800 Mon Sep 17 00:00:00 2001 From: Vivek1898 <72497717+Vivek1898@users.noreply.github.com> Date: Sun, 18 Oct 2026 14:39:52 +0530 Subject: [PATCH] test(guardrails): fold the #45543 cases into the mapped test file AGENTS.md asks bug fixes to extend the existing mapped test file rather than add one. The mapped file is also the better home on its merits: it already carries TestStreamingGuardrailInformationBucket, which pins the same bucket-flip shape from #38722, so these five cases belong beside it instead of in a file a reader would not think to open. Moved as TestEnsureLitellmMetadataBucket, following that file's conventions: the regression context in the class docstring, the user_api_key_dict fixture as a method, and the module's existing unified_module alias instead of a fresh private import. The `# type: ignore[arg-type]` on the None case is replaced with `# pyright: ignore[reportArgumentType]`, trailing on the call so it attaches to the right node. pyrightconfig.json sets enableTypeIgnoreComments to false, so the old form suppressed nothing and named no rule. --- .../test_unified_guardrail_metadata_bucket.py | 96 ------------------- .../test_unified_guardrail.py | 72 ++++++++++++++ 2 files changed, 72 insertions(+), 96 deletions(-) delete mode 100644 tests/unit/proxy/guardrails/guardrail_hooks/test_unified_guardrail_metadata_bucket.py diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_unified_guardrail_metadata_bucket.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_unified_guardrail_metadata_bucket.py deleted file mode 100644 index 7803d379d74..00000000000 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_unified_guardrail_metadata_bucket.py +++ /dev/null @@ -1,96 +0,0 @@ -"""Regression tests for #45543. - -A pre_call guardrail used to write ``data["litellm_metadata"]`` directly. Every -reader resolves the proxy-internal bucket through -``get_metadata_variable_name_from_kwargs``, which answers ``litellm_metadata`` as -soon as that key merely exists, so creating it on a chat-completions request (whose -bucket is ``metadata``) moved every later read onto a dict holding only the -``user_api_key_*`` keys. Tags and router fields stayed behind in ``metadata``: -tag-based routing found no tags and picked a deployment of another model, and -SpendLogs lost ``model_group`` and ``model_id``. Nothing raised. -""" - -from typing import Final - -import pytest - -from litellm.litellm_core_utils.core_helpers import ( - get_metadata_variable_name_from_kwargs, -) -from litellm.proxy._types import UserAPIKeyAuth -from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import ( - _ensure_litellm_metadata, -) - - -@pytest.fixture -def user_api_key_dict() -> UserAPIKeyAuth: - return UserAPIKeyAuth(api_key="sk-test-45543", user_id="user-1", team_id="team-1") - - -def test_chat_completions_request_keeps_its_metadata_bucket(user_api_key_dict: UserAPIKeyAuth) -> None: - """The reported shape: proxy state already lives in ``metadata``.""" - data: Final = { - "model": "all-models", - "metadata": {"tags": ["model:glm-5"], "model_group": "all-models"}, - } - - _ensure_litellm_metadata(data, user_api_key_dict) - - # The bug was creating this key at all; its presence is what flips the readers. - assert "litellm_metadata" not in data - assert get_metadata_variable_name_from_kwargs(data) == "metadata" - # The fields the routing and spend-logging readers need are still where they look. - assert data["metadata"]["tags"] == ["model:glm-5"] - assert data["metadata"]["model_group"] == "all-models" - # And the guardrail still gets the key metadata it asked for. - assert data["metadata"]["user_api_key_team_id"] == "team-1" - assert data["metadata"]["user_api_key_user_id"] == "user-1" - - -def test_existing_metadata_values_are_not_overwritten(user_api_key_dict: UserAPIKeyAuth) -> None: - """Merging is ``setdefault``: a value already resolved upstream wins.""" - data: Final = { - "model": "all-models", - "metadata": {"user_api_key_user_id": "resolved-upstream"}, - } - - _ensure_litellm_metadata(data, user_api_key_dict) - - assert data["metadata"]["user_api_key_user_id"] == "resolved-upstream" - assert data["metadata"]["user_api_key_team_id"] == "team-1" - - -def test_litellm_metadata_route_still_uses_litellm_metadata(user_api_key_dict: UserAPIKeyAuth) -> None: - """Batch and file routes really do keep proxy state in ``litellm_metadata``. - - The fix must not flip those the other way, so the bucket is resolved rather - than hardcoded in either direction. - """ - data: Final = {"model": "m", "litellm_metadata": {"tags": ["keep-me"]}} - - _ensure_litellm_metadata(data, user_api_key_dict) - - assert "metadata" not in data - assert get_metadata_variable_name_from_kwargs(data) == "litellm_metadata" - assert data["litellm_metadata"]["tags"] == ["keep-me"] - assert data["litellm_metadata"]["user_api_key_team_id"] == "team-1" - - -def test_request_without_any_metadata_gets_the_default_bucket(user_api_key_dict: UserAPIKeyAuth) -> None: - data: Final = {"model": "m"} - - _ensure_litellm_metadata(data, user_api_key_dict) - - assert "litellm_metadata" not in data - assert data["metadata"]["user_api_key_team_id"] == "team-1" - - -def test_no_user_api_key_metadata_creates_no_bucket() -> None: - """``None`` yields no metadata, so nothing should be written at all.""" - data: Final = {"model": "m"} - - _ensure_litellm_metadata(data, None) # type: ignore[arg-type] # reason: the hook is reached with no key in tests and on unauthenticated routes - - assert "litellm_metadata" not in data - assert "metadata" not in data diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py b/tests/unit/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py index 84f4d7a4173..b1796330dd7 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py @@ -21,6 +21,9 @@ from litellm.integrations.custom_guardrail import ( log_guardrail_information, ) from litellm.litellm_core_utils.api_route_to_call_types import get_call_types_for_route +from litellm.litellm_core_utils.core_helpers import ( + get_metadata_variable_name_from_kwargs, +) from litellm.llms import discover_guardrail_translation_mappings, load_guardrail_translation_mappings from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation from litellm.llms.base_llm.guardrail_translation.utils import ( @@ -3291,3 +3294,72 @@ async def test_pass_through_guardrail_text_is_the_client_body_without_the_inject texts: Final = [call["inputs"]["texts"] for call in guardrail.apply_calls] assert [[json.loads(text) for text in call_texts] for call_texts in texts] == [[body]], texts assert data["litellm_metadata"]["user_api_key_user_email"] == "someone@example.com", data["litellm_metadata"] + + +class TestEnsureLitellmMetadataBucket: + """``_ensure_litellm_metadata`` must merge into the bucket the request + already resolves to, not plant ``litellm_metadata``. Creating that key on a + chat-completions request flips every reader onto a dict holding only the + ``user_api_key_*`` keys, so tags and router fields are left unread: tag + routing picks a deployment of another model and SpendLogs lose + ``model_group``/``model_id``, without raising (#45543). Same bucket-flip + shape as TestStreamingGuardrailInformationBucket above.""" + + @pytest.fixture + def user_api_key_dict(self) -> UserAPIKeyAuth: + return UserAPIKeyAuth(api_key="sk-test-45543", user_id="user-1", team_id="team-1") + + def test_chat_completions_request_keeps_its_metadata_bucket(self, user_api_key_dict: UserAPIKeyAuth) -> None: + data: Final = { + "model": "all-models", + "metadata": {"tags": ["model:glm-5"], "model_group": "all-models"}, + } + + unified_module._ensure_litellm_metadata(data, user_api_key_dict) + + # Creating the key at all is the bug; its presence is what flips the readers. + assert "litellm_metadata" not in data + assert get_metadata_variable_name_from_kwargs(data) == "metadata" + assert data["metadata"]["tags"] == ["model:glm-5"] + assert data["metadata"]["model_group"] == "all-models" + assert data["metadata"]["user_api_key_team_id"] == "team-1" + assert data["metadata"]["user_api_key_user_id"] == "user-1" + + def test_existing_metadata_values_are_not_overwritten(self, user_api_key_dict: UserAPIKeyAuth) -> None: + data: Final = { + "model": "all-models", + "metadata": {"user_api_key_user_id": "resolved-upstream"}, + } + + unified_module._ensure_litellm_metadata(data, user_api_key_dict) + + assert data["metadata"]["user_api_key_user_id"] == "resolved-upstream" + assert data["metadata"]["user_api_key_team_id"] == "team-1" + + def test_litellm_metadata_route_still_uses_litellm_metadata(self, user_api_key_dict: UserAPIKeyAuth) -> None: + """Batch and file routes do keep proxy state in ``litellm_metadata``, so + the bucket is resolved rather than hardcoded in either direction.""" + data: Final = {"model": "m", "litellm_metadata": {"tags": ["keep-me"]}} + + unified_module._ensure_litellm_metadata(data, user_api_key_dict) + + assert "metadata" not in data + assert get_metadata_variable_name_from_kwargs(data) == "litellm_metadata" + assert data["litellm_metadata"]["tags"] == ["keep-me"] + assert data["litellm_metadata"]["user_api_key_team_id"] == "team-1" + + def test_request_without_any_metadata_gets_the_default_bucket(self, user_api_key_dict: UserAPIKeyAuth) -> None: + data: Final = {"model": "m"} + + unified_module._ensure_litellm_metadata(data, user_api_key_dict) + + assert "litellm_metadata" not in data + assert data["metadata"]["user_api_key_team_id"] == "team-1" + + def test_no_user_api_key_metadata_creates_no_bucket(self) -> None: + data: Final = {"model": "m"} + + unified_module._ensure_litellm_metadata(data, None) # pyright: ignore[reportArgumentType] # the hook is reached with no key on unauthenticated routes + + assert "litellm_metadata" not in data + assert "metadata" not in data