113 lines
3.5 KiB
Python
113 lines
3.5 KiB
Python
from __future__ import annotations
|
|
|
|
from fastapi import Depends, FastAPI
|
|
from fastapi.testclient import TestClient
|
|
from option_mcp_common.auth import Principal, TokenStore, require_principal
|
|
from option_mcp_common.mcp_bridge import _derive_input_schemas, mount_mcp_endpoint
|
|
from option_mcp_common.server import build_app
|
|
from pydantic import BaseModel
|
|
|
|
|
|
class EchoBody(BaseModel):
|
|
msg: str
|
|
n: int = 1
|
|
|
|
|
|
def _make_app() -> tuple[FastAPI, TokenStore]:
|
|
store = TokenStore(tokens={"t": Principal("obs", {"observer"})})
|
|
app = build_app(name="t", version="v", token_store=store)
|
|
|
|
@app.post("/tools/echo")
|
|
def echo(body: EchoBody, principal: Principal = Depends(require_principal)):
|
|
return {"echo": body.msg, "n": body.n}
|
|
|
|
@app.post("/tools/ping")
|
|
def ping(principal: Principal = Depends(require_principal)):
|
|
return {"pong": True}
|
|
|
|
return app, store
|
|
|
|
|
|
def test_derive_input_schemas_resolves_lazy_annotations():
|
|
app, _ = _make_app()
|
|
schemas = _derive_input_schemas(app, ["echo", "ping"])
|
|
assert "echo" in schemas
|
|
echo_schema = schemas["echo"]
|
|
assert echo_schema["type"] == "object"
|
|
assert "msg" in echo_schema["properties"]
|
|
assert "n" in echo_schema["properties"]
|
|
assert "msg" in echo_schema["required"]
|
|
# ping has no Pydantic body → not in map (fallback applied by caller)
|
|
assert "ping" not in schemas
|
|
|
|
|
|
def test_mount_mcp_endpoint_exposes_derived_schemas():
|
|
app, store = _make_app()
|
|
mount_mcp_endpoint(
|
|
app,
|
|
name="test",
|
|
version="1.0",
|
|
token_store=store,
|
|
internal_base_url="http://localhost:0",
|
|
tools=[
|
|
{"name": "echo", "description": "Echo a message."},
|
|
{"name": "ping", "description": "Ping."},
|
|
],
|
|
)
|
|
c = TestClient(app)
|
|
r = c.post(
|
|
"/mcp",
|
|
headers={"Authorization": "Bearer t"},
|
|
json={"jsonrpc": "2.0", "id": 1, "method": "tools/list"},
|
|
)
|
|
assert r.status_code == 200
|
|
tools = r.json()["result"]["tools"]
|
|
by_name = {t["name"]: t for t in tools}
|
|
assert set(by_name["echo"]["inputSchema"]["required"]) == {"msg"}
|
|
# ping fallback su schema generico
|
|
assert by_name["ping"]["inputSchema"] == {
|
|
"type": "object",
|
|
"additionalProperties": True,
|
|
}
|
|
|
|
|
|
def test_mount_mcp_endpoint_requires_auth():
|
|
app, store = _make_app()
|
|
mount_mcp_endpoint(
|
|
app,
|
|
name="test",
|
|
version="1.0",
|
|
token_store=store,
|
|
internal_base_url="http://localhost:0",
|
|
tools=[{"name": "echo"}],
|
|
)
|
|
c = TestClient(app)
|
|
r = c.post("/mcp", json={"jsonrpc": "2.0", "id": 1, "method": "tools/list"})
|
|
assert r.status_code == 401
|
|
r = c.post(
|
|
"/mcp",
|
|
headers={"Authorization": "Bearer WRONG"},
|
|
json={"jsonrpc": "2.0", "id": 1, "method": "tools/list"},
|
|
)
|
|
assert r.status_code == 403
|
|
|
|
|
|
def test_explicit_input_schema_overrides_derived():
|
|
app, store = _make_app()
|
|
custom = {"type": "object", "properties": {"custom": {"type": "string"}}, "required": ["custom"]}
|
|
mount_mcp_endpoint(
|
|
app,
|
|
name="test",
|
|
version="1.0",
|
|
token_store=store,
|
|
internal_base_url="http://localhost:0",
|
|
tools=[{"name": "echo", "input_schema": custom}],
|
|
)
|
|
c = TestClient(app)
|
|
r = c.post(
|
|
"/mcp",
|
|
headers={"Authorization": "Bearer t"},
|
|
json={"jsonrpc": "2.0", "id": 1, "method": "tools/list"},
|
|
)
|
|
assert r.json()["result"]["tools"][0]["inputSchema"] == custom
|