From 715e3b6b017962709d6bce5114c6f1f1193cb63a Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Thu, 22 Jan 2026 14:34:39 -0800 Subject: [PATCH] TestPolicyMatcherGetMatchingPolicies --- .../policy_engine/test_policy_resolver.py | 96 +++++++++++++++++++ 1 file changed, 96 insertions(+) create mode 100644 tests/test_litellm/proxy/policy_engine/test_policy_resolver.py diff --git a/tests/test_litellm/proxy/policy_engine/test_policy_resolver.py b/tests/test_litellm/proxy/policy_engine/test_policy_resolver.py new file mode 100644 index 00000000000..ba85d303eef --- /dev/null +++ b/tests/test_litellm/proxy/policy_engine/test_policy_resolver.py @@ -0,0 +1,96 @@ +""" +Unit tests for PolicyResolver - tests guardrail resolution for request contexts. +""" + +import pytest + +from litellm.proxy.policy_engine.policy_resolver import PolicyResolver +from litellm.types.proxy.policy_engine import ( + Policy, + PolicyGuardrails, + PolicyMatchContext, + PolicyScope, +) + + +class TestPolicyMatcherGetMatchingPolicies: + """Test resolve_guardrails_for_context - the main entry point.""" + + def test_resolve_guardrails_simple_match(self): + """Test resolving guardrails for a simple matching policy.""" + policies = { + "global": Policy( + guardrails=PolicyGuardrails(add=["pii_blocker", "toxicity_filter"]), + scope=PolicyScope(teams=["*"]), + ), + } + + context = PolicyMatchContext(team_alias="any-team", key_alias="k", model="gpt-4") + guardrails = PolicyResolver.resolve_guardrails_for_context( + context=context, policies=policies + ) + + assert set(guardrails) == {"pii_blocker", "toxicity_filter"} + + def test_resolve_guardrails_with_inheritance(self): + """Test child policy inherits and adds guardrails from parent.""" + policies = { + "base": Policy( + guardrails=PolicyGuardrails(add=["pii_blocker"]), + scope=PolicyScope(teams=["*"]), + ), + "healthcare": Policy( + inherit="base", + guardrails=PolicyGuardrails(add=["hipaa_audit"]), + scope=PolicyScope(teams=["healthcare-team"]), + ), + } + + context = PolicyMatchContext(team_alias="healthcare-team", key_alias="k", model="gpt-4") + guardrails = PolicyResolver.resolve_guardrails_for_context( + context=context, policies=policies + ) + + # Both base and healthcare match, healthcare inherits from base + assert set(guardrails) == {"pii_blocker", "hipaa_audit"} + + def test_resolve_guardrails_with_remove(self): + """Test child policy can remove guardrails from parent in its inheritance chain.""" + policies = { + "base": Policy( + guardrails=PolicyGuardrails(add=["pii_blocker", "phi_blocker"]), + scope=PolicyScope(teams=["internal-only"]), # Does NOT match dev-team + ), + "dev": Policy( + inherit="base", + guardrails=PolicyGuardrails(add=["toxicity_filter"], remove=["phi_blocker"]), + scope=PolicyScope(teams=["dev-team"]), # Only this matches + ), + } + + # Only dev policy matches (base scope doesn't match) + context = PolicyMatchContext(team_alias="dev-team", key_alias="k", model="gpt-4") + guardrails = PolicyResolver.resolve_guardrails_for_context( + context=context, policies=policies + ) + + # dev inherits pii_blocker from base, adds toxicity_filter, removes phi_blocker + assert "pii_blocker" in guardrails + assert "toxicity_filter" in guardrails + assert "phi_blocker" not in guardrails + + def test_resolve_guardrails_no_match(self): + """Test returns empty list when no policies match.""" + policies = { + "healthcare": Policy( + guardrails=PolicyGuardrails(add=["hipaa_audit"]), + scope=PolicyScope(teams=["healthcare-team"]), + ), + } + + context = PolicyMatchContext(team_alias="finance-team", key_alias="k", model="gpt-4") + guardrails = PolicyResolver.resolve_guardrails_for_context( + context=context, policies=policies + ) + + assert guardrails == []