228 lines
6.8 KiB
Python
228 lines
6.8 KiB
Python
"""权限模块测试"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import pytest
|
|
|
|
from minenasai.agent.permissions import (
|
|
ConfirmationStatus,
|
|
DangerLevel,
|
|
PermissionManager,
|
|
ToolPermission,
|
|
)
|
|
|
|
|
|
class TestDangerLevel:
|
|
"""DangerLevel 枚举测试"""
|
|
|
|
def test_danger_levels(self):
|
|
"""测试危险等级"""
|
|
assert DangerLevel.SAFE.value == "safe"
|
|
assert DangerLevel.CRITICAL.value == "critical"
|
|
|
|
|
|
class TestToolPermission:
|
|
"""ToolPermission 测试"""
|
|
|
|
def test_default_permission(self):
|
|
"""测试默认权限"""
|
|
perm = ToolPermission(
|
|
name="test_tool",
|
|
danger_level=DangerLevel.SAFE,
|
|
)
|
|
|
|
assert perm.requires_confirmation is False
|
|
assert perm.rate_limit is None
|
|
|
|
|
|
class TestPermissionManager:
|
|
"""PermissionManager 测试"""
|
|
|
|
def setup_method(self):
|
|
"""初始化"""
|
|
self.manager = PermissionManager()
|
|
|
|
def test_get_default_permission(self):
|
|
"""测试获取默认权限"""
|
|
perm = self.manager.get_permission("read_file")
|
|
|
|
assert perm is not None
|
|
assert perm.danger_level == DangerLevel.SAFE
|
|
|
|
def test_register_tool(self):
|
|
"""测试注册工具权限"""
|
|
perm = ToolPermission(
|
|
name="custom_tool",
|
|
danger_level=DangerLevel.MEDIUM,
|
|
description="自定义工具",
|
|
)
|
|
self.manager.register_tool(perm)
|
|
|
|
result = self.manager.get_permission("custom_tool")
|
|
assert result is not None
|
|
assert result.danger_level == DangerLevel.MEDIUM
|
|
|
|
def test_check_permission_allowed(self):
|
|
"""测试权限检查 - 允许"""
|
|
allowed, reason = self.manager.check_permission("read_file")
|
|
|
|
assert allowed is True
|
|
|
|
def test_check_permission_unknown_tool(self):
|
|
"""测试权限检查 - 未知工具"""
|
|
allowed, reason = self.manager.check_permission("unknown_tool")
|
|
|
|
assert allowed is False
|
|
assert "未知工具" in reason
|
|
|
|
def test_check_permission_denied_path(self):
|
|
"""测试权限检查 - 禁止路径"""
|
|
perm = ToolPermission(
|
|
name="restricted_read",
|
|
danger_level=DangerLevel.SAFE,
|
|
denied_paths=["/etc/", "/root/"],
|
|
)
|
|
self.manager.register_tool(perm)
|
|
|
|
allowed, reason = self.manager.check_permission(
|
|
"restricted_read",
|
|
params={"path": "/etc/passwd"},
|
|
)
|
|
|
|
assert allowed is False
|
|
assert "禁止访问" in reason
|
|
|
|
def test_requires_confirmation_by_level(self):
|
|
"""测试确认要求 - 按等级"""
|
|
# HIGH 级别需要确认
|
|
assert self.manager.requires_confirmation("delete_file") is True
|
|
|
|
# SAFE 级别不需要确认
|
|
assert self.manager.requires_confirmation("read_file") is False
|
|
|
|
def test_requires_confirmation_explicit(self):
|
|
"""测试确认要求 - 显式设置"""
|
|
perm = ToolPermission(
|
|
name="explicit_confirm",
|
|
danger_level=DangerLevel.LOW,
|
|
requires_confirmation=True,
|
|
)
|
|
self.manager.register_tool(perm)
|
|
|
|
assert self.manager.requires_confirmation("explicit_confirm") is True
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_confirmation_flow(self):
|
|
"""测试确认流程"""
|
|
request = await self.manager.request_confirmation(
|
|
request_id="req-1",
|
|
tool_name="delete_file",
|
|
params={"path": "/test.txt"},
|
|
)
|
|
|
|
assert request.status == ConfirmationStatus.PENDING
|
|
|
|
# 批准
|
|
self.manager.approve_confirmation("req-1")
|
|
assert request.status == ConfirmationStatus.APPROVED
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_confirmation_deny(self):
|
|
"""测试拒绝确认"""
|
|
request = await self.manager.request_confirmation(
|
|
request_id="req-2",
|
|
tool_name="delete_file",
|
|
params={"path": "/test.txt"},
|
|
)
|
|
|
|
self.manager.deny_confirmation("req-2")
|
|
assert request.status == ConfirmationStatus.DENIED
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_pending_confirmations(self):
|
|
"""测试获取待处理确认"""
|
|
await self.manager.request_confirmation(
|
|
request_id="req-3",
|
|
tool_name="test",
|
|
params={},
|
|
)
|
|
|
|
pending = self.manager.get_pending_confirmations()
|
|
assert len(pending) >= 1
|
|
|
|
|
|
class TestToolRegistry:
|
|
"""ToolRegistry 测试"""
|
|
|
|
def test_import_registry(self):
|
|
"""测试导入注册中心"""
|
|
from minenasai.agent import ToolRegistry, get_tool_registry
|
|
|
|
registry = get_tool_registry()
|
|
assert isinstance(registry, ToolRegistry)
|
|
|
|
def test_register_builtin_tools(self):
|
|
"""测试注册内置工具"""
|
|
from minenasai.agent import get_tool_registry, register_builtin_tools
|
|
|
|
registry = get_tool_registry()
|
|
initial_count = len(registry.list_tools())
|
|
|
|
register_builtin_tools()
|
|
|
|
# 应该有更多工具
|
|
new_count = len(registry.list_tools())
|
|
assert new_count >= initial_count
|
|
|
|
def test_tool_decorator(self):
|
|
"""测试工具装饰器"""
|
|
from minenasai.agent import tool, get_tool_registry
|
|
|
|
@tool(name="decorated_tool", description="装饰器测试")
|
|
async def decorated_tool(param: str) -> dict:
|
|
return {"result": param}
|
|
|
|
registry = get_tool_registry()
|
|
tool_obj = registry.get("decorated_tool")
|
|
|
|
assert tool_obj is not None
|
|
assert tool_obj.description == "装饰器测试"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_execute_tool(self):
|
|
"""测试执行工具"""
|
|
from minenasai.agent import get_tool_registry, DangerLevel
|
|
|
|
registry = get_tool_registry()
|
|
|
|
# 注册测试工具
|
|
async def echo(message: str) -> dict:
|
|
return {"echo": message}
|
|
|
|
registry.register(
|
|
name="echo",
|
|
description="回显消息",
|
|
func=echo,
|
|
parameters={
|
|
"type": "object",
|
|
"properties": {"message": {"type": "string"}},
|
|
"required": ["message"],
|
|
},
|
|
danger_level=DangerLevel.SAFE,
|
|
)
|
|
|
|
result = await registry.execute("echo", {"message": "hello"})
|
|
|
|
assert result["success"] is True
|
|
assert result["result"]["echo"] == "hello"
|
|
|
|
def test_get_stats(self):
|
|
"""测试获取统计"""
|
|
from minenasai.agent import get_tool_registry
|
|
|
|
registry = get_tool_registry()
|
|
stats = registry.get_stats()
|
|
|
|
assert "total_tools" in stats
|
|
assert "categories" in stats
|