222 lines
7.4 KiB
Python
222 lines
7.4 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import os
|
|
import stat
|
|
import tempfile
|
|
import threading
|
|
import unittest
|
|
from datetime import date
|
|
from pathlib import Path
|
|
from unittest.mock import patch
|
|
|
|
from arr_ingestion.contracts import IngestionError
|
|
from arr_ingestion.direct_contracts import DirectSubmissionReceipt
|
|
from arr_mcp.auth import BearerAuthASGI, BearerAuthConfig
|
|
from arr_mcp.database import _read_private_database_config
|
|
from arr_mcp.gateway import DirectResultGateway
|
|
from arr_mcp.run import _allowed_hosts
|
|
from arr_mcp.server import load_tool_schemas
|
|
|
|
|
|
TOKEN = "M" * 48
|
|
|
|
|
|
def receipt() -> DirectSubmissionReceipt:
|
|
return DirectSubmissionReceipt(
|
|
"committed",
|
|
"arrjob-mcp-test",
|
|
1,
|
|
date(2026, 7, 27),
|
|
71,
|
|
2,
|
|
5,
|
|
)
|
|
|
|
|
|
class FakeService:
|
|
def __init__(self) -> None:
|
|
self.requests: list[object] = []
|
|
|
|
def submit(self, raw_request: object) -> DirectSubmissionReceipt:
|
|
self.requests.append(raw_request)
|
|
return receipt()
|
|
|
|
|
|
class BlockingService(FakeService):
|
|
def __init__(self) -> None:
|
|
super().__init__()
|
|
self.entered = threading.Event()
|
|
self.release = threading.Event()
|
|
|
|
def submit(self, raw_request: object) -> DirectSubmissionReceipt:
|
|
self.requests.append(raw_request)
|
|
self.entered.set()
|
|
if not self.release.wait(timeout=5):
|
|
raise AssertionError("blocking test timed out")
|
|
return receipt()
|
|
|
|
|
|
async def invoke_asgi(
|
|
app: object,
|
|
headers: list[tuple[bytes, bytes]],
|
|
) -> list[dict[str, object]]:
|
|
messages: list[dict[str, object]] = []
|
|
|
|
async def receive() -> dict[str, object]:
|
|
return {"type": "http.disconnect"}
|
|
|
|
async def send(message: dict[str, object]) -> None:
|
|
messages.append(message)
|
|
|
|
await app( # type: ignore[operator]
|
|
{"type": "http", "method": "GET", "path": "/mcp", "headers": headers},
|
|
receive,
|
|
send,
|
|
)
|
|
return messages
|
|
|
|
|
|
class DirectResultGatewayTests(unittest.TestCase):
|
|
def test_gateway_returns_only_the_service_receipt(self) -> None:
|
|
service = FakeService()
|
|
gateway = DirectResultGateway(service)
|
|
arguments = {"opaque": "request"}
|
|
result = gateway.submit_processing_result(arguments)
|
|
self.assertEqual(service.requests, [arguments])
|
|
self.assertEqual(result, receipt().to_dict())
|
|
|
|
def test_gateway_fails_fast_when_replay_capacity_is_full(self) -> None:
|
|
service = BlockingService()
|
|
gateway = DirectResultGateway(service, max_concurrency=1)
|
|
failures: list[BaseException] = []
|
|
|
|
def first_call() -> None:
|
|
try:
|
|
gateway.submit_processing_result({"call": 1})
|
|
except BaseException as error: # pragma: no cover - assertion aid
|
|
failures.append(error)
|
|
|
|
worker = threading.Thread(target=first_call)
|
|
worker.start()
|
|
self.assertTrue(service.entered.wait(timeout=2))
|
|
try:
|
|
with self.assertRaises(IngestionError) as raised:
|
|
gateway.submit_processing_result({"call": 2})
|
|
self.assertEqual(raised.exception.code, "DIRECT_GATEWAY_BUSY")
|
|
self.assertTrue(raised.exception.retryable)
|
|
finally:
|
|
service.release.set()
|
|
worker.join(timeout=2)
|
|
self.assertFalse(worker.is_alive())
|
|
self.assertEqual(failures, [])
|
|
|
|
|
|
class BearerAuthTests(unittest.TestCase):
|
|
def test_config_is_fail_closed_and_secret_is_not_represented(self) -> None:
|
|
with patch.dict(os.environ, {"ARR_MCP_BEARER_TOKEN": TOKEN}, clear=False):
|
|
config = BearerAuthConfig.from_environment()
|
|
self.assertNotIn(TOKEN, repr(config))
|
|
with patch.dict(os.environ, {}, clear=True):
|
|
with self.assertRaises(ValueError):
|
|
BearerAuthConfig.from_environment()
|
|
|
|
def test_middleware_requires_one_exact_bearer_header(self) -> None:
|
|
async def inner(scope, receive, send):
|
|
del scope, receive
|
|
await send({"type": "http.response.start", "status": 204, "headers": []})
|
|
await send({"type": "http.response.body", "body": b""})
|
|
|
|
app = BearerAuthASGI(inner, BearerAuthConfig(TOKEN))
|
|
accepted = asyncio.run(
|
|
invoke_asgi(app, [(b"authorization", ("Bearer " + TOKEN).encode())])
|
|
)
|
|
self.assertEqual(accepted[0]["status"], 204)
|
|
rejected = asyncio.run(invoke_asgi(app, []))
|
|
self.assertEqual(rejected[0]["status"], 401)
|
|
duplicated = asyncio.run(
|
|
invoke_asgi(
|
|
app,
|
|
[
|
|
(b"authorization", ("Bearer " + TOKEN).encode()),
|
|
(b"authorization", ("Bearer " + TOKEN).encode()),
|
|
],
|
|
)
|
|
)
|
|
self.assertEqual(duplicated[0]["status"], 401)
|
|
|
|
|
|
class McpConfigurationTests(unittest.TestCase):
|
|
def test_advertised_schema_is_portable_and_server_contract_stays_exact(self) -> None:
|
|
request, response = load_tool_schemas()
|
|
self.assertEqual(
|
|
set(request["properties"]),
|
|
{"submission_grant", "job_id", "attempt_no", "payload"},
|
|
)
|
|
payload = request["properties"]["payload"]
|
|
self.assertEqual(
|
|
set(payload["required"]),
|
|
set(payload["properties"]),
|
|
)
|
|
self.assertEqual(payload["properties"]["status"]["const"], "success")
|
|
self.assertIs(
|
|
payload["properties"]["activation_eligible"]["const"],
|
|
True,
|
|
)
|
|
self.assertEqual(
|
|
payload["properties"]["records"]["items"],
|
|
{"type": "object"},
|
|
)
|
|
forbidden_keys: list[str] = []
|
|
|
|
def collect(value: object) -> None:
|
|
if isinstance(value, dict):
|
|
for key, nested in value.items():
|
|
if key in {"$ref", "$defs", "$id"}:
|
|
forbidden_keys.append(key)
|
|
collect(nested)
|
|
elif isinstance(value, list):
|
|
for nested in value:
|
|
collect(nested)
|
|
|
|
collect(request)
|
|
self.assertEqual(forbidden_keys, [])
|
|
self.assertEqual(
|
|
response["properties"]["contract_version"]["const"],
|
|
"arr-direct-ingestion-1",
|
|
)
|
|
|
|
def test_allowed_hosts_are_fail_closed_beyond_loopback(self) -> None:
|
|
self.assertIn("127.0.0.1:*", _allowed_hosts("127.0.0.1", []))
|
|
with self.assertRaises(ValueError):
|
|
_allowed_hosts("0.0.0.0", [])
|
|
self.assertEqual(
|
|
_allowed_hosts("0.0.0.0", ["mcp.example.test", "mcp.example.test"]),
|
|
["mcp.example.test"],
|
|
)
|
|
|
|
def test_database_config_requires_owned_mode_0600_file(self) -> None:
|
|
content = "\n".join(
|
|
[
|
|
"ARR_DB_HOST=127.0.0.1",
|
|
"ARR_DB_PORT=5432",
|
|
"ARR_DB_USER=synthetic",
|
|
"ARR_DB_PASSWORD=synthetic",
|
|
"ARR_DB_NAME=booking_test",
|
|
]
|
|
)
|
|
with tempfile.TemporaryDirectory() as temporary:
|
|
path = Path(temporary) / "database.env"
|
|
path.write_text(content, encoding="utf-8")
|
|
path.chmod(0o600)
|
|
values = _read_private_database_config(path)
|
|
self.assertEqual(values["dbname"], "booking_test")
|
|
path.chmod(0o644)
|
|
with self.assertRaises(ValueError):
|
|
_read_private_database_config(path)
|
|
path.chmod(stat.S_IRUSR | stat.S_IWUSR)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main(verbosity=2)
|