127 lines
5 KiB
Python
127 lines
5 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from pathlib import Path
|
|
from typing import Any, cast
|
|
|
|
import pytest
|
|
|
|
import pyro_mcp.server as server_module
|
|
from pyro_mcp.server import create_server
|
|
from pyro_mcp.vm_manager import VmManager
|
|
|
|
|
|
def test_create_server_registers_vm_tools(tmp_path: Path) -> None:
|
|
manager = VmManager(backend_name="mock", base_dir=tmp_path / "vms")
|
|
|
|
async def _run() -> list[str]:
|
|
server = create_server(manager=manager)
|
|
tools = await server.list_tools()
|
|
return sorted(tool.name for tool in tools)
|
|
|
|
tool_names = asyncio.run(_run())
|
|
assert "vm_create" in tool_names
|
|
assert "vm_exec" in tool_names
|
|
assert "vm_list_profiles" in tool_names
|
|
assert "vm_status" in tool_names
|
|
|
|
|
|
def test_vm_tools_lifecycle_round_trip(tmp_path: Path) -> None:
|
|
manager = VmManager(backend_name="mock", base_dir=tmp_path / "vms")
|
|
|
|
def _extract_structured(raw_result: object) -> dict[str, Any]:
|
|
if not isinstance(raw_result, tuple) or len(raw_result) != 2:
|
|
raise TypeError("unexpected call_tool result shape")
|
|
_, structured = raw_result
|
|
if not isinstance(structured, dict):
|
|
raise TypeError("expected structured dictionary result")
|
|
return cast(dict[str, Any], structured)
|
|
|
|
async def _run() -> dict[str, Any]:
|
|
server = create_server(manager=manager)
|
|
created = _extract_structured(
|
|
await server.call_tool(
|
|
"vm_create",
|
|
{"profile": "debian-git", "vcpu_count": 1, "mem_mib": 512, "ttl_seconds": 600},
|
|
)
|
|
)
|
|
vm_id = str(created["vm_id"])
|
|
await server.call_tool("vm_start", {"vm_id": vm_id})
|
|
executed = _extract_structured(
|
|
await server.call_tool(
|
|
"vm_exec", {"vm_id": vm_id, "command": "printf 'git version 2.0\\n'"}
|
|
)
|
|
)
|
|
return executed
|
|
|
|
executed = asyncio.run(_run())
|
|
assert int(executed["exit_code"]) == 0
|
|
assert "git version" in str(executed["stdout"])
|
|
|
|
|
|
def test_vm_tools_status_stop_delete_and_reap(tmp_path: Path) -> None:
|
|
manager = VmManager(backend_name="mock", base_dir=tmp_path / "vms")
|
|
manager.MIN_TTL_SECONDS = 1
|
|
|
|
def _extract_structured(raw_result: object) -> dict[str, Any]:
|
|
if not isinstance(raw_result, tuple) or len(raw_result) != 2:
|
|
raise TypeError("unexpected call_tool result shape")
|
|
_, structured = raw_result
|
|
if not isinstance(structured, dict):
|
|
raise TypeError("expected structured dictionary result")
|
|
return cast(dict[str, Any], structured)
|
|
|
|
async def _run() -> tuple[
|
|
dict[str, Any], dict[str, Any], dict[str, Any], list[dict[str, object]], dict[str, Any]
|
|
]:
|
|
server = create_server(manager=manager)
|
|
profiles_raw = await server.call_tool("vm_list_profiles", {})
|
|
if not isinstance(profiles_raw, tuple) or len(profiles_raw) != 2:
|
|
raise TypeError("unexpected profiles result")
|
|
_, profiles_structured = profiles_raw
|
|
if not isinstance(profiles_structured, dict):
|
|
raise TypeError("profiles tool should return a dictionary")
|
|
raw_profiles = profiles_structured.get("result")
|
|
if not isinstance(raw_profiles, list):
|
|
raise TypeError("profiles tool did not contain a result list")
|
|
created = _extract_structured(
|
|
await server.call_tool(
|
|
"vm_create",
|
|
{"profile": "debian-base", "vcpu_count": 1, "mem_mib": 512, "ttl_seconds": 600},
|
|
)
|
|
)
|
|
vm_id = str(created["vm_id"])
|
|
await server.call_tool("vm_start", {"vm_id": vm_id})
|
|
status = _extract_structured(await server.call_tool("vm_status", {"vm_id": vm_id}))
|
|
stopped = _extract_structured(await server.call_tool("vm_stop", {"vm_id": vm_id}))
|
|
deleted = _extract_structured(await server.call_tool("vm_delete", {"vm_id": vm_id}))
|
|
|
|
expiring = _extract_structured(
|
|
await server.call_tool(
|
|
"vm_create",
|
|
{"profile": "debian-base", "vcpu_count": 1, "mem_mib": 512, "ttl_seconds": 1},
|
|
)
|
|
)
|
|
expiring_id = str(expiring["vm_id"])
|
|
manager._instances[expiring_id].expires_at = 0.0 # noqa: SLF001
|
|
reaped = _extract_structured(await server.call_tool("vm_reap_expired", {}))
|
|
return status, stopped, deleted, cast(list[dict[str, object]], raw_profiles), reaped
|
|
|
|
status, stopped, deleted, profiles, reaped = asyncio.run(_run())
|
|
assert status["state"] == "started"
|
|
assert stopped["state"] == "stopped"
|
|
assert bool(deleted["deleted"]) is True
|
|
assert profiles[0]["name"] == "debian-base"
|
|
assert int(reaped["count"]) == 1
|
|
|
|
|
|
def test_server_main_runs_stdio_transport(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
called: dict[str, str] = {}
|
|
|
|
class StubServer:
|
|
def run(self, transport: str) -> None:
|
|
called["transport"] = transport
|
|
|
|
monkeypatch.setattr(server_module, "create_server", lambda: StubServer())
|
|
server_module.main()
|
|
assert called == {"transport": "stdio"}
|