scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-06 02:48:13 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 103

test(proxy): cover _enforce_guardrail_added_tag_budgets directly
Some checks failed
LiteLLM Rust / rust-lint (push) Waiting to run
Details
LiteLLM Rust / rust-test (push) Waiting to run
Details
Terraform Provider / gofmt, vet, build, test (push) Waiting to run
Details
Terraform Provider / Provider endpoints vs proxy OpenAPI schema (push) Waiting to run
Details
ai-gateway image / ai-gateway release image (push) Has been cancelled
Details

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Devin AI 2026-09-12 08:29:32 +00:00 • committed by yassin
parent e79cf49722
commit 6a5127076d
1 changed files with 41 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

41
tests/test_litellm/proxy/test_common_request_processing.py
View file

@ -485,6 +485,47 @@ class TestProxyBaseLLMRequestProcessing:
assert returned_data["metadata"]["tags"] == ["existing-tag"]
tag_budget_check.assert_not_awaited()
@pytest.mark.asyncio
async def test_enforce_guardrail_added_tag_budgets_checks_only_added_tags(self, monkeypatch):
"""
The helper must forward only tags that pre_call_hook added, so tags the
caller already sent (and auth already checked) are not budget-checked twice.
"""
from litellm.proxy.common_request_processing import _enforce_guardrail_added_tag_budgets
tag_budget_check = AsyncMock()
monkeypatch.setattr(
litellm.proxy.common_request_processing,
"_tag_max_budget_check",
tag_budget_check,
)
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", MagicMock())
monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", MagicMock())
mock_user_api_key_dict = MagicMock(spec=UserAPIKeyAuth)
mock_proxy_logging_obj = MagicMock(spec=ProxyLogging)
await _enforce_guardrail_added_tag_budgets(
data={"metadata": {"tags": ["existing-tag", "guardrail-tag"]}},
tags_before_guardrails=frozenset({"existing-tag"}),
user_api_key_dict=mock_user_api_key_dict,
proxy_logging_obj=mock_proxy_logging_obj,
)
tag_budget_check.assert_awaited_once()
_, call_kwargs = tag_budget_check.call_args
assert call_kwargs["request_body"] == {"metadata": {"tags": ["guardrail-tag"]}}
assert call_kwargs["valid_token"] is mock_user_api_key_dict
tag_budget_check.reset_mock()
await _enforce_guardrail_added_tag_budgets(
data={"metadata": {"tags": ["existing-tag"]}},
tags_before_guardrails=frozenset({"existing-tag"}),
user_api_key_dict=mock_user_api_key_dict,
proxy_logging_obj=mock_proxy_logging_obj,
)
tag_budget_check.assert_not_awaited()
@pytest.mark.asyncio
async def test_common_processing_pre_call_logic_arms_auto_router_compression_before_guardrails(
self, monkeypatch