151 lines
5.2 KiB
Python
151 lines
5.2 KiB
Python
"""工具与确认 API 测试：TestClient + respx 模拟沙箱。"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from collections.abc import Iterator
|
|
|
|
import httpx
|
|
import pytest
|
|
import respx
|
|
from fastapi.testclient import TestClient
|
|
|
|
from app.api.v1.deps import get_confirmation_store, get_tool_runtime
|
|
from app.application.tool_runtime.runtime import ToolRuntime
|
|
from app.core.config import SandboxSettings
|
|
from app.integrations.sandbox.client import SandboxClient
|
|
from app.main import app
|
|
from app.modules.runtime.events import InMemoryEventStore
|
|
from app.modules.safety.confirmation import InMemoryConfirmationStore
|
|
from app.modules.safety.governor import PermissionPolicy, SafetyGovernor, ToolPolicy
|
|
|
|
|
|
def _settings() -> SandboxSettings:
|
|
return SandboxSettings(
|
|
rust_sandbox_url="http://sandbox.test",
|
|
sandbox_auth_enabled=True,
|
|
sandbox_auth_secret="shared-secret",
|
|
sandbox_default_timeout_ms=30000,
|
|
sandbox_auth_ttl_seconds=30,
|
|
)
|
|
|
|
|
|
def _ok() -> httpx.Response:
|
|
return httpx.Response(
|
|
200,
|
|
json={
|
|
"ok": True,
|
|
"exit_code": 0,
|
|
"stdout": "42\n",
|
|
"stderr": "",
|
|
"duration_ms": 3,
|
|
"timeout": False,
|
|
"truncated": False,
|
|
"trace_id": "t1",
|
|
},
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def runtime_bundle() -> Iterator[None]:
|
|
"""注入 mock ToolRuntime / 确认存储到 app 依赖覆盖；返回值无意义。"""
|
|
store = InMemoryEventStore()
|
|
confirmations = InMemoryConfirmationStore()
|
|
http = httpx.AsyncClient()
|
|
sandbox = SandboxClient(http, _settings())
|
|
runtime = ToolRuntime(sandbox, SafetyGovernor(), store)
|
|
app.dependency_overrides[get_tool_runtime] = lambda: runtime
|
|
app.dependency_overrides[get_confirmation_store] = lambda: confirmations
|
|
try:
|
|
yield
|
|
finally:
|
|
app.dependency_overrides.clear()
|
|
|
|
|
|
@respx.mock
|
|
def test_list_tools(runtime_bundle: None) -> None:
|
|
_ = runtime_bundle # fixture 仅用于副作用（注入依赖覆盖）
|
|
with TestClient(app) as client:
|
|
resp = client.get("/api/v1/tools")
|
|
assert resp.status_code == 200
|
|
names = {tool["name"] for tool in resp.json()}
|
|
assert {"datetime", "shell", "fetch_url"} <= names
|
|
|
|
|
|
@respx.mock
|
|
def test_invoke_datetime_completed(runtime_bundle: None) -> None:
|
|
_ = runtime_bundle
|
|
respx.post("http://sandbox.test/v1/exec").mock(return_value=_ok())
|
|
with TestClient(app) as client:
|
|
resp = client.post(
|
|
"/api/v1/tools/datetime/invoke",
|
|
json={"input": {}, "trace_id": "t1"},
|
|
)
|
|
assert resp.status_code == 200
|
|
body = resp.json()
|
|
assert body["status"] == "completed"
|
|
assert body["result"]["stdout"] == "42\n"
|
|
|
|
|
|
@respx.mock
|
|
def test_shell_invoke_then_approve(runtime_bundle: None) -> None:
|
|
_ = runtime_bundle
|
|
respx.post("http://sandbox.test/v1/exec").mock(return_value=_ok())
|
|
with TestClient(app) as client:
|
|
resp = client.post(
|
|
"/api/v1/tools/shell/invoke",
|
|
json={"input": {"command": "echo hi"}, "trace_id": "t1", "run_id": "run_tools_test"},
|
|
)
|
|
assert resp.status_code == 200
|
|
ask = resp.json()
|
|
assert ask["status"] == "requires_confirmation"
|
|
assert ask["risk_level"] == "L3"
|
|
confirmation_id = ask["confirmation_id"]
|
|
|
|
approved = client.post(f"/api/v1/safety/confirmations/{confirmation_id}/approve")
|
|
assert approved.status_code == 200
|
|
assert approved.json()["status"] == "resolved"
|
|
|
|
|
|
@respx.mock
|
|
def test_deny_then_approve_returns_404(runtime_bundle: None) -> None:
|
|
_ = runtime_bundle
|
|
respx.post("http://sandbox.test/v1/exec").mock(return_value=_ok())
|
|
with TestClient(app) as client:
|
|
resp = client.post(
|
|
"/api/v1/tools/shell/invoke",
|
|
json={"input": {"command": "echo hi"}, "run_id": "run_tools_test2"},
|
|
)
|
|
confirmation_id = resp.json()["confirmation_id"]
|
|
|
|
denied = client.post(f"/api/v1/safety/confirmations/{confirmation_id}/deny")
|
|
assert denied.status_code == 200
|
|
assert denied.json()["status"] in ("cancelled", "resolved")
|
|
|
|
again = client.post(f"/api/v1/safety/confirmations/{confirmation_id}/approve")
|
|
assert again.status_code == 404
|
|
assert again.json()["error"]["code"] == "not_found"
|
|
|
|
|
|
@respx.mock
|
|
def test_forbidden_tool_returns_403() -> None:
|
|
respx.post("http://sandbox.test/v1/exec").mock(return_value=_ok())
|
|
store = InMemoryEventStore()
|
|
confirmations = InMemoryConfirmationStore()
|
|
http = httpx.AsyncClient()
|
|
sandbox = SandboxClient(http, _settings())
|
|
governor = SafetyGovernor(
|
|
{"shell": ToolPolicy(risk_level="L3", permission_policy=PermissionPolicy.forbidden)}
|
|
)
|
|
runtime = ToolRuntime(sandbox, governor, store)
|
|
app.dependency_overrides[get_tool_runtime] = lambda: runtime
|
|
app.dependency_overrides[get_confirmation_store] = lambda: confirmations
|
|
try:
|
|
with TestClient(app) as client:
|
|
resp = client.post(
|
|
"/api/v1/tools/shell/invoke", json={"input": {"command": "rm -rf /"}}
|
|
)
|
|
assert resp.status_code == 403
|
|
assert resp.json()["error"]["code"] == "permission_denied"
|
|
finally:
|
|
app.dependency_overrides.clear()
|