scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-15 09:12:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 311

test(guardrails): fold the #45543 cases into the mapped test file

Browse source 
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.
This commit is contained in:
Vivek1898 2026-10-18 14:39:52 +05:30 • committed by Vivek1898
parent aa3faa947b
commit f90ca3ff17
2 changed files with 72 additions and 96 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

96
tests/unit/proxy/guardrails/guardrail_hooks/test_unified_guardrail_metadata_bucket.py
View file

@ -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

72
tests/unit/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py
View file

@ -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