diff --git a/tests/test_litellm/proxy/policy_engine/test_attachment_registry.py b/tests/test_litellm/proxy/policy_engine/test_attachment_registry.py index 9ebdd6d6ecc..1ed956fe99f 100644 --- a/tests/test_litellm/proxy/policy_engine/test_attachment_registry.py +++ b/tests/test_litellm/proxy/policy_engine/test_attachment_registry.py @@ -1,11 +1,7 @@ """ -Unit tests for AttachmentRegistry - tests policy attachment management. +Unit tests for AttachmentRegistry - tests policy attachment matching. -Tests: -- Loading attachments from config -- Getting attached policies for a context -- Global scope attachments -- Team/key/model specific attachments +Tests the main entry point: get_attached_policies() """ import pytest @@ -14,69 +10,20 @@ from litellm.proxy.policy_engine.attachment_registry import ( AttachmentRegistry, get_attachment_registry, ) -from litellm.types.proxy.policy_engine import ( - PolicyAttachment, - PolicyMatchContext, -) - - -class TestAttachmentRegistryLoading: - """Test loading attachments from configuration.""" - - def test_load_attachments_simple(self): - """Test loading simple attachments.""" - registry = AttachmentRegistry() - config = [ - {"policy": "global-baseline", "scope": "*"}, - {"policy": "healthcare-policy", "teams": ["healthcare-team"]}, - ] - registry.load_attachments(config) - - assert registry.is_initialized() - assert len(registry.get_all_attachments()) == 2 - - def test_load_attachments_with_multiple_scopes(self): - """Test loading attachments with multiple scope types.""" - registry = AttachmentRegistry() - config = [ - {"policy": "global-baseline", "scope": "*"}, - {"policy": "team-policy", "teams": ["team-a", "team-b"]}, - {"policy": "key-policy", "keys": ["dev-key-*"]}, - {"policy": "model-policy", "models": ["gpt-4", "gpt-4-turbo"]}, - ] - registry.load_attachments(config) - - assert len(registry.get_all_attachments()) == 4 - - def test_load_attachments_empty_list(self): - """Test loading empty attachments list.""" - registry = AttachmentRegistry() - registry.load_attachments([]) - - assert registry.is_initialized() - assert len(registry.get_all_attachments()) == 0 - - def test_clear_attachments(self): - """Test clearing attachments.""" - registry = AttachmentRegistry() - registry.load_attachments([{"policy": "test", "scope": "*"}]) - assert registry.is_initialized() - - registry.clear() - assert not registry.is_initialized() - assert len(registry.get_all_attachments()) == 0 +from litellm.types.proxy.policy_engine import PolicyMatchContext class TestGetAttachedPolicies: - """Test getting attached policies for a context.""" + """Test get_attached_policies - the main entry point.""" - def test_global_scope_matches_all(self): - """Test global scope (*) matches all contexts.""" + def test_global_scope_matches_all_requests(self): + """Test global scope (*) matches any request context.""" registry = AttachmentRegistry() registry.load_attachments([ {"policy": "global-baseline", "scope": "*"}, ]) + # Should match any context context = PolicyMatchContext( team_alias="any-team", key_alias="any-key", model="any-model" ) @@ -94,36 +41,32 @@ class TestGetAttachedPolicies: context = PolicyMatchContext( team_alias="healthcare-team", key_alias="key", model="gpt-4" ) - attached = registry.get_attached_policies(context) - assert "healthcare-policy" in attached + assert "healthcare-policy" in registry.get_attached_policies(context) - # No match + # No match - different team context_other = PolicyMatchContext( team_alias="finance-team", key_alias="key", model="gpt-4" ) - attached_other = registry.get_attached_policies(context_other) - assert "healthcare-policy" not in attached_other + assert "healthcare-policy" not in registry.get_attached_policies(context_other) - def test_key_pattern_attachment(self): - """Test key pattern attachment matches wildcard.""" + def test_key_wildcard_pattern_attachment(self): + """Test key pattern attachment with wildcard.""" registry = AttachmentRegistry() registry.load_attachments([ {"policy": "dev-policy", "keys": ["dev-key-*"]}, ]) - # Match + # Match - key starts with dev-key- context = PolicyMatchContext( team_alias="team", key_alias="dev-key-123", model="gpt-4" ) - attached = registry.get_attached_policies(context) - assert "dev-policy" in attached + assert "dev-policy" in registry.get_attached_policies(context) - # No match + # No match - different prefix context_prod = PolicyMatchContext( team_alias="team", key_alias="prod-key-123", model="gpt-4" ) - attached_prod = registry.get_attached_policies(context_prod) - assert "dev-policy" not in attached_prod + assert "dev-policy" not in registry.get_attached_policies(context_prod) def test_model_specific_attachment(self): """Test model-specific attachment.""" @@ -136,18 +79,35 @@ class TestGetAttachedPolicies: context = PolicyMatchContext( team_alias="team", key_alias="key", model="gpt-4" ) - attached = registry.get_attached_policies(context) - assert "gpt4-policy" in attached + assert "gpt4-policy" in registry.get_attached_policies(context) # No match context_other = PolicyMatchContext( team_alias="team", key_alias="key", model="gpt-3.5" ) - attached_other = registry.get_attached_policies(context_other) - assert "gpt4-policy" not in attached_other + assert "gpt4-policy" not in registry.get_attached_policies(context_other) - def test_multiple_attachments_match(self): - """Test multiple attachments can match same context.""" + def test_model_wildcard_pattern(self): + """Test model wildcard pattern like bedrock/*.""" + registry = AttachmentRegistry() + registry.load_attachments([ + {"policy": "bedrock-policy", "models": ["bedrock/*"]}, + ]) + + # Match + context = PolicyMatchContext( + team_alias="team", key_alias="key", model="bedrock/claude-3" + ) + assert "bedrock-policy" in registry.get_attached_policies(context) + + # No match + context_other = PolicyMatchContext( + team_alias="team", key_alias="key", model="openai/gpt-4" + ) + assert "bedrock-policy" not in registry.get_attached_policies(context_other) + + def test_multiple_attachments_match_same_context(self): + """Test multiple attachments can match the same context.""" registry = AttachmentRegistry() registry.load_attachments([ {"policy": "global-baseline", "scope": "*"}, @@ -160,12 +120,13 @@ class TestGetAttachedPolicies: ) attached = registry.get_attached_policies(context) + # All three should match assert "global-baseline" in attached assert "healthcare-policy" in attached assert "gpt4-policy" in attached assert len(attached) == 3 - def test_no_duplicate_policies(self): + def test_same_policy_multiple_attachments_no_duplicates(self): """Test same policy attached multiple ways doesn't duplicate.""" registry = AttachmentRegistry() registry.load_attachments([ @@ -181,108 +142,60 @@ class TestGetAttachedPolicies: # Should only appear once assert attached.count("multi-policy") == 1 - -class TestIsPolicyAttached: - """Test is_policy_attached method.""" - - def test_policy_is_attached(self): - """Test checking if a specific policy is attached.""" - registry = AttachmentRegistry() - registry.load_attachments([ - {"policy": "global-baseline", "scope": "*"}, - ]) - - context = PolicyMatchContext(team_alias="team", key_alias="key", model="gpt-4") - assert registry.is_policy_attached("global-baseline", context) is True - assert registry.is_policy_attached("other-policy", context) is False - - -class TestGetAttachmentsForPolicy: - """Test getting attachments for a specific policy.""" - - def test_get_attachments_for_policy(self): - """Test getting all attachments for a policy.""" - registry = AttachmentRegistry() - registry.load_attachments([ - {"policy": "multi-policy", "scope": "*"}, - {"policy": "multi-policy", "teams": ["team-a"]}, - {"policy": "other-policy", "teams": ["team-b"]}, - ]) - - attachments = registry.get_attachments_for_policy("multi-policy") - assert len(attachments) == 2 - - attachments_other = registry.get_attachments_for_policy("other-policy") - assert len(attachments_other) == 1 - - -class TestAddAndRemoveAttachments: - """Test adding and removing individual attachments.""" - - def test_add_attachment(self): - """Test adding a single attachment.""" + def test_no_attachments_returns_empty(self): + """Test empty attachments returns empty list.""" registry = AttachmentRegistry() registry.load_attachments([]) - attachment = PolicyAttachment(policy="new-policy", scope="*") - registry.add_attachment(attachment) + context = PolicyMatchContext( + team_alias="team", key_alias="key", model="gpt-4" + ) + attached = registry.get_attached_policies(context) + assert attached == [] - assert len(registry.get_all_attachments()) == 1 - - def test_remove_attachments_for_policy(self): - """Test removing all attachments for a policy.""" + def test_no_matching_attachments_returns_empty(self): + """Test no matching attachments returns empty list.""" registry = AttachmentRegistry() registry.load_attachments([ - {"policy": "policy-a", "scope": "*"}, - {"policy": "policy-a", "teams": ["team-a"]}, - {"policy": "policy-b", "teams": ["team-b"]}, + {"policy": "healthcare-policy", "teams": ["healthcare-team"]}, ]) - removed = registry.remove_attachments_for_policy("policy-a") - assert removed == 2 - assert len(registry.get_all_attachments()) == 1 - assert len(registry.get_attachments_for_policy("policy-a")) == 0 - - -class TestPolicyAttachmentModel: - """Test PolicyAttachment model methods.""" - - def test_is_global(self): - """Test is_global method.""" - global_attachment = PolicyAttachment(policy="test", scope="*") - assert global_attachment.is_global() is True - - team_attachment = PolicyAttachment(policy="test", teams=["team-a"]) - assert team_attachment.is_global() is False - - def test_to_policy_scope_global(self): - """Test converting global attachment to PolicyScope.""" - attachment = PolicyAttachment(policy="test", scope="*") - scope = attachment.to_policy_scope() - - assert scope.get_teams() == ["*"] - assert scope.get_keys() == ["*"] - assert scope.get_models() == ["*"] - - def test_to_policy_scope_specific(self): - """Test converting specific attachment to PolicyScope.""" - attachment = PolicyAttachment( - policy="test", - teams=["team-a", "team-b"], - keys=["key-*"], - models=["gpt-4"], + context = PolicyMatchContext( + team_alias="finance-team", key_alias="key", model="gpt-4" ) - scope = attachment.to_policy_scope() + attached = registry.get_attached_policies(context) + assert attached == [] - assert scope.teams == ["team-a", "team-b"] - assert scope.keys == ["key-*"] - assert scope.models == ["gpt-4"] + def test_combined_team_and_model_attachment(self): + """Test attachment with both team and model constraints.""" + registry = AttachmentRegistry() + registry.load_attachments([ + {"policy": "strict-policy", "teams": ["healthcare-team"], "models": ["gpt-4"]}, + ]) + + # Match - both team and model match + context = PolicyMatchContext( + team_alias="healthcare-team", key_alias="key", model="gpt-4" + ) + assert "strict-policy" in registry.get_attached_policies(context) + + # No match - team matches but model doesn't + context_wrong_model = PolicyMatchContext( + team_alias="healthcare-team", key_alias="key", model="gpt-3.5" + ) + assert "strict-policy" not in registry.get_attached_policies(context_wrong_model) + + # No match - model matches but team doesn't + context_wrong_team = PolicyMatchContext( + team_alias="finance-team", key_alias="key", model="gpt-4" + ) + assert "strict-policy" not in registry.get_attached_policies(context_wrong_team) -class TestGlobalSingleton: +class TestAttachmentRegistrySingleton: """Test global singleton behavior.""" - def test_get_attachment_registry_singleton(self): + def test_get_attachment_registry_returns_same_instance(self): """Test get_attachment_registry returns same instance.""" registry1 = get_attachment_registry() registry2 = get_attachment_registry()