diff --git a/tests/test_litellm/proxy/policy_engine/test_policy_matcher.py b/tests/test_litellm/proxy/policy_engine/test_policy_matcher.py new file mode 100644 index 00000000000..f0d67bc2bb8 --- /dev/null +++ b/tests/test_litellm/proxy/policy_engine/test_policy_matcher.py @@ -0,0 +1,84 @@ +""" +Unit tests for PolicyMatcher - tests wildcard pattern matching for policies. + +Tests: +- Wildcard matching (*, prefix-*) +- Scope matching (teams, keys, models) +""" + +import pytest + +from litellm.proxy.policy_engine.policy_matcher import PolicyMatcher +from litellm.types.proxy.policy_engine import ( + Policy, + PolicyGuardrails, + PolicyMatchContext, + PolicyScope, +) + + +class TestPolicyMatcherGetMatchingPolicies: + """Test getting matching policies from a set of policies.""" + + def test_get_matching_policies_by_team(self): + """Test matching policies by team alias.""" + policies = { + "healthcare": Policy( + guardrails=PolicyGuardrails(add=["hipaa_audit"]), + scope=PolicyScope(teams=["healthcare-team"]), + ), + } + + # Match + context = PolicyMatchContext(team_alias="healthcare-team", key_alias="k", model="gpt-4") + assert "healthcare" in PolicyMatcher.get_matching_policies(policies=policies, context=context) + + # No match + context = PolicyMatchContext(team_alias="finance-team", key_alias="k", model="gpt-4") + assert len(PolicyMatcher.get_matching_policies(policies=policies, context=context)) == 0 + + def test_get_matching_policies_by_model_wildcard(self): + """Test matching policies by model with wildcard pattern.""" + policies = { + "bedrock-only": Policy( + guardrails=PolicyGuardrails(add=["pii_blocker"]), + scope=PolicyScope(models=["bedrock/*"]), + ), + } + + # Match - bedrock model + context = PolicyMatchContext(team_alias="t", key_alias="k", model="bedrock/claude-3") + assert "bedrock-only" in PolicyMatcher.get_matching_policies(policies=policies, context=context) + + # No match - different provider + context = PolicyMatchContext(team_alias="t", key_alias="k", model="openai/gpt-4") + assert len(PolicyMatcher.get_matching_policies(policies=policies, context=context)) == 0 + + def test_get_matching_policies_by_key_pattern(self): + """Test matching policies by key alias pattern.""" + policies = { + "dev-keys": Policy( + guardrails=PolicyGuardrails(add=["toxicity_filter"]), + scope=PolicyScope(keys=["dev-key-*"]), + ), + } + + # Match + context = PolicyMatchContext(team_alias="t", key_alias="dev-key-123", model="gpt-4") + assert "dev-keys" in PolicyMatcher.get_matching_policies(policies=policies, context=context) + + # No match + context = PolicyMatchContext(team_alias="t", key_alias="prod-key-123", model="gpt-4") + assert len(PolicyMatcher.get_matching_policies(policies=policies, context=context)) == 0 + + def test_get_matching_policies_global_wildcard(self): + """Test global policy with '*' matches everything.""" + policies = { + "global": Policy( + guardrails=PolicyGuardrails(add=["pii_blocker"]), + scope=PolicyScope(teams=["*"], keys=["*"], models=["*"]), + ), + } + + context = PolicyMatchContext(team_alias="any-team", key_alias="any-key", model="any-model") + assert "global" in PolicyMatcher.get_matching_policies(policies=policies, context=context)