74 lines
3.0 KiB
Python
74 lines
3.0 KiB
Python
import pytest
|
|
from pydantic import BaseModel
|
|
|
|
from app.core.contracts import RequestContext
|
|
from app.core.errors import ForbiddenAgentError, ValidationAgentError
|
|
from app.service.tool_executor import ToolDefinition, ToolExecutor, ToolRegistry
|
|
|
|
|
|
class QueryInput(BaseModel):
|
|
customer_id: int
|
|
|
|
|
|
async def query(arguments: QueryInput, context: RequestContext):
|
|
return {"customer_id": arguments.customer_id, "value": "readonly"}
|
|
|
|
|
|
def executor() -> ToolExecutor:
|
|
registry = ToolRegistry()
|
|
registry.register(ToolDefinition(
|
|
name="query_profile", input_model=QueryInput, handler=query,
|
|
required_permission="profile:read", allowed_roles=("advisor",),
|
|
))
|
|
tool = ToolExecutor(registry)
|
|
tool._audit = lambda *args: __import__("asyncio").sleep(0) # type: ignore[method-assign]
|
|
return tool
|
|
|
|
|
|
def context() -> RequestContext:
|
|
return RequestContext(user_id="1", trace_id="trace", roles=("advisor",),
|
|
permissions=("profile:read",))
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_tool_validates_acl_schema_and_returns_reference():
|
|
result = await executor().execute(
|
|
name="query_profile", arguments={"customer_id": 1}, intent="profile",
|
|
configured_tools={"profile": ("query_profile",)}, context=context())
|
|
assert result.output["value"] == "readonly"
|
|
assert result.record.status == "succeeded"
|
|
assert result.references[0].source_id == "trace:query_profile"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_tool_denies_not_in_intent_allowlist_and_invalid_schema():
|
|
tool = executor()
|
|
with pytest.raises(ForbiddenAgentError):
|
|
await tool.execute(name="query_profile", arguments={"customer_id": 1}, intent="other",
|
|
configured_tools={"other": ()}, context=context())
|
|
with pytest.raises(ValidationAgentError):
|
|
await tool.execute(name="query_profile", arguments={"customer_id": "bad"}, intent="profile",
|
|
configured_tools={"profile": ("query_profile",)}, context=context())
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_tool_denies_missing_permission_and_unknown_tool():
|
|
user = context().model_copy(update={"permissions": ()})
|
|
with pytest.raises(ForbiddenAgentError):
|
|
await executor().execute(
|
|
name="query_profile", arguments={"customer_id": 1}, intent="profile",
|
|
configured_tools={"profile": ("query_profile",)}, context=user)
|
|
with pytest.raises(ForbiddenAgentError):
|
|
await executor().execute(name="unknown", arguments={}, intent="profile",
|
|
configured_tools={"profile": ("unknown",)}, context=context())
|
|
|
|
|
|
def test_registry_rejects_writable_tools_and_duplicates():
|
|
registry = ToolRegistry()
|
|
definition = ToolDefinition("query", QueryInput, query, "x", ("a",))
|
|
registry.register(definition)
|
|
with pytest.raises(ValidationAgentError):
|
|
registry.register(definition)
|
|
with pytest.raises(ValidationAgentError):
|
|
registry.register(ToolDefinition("write", QueryInput, query, "x", ("a",), read_only=False))
|