92 lines
3.1 KiB
Python
92 lines
3.1 KiB
Python
"""Local protocol fixtures; no devices, model calls or user data."""
|
|
|
|
import asyncio
|
|
from contextlib import asynccontextmanager
|
|
import os
|
|
import socket
|
|
import tempfile
|
|
import threading
|
|
import time
|
|
|
|
from fastapi import FastAPI, HTTPException, Request
|
|
import pytest
|
|
import uvicorn
|
|
|
|
from artemis.service_registry import ServiceRegistration, TOKEN_HEADER
|
|
from mcp_server.shared import LocalMcpBoundary, create_shared_mcp
|
|
|
|
|
|
@pytest.fixture
|
|
def shared_mcp_backend(tmp_path, monkeypatch):
|
|
monkeypatch.setattr(tempfile, "tempdir", str(tmp_path))
|
|
monkeypatch.setenv("TEMP", str(tmp_path))
|
|
monkeypatch.setenv("TMP", str(tmp_path))
|
|
monkeypatch.setenv("TMPDIR", str(tmp_path))
|
|
monkeypatch.setenv("ARTEMIS_APP_DIR", str(tmp_path / "app"))
|
|
monkeypatch.setenv("ARTEMIS_TRACES_DIR", str(tmp_path / "traces"))
|
|
monkeypatch.delenv("ARTEMIS_DAEMON_PORT", raising=False)
|
|
monkeypatch.delenv("ARTEMIS_DAEMON_HOST", raising=False)
|
|
listener = socket.socket()
|
|
listener.bind(("127.0.0.1", 0))
|
|
port = listener.getsockname()[1]
|
|
registration = ServiceRegistration(
|
|
port=port, token="fixture-token", app_dir=tmp_path / "app", traces_dir=tmp_path / "traces"
|
|
)
|
|
shared, executor = create_shared_mcp()
|
|
calls = []
|
|
|
|
async def record(delay: float):
|
|
calls.append(delay)
|
|
await asyncio.sleep(delay)
|
|
return {"pid": os.getpid(), "thread": threading.get_ident(), "count": len(calls)}
|
|
|
|
@shared.tool()
|
|
async def shared_probe(delay: float = 0.0) -> dict[str, int]:
|
|
return await executor.invoke(record, delay)
|
|
|
|
@asynccontextmanager
|
|
async def lifespan(app):
|
|
registration.acquire()
|
|
executor.start()
|
|
try:
|
|
async with shared.session_manager.run():
|
|
yield
|
|
finally:
|
|
await executor.stop()
|
|
registration.close()
|
|
|
|
app = FastAPI(lifespan=lifespan)
|
|
app.mount("/api/mcp", LocalMcpBoundary(shared.streamable_http_app(), "fixture-token"))
|
|
|
|
@app.get("/api/system/shared-service")
|
|
async def identity(request: Request):
|
|
if request.headers.get(TOKEN_HEADER) != "fixture-token":
|
|
raise HTTPException(403)
|
|
return {
|
|
key: registration.info[key] for key in ("protocol", "installation", "instance", "pid")
|
|
}
|
|
|
|
server = uvicorn.Server(uvicorn.Config(app, host="127.0.0.1", port=port, log_level="error"))
|
|
thread = threading.Thread(target=server.run, kwargs={"sockets": [listener]}, daemon=True)
|
|
thread.start()
|
|
deadline = time.monotonic() + 8
|
|
while not server.started and thread.is_alive() and time.monotonic() < deadline:
|
|
time.sleep(0.01)
|
|
assert server.started
|
|
try:
|
|
yield {
|
|
"port": port,
|
|
"registration": registration,
|
|
"server": server,
|
|
"calls": calls,
|
|
"env": {
|
|
key: os.environ[key]
|
|
for key in ("TEMP", "TMP", "TMPDIR", "ARTEMIS_APP_DIR", "ARTEMIS_TRACES_DIR")
|
|
},
|
|
}
|
|
finally:
|
|
server.should_exit = True
|
|
thread.join(timeout=8)
|
|
listener.close()
|
|
assert not thread.is_alive()
|