142 lines
5.9 KiB
Python
142 lines
5.9 KiB
Python
from __future__ import annotations
|
|
|
|
from app.ai_provider.routing.permissions import CapabilityPermissionController, PermissionRule
|
|
|
|
|
|
class TestPermissionRule:
|
|
def test_exact_match(self):
|
|
rule = PermissionRule(model_pattern="fast-model", capability="vision", allowed=True)
|
|
assert rule.match("fast-model") is True
|
|
assert rule.match("pro-model") is False
|
|
|
|
def test_wildcard_match(self):
|
|
rule = PermissionRule(model_pattern="deepseek-*", capability="thinking", allowed=False)
|
|
assert rule.match("deepseek-v4-flash") is True
|
|
assert rule.match("deepseek-v4-pro") is True
|
|
assert rule.match("claude-sonnet-4-6") is False
|
|
|
|
def test_universal_match(self):
|
|
rule = PermissionRule(model_pattern="*", capability="vision", allowed=True)
|
|
assert rule.match("any-model") is True
|
|
assert rule.match("another-one") is True
|
|
|
|
def test_suffix_wildcard(self):
|
|
rule = PermissionRule(model_pattern="*-pro", capability="thinking", allowed=True)
|
|
assert rule.match("deepseek-v4-pro") is True
|
|
assert rule.match("fast-model") is False
|
|
|
|
def test_match_method_called_explicitly(self):
|
|
rule = PermissionRule(model_pattern="vision-*", capability="vision")
|
|
assert rule.match("vision-model") is True
|
|
|
|
def test_allowed_defaults_to_true(self):
|
|
rule = PermissionRule(model_pattern="*", capability="vision")
|
|
assert rule.allowed is True
|
|
|
|
def test_allowed_set_to_false(self):
|
|
rule = PermissionRule(model_pattern="*", capability="vision", allowed=False)
|
|
assert rule.allowed is False
|
|
|
|
def test_repr(self):
|
|
rule = PermissionRule(model_pattern="*", capability="vision", allowed=True)
|
|
r = repr(rule)
|
|
assert "PermissionRule" in r
|
|
assert "*" in r
|
|
assert "vision" in r
|
|
|
|
def test_eq_same_rules(self):
|
|
rule1 = PermissionRule(model_pattern="foo-*", capability="vision", allowed=True)
|
|
rule2 = PermissionRule(model_pattern="foo-*", capability="vision", allowed=True)
|
|
assert rule1 == rule2
|
|
|
|
def test_eq_different_rules(self):
|
|
rule1 = PermissionRule(model_pattern="foo-*", capability="vision", allowed=True)
|
|
rule2 = PermissionRule(model_pattern="bar-*", capability="vision", allowed=True)
|
|
assert rule1 != rule2
|
|
|
|
|
|
class TestCapabilityPermissionController:
|
|
def test_allow_all_by_default(self):
|
|
controller = CapabilityPermissionController([])
|
|
assert controller.check("fast-model", "vision") is True
|
|
|
|
def test_deny_specific_capability_on_model(self):
|
|
controller = CapabilityPermissionController(
|
|
[
|
|
PermissionRule(model_pattern="fast-*", capability="vision", allowed=False),
|
|
]
|
|
)
|
|
assert controller.check("fast-model", "vision") is False
|
|
assert controller.check("pro-model", "vision") is True
|
|
|
|
def test_deny_then_allow_more_specific(self):
|
|
"""A more specific deny should override a broader allow."""
|
|
controller = CapabilityPermissionController(
|
|
[
|
|
PermissionRule(model_pattern="*", capability="thinking", allowed=True),
|
|
PermissionRule(model_pattern="fast-*", capability="thinking", allowed=False),
|
|
]
|
|
)
|
|
assert controller.check("fast-model", "thinking") is False
|
|
assert controller.check("pro-model", "thinking") is True
|
|
|
|
def test_filter_capabilities(self):
|
|
controller = CapabilityPermissionController(
|
|
[
|
|
PermissionRule(model_pattern="*", capability="vision", allowed=False),
|
|
PermissionRule(model_pattern="vision-*", capability="vision", allowed=True),
|
|
]
|
|
)
|
|
filtered = controller.filter_capabilities("fast-model", {"chat", "vision", "thinking"})
|
|
assert filtered == {"chat", "thinking"}
|
|
|
|
filtered_vision = controller.filter_capabilities(
|
|
"vision-model", {"chat", "vision", "thinking"}
|
|
)
|
|
assert filtered_vision == {"chat", "vision", "thinking"}
|
|
|
|
def test_filter_with_no_rules(self):
|
|
controller = CapabilityPermissionController([])
|
|
filtered = controller.filter_capabilities("any-model", {"chat", "vision"})
|
|
assert filtered == {"chat", "vision"}
|
|
|
|
def test_check_nonexistent_model(self):
|
|
"""Non-existent models should still be checked against rules."""
|
|
controller = CapabilityPermissionController(
|
|
[
|
|
PermissionRule(model_pattern="*", capability="vision", allowed=False),
|
|
]
|
|
)
|
|
assert controller.check("unknown-model", "vision") is False
|
|
|
|
def test_multiple_rules_matching(self):
|
|
"""When multiple rules match, the last one wins (most specific wins by order)."""
|
|
controller = CapabilityPermissionController(
|
|
[
|
|
PermissionRule(model_pattern="*", capability="tool_use", allowed=True),
|
|
PermissionRule(model_pattern="basic-*", capability="tool_use", allowed=False),
|
|
PermissionRule(model_pattern="basic-pro", capability="tool_use", allowed=True),
|
|
]
|
|
)
|
|
assert controller.check("basic-lite", "tool_use") is False
|
|
assert controller.check("basic-pro", "tool_use") is True
|
|
assert controller.check("pro-model", "tool_use") is True
|
|
|
|
def test_check_capability_not_in_rules(self):
|
|
controller = CapabilityPermissionController(
|
|
[
|
|
PermissionRule(model_pattern="*", capability="vision", allowed=False),
|
|
]
|
|
)
|
|
# 'chat' is not in any rule, so it's allowed by default
|
|
assert controller.check("fast-model", "chat") is True
|
|
|
|
def test_is_allowed_model(self):
|
|
controller = CapabilityPermissionController(
|
|
[
|
|
PermissionRule(model_pattern="*", capability="vision", allowed=False),
|
|
]
|
|
)
|
|
assert controller.is_allowed("fast-model", "vision") is False
|
|
assert controller.is_allowed("fast-model", "chat") is True
|