pyro-mcp/tests/test_server.py
Thales Maciel 287f6d100f Add stopped-workspace disk export and inspection
Finish the 3.1.0 secondary disk-tools milestone so stable workspaces can be
stopped, inspected offline, exported as raw ext4 images, and started again
without changing the primary workspace-first interaction model.

Add workspace stop/start plus workspace disk export/list/read across the CLI,
SDK, and MCP, backed by a new offline debugfs inspection helper and guest-only
validation. Scrub runtime-only guest state before disk inspection/export, and
fix the real guest reliability gaps by flushing the filesystem on stop and
removing stale Firecracker socket files before restart.

Update the docs, examples, changelog, and roadmap to mark 3.1.0 done, and
cover the new lifecycle/disk paths with API, CLI, manager, contract, and
package-surface tests.

Validation: uv lock; UV_CACHE_DIR=.uv-cache make check; UV_CACHE_DIR=.uv-cache
make dist-check; real guest-backed smoke for create, shell/service activity,
stop, workspace disk list/read/export, start, exec, and delete.
2026-03-12 20:57:16 -03:00

402 lines
14 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
from pyro_mcp.vm_network import TapNetworkManager
def test_create_server_registers_vm_tools(tmp_path: Path) -> None:
manager = VmManager(
backend_name="mock",
base_dir=tmp_path / "vms",
network_manager=TapNetworkManager(enabled=False),
)
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_environments" in tool_names
assert "vm_network_info" in tool_names
assert "vm_run" in tool_names
assert "vm_status" in tool_names
assert "workspace_create" in tool_names
assert "workspace_start" in tool_names
assert "workspace_stop" in tool_names
assert "workspace_diff" in tool_names
assert "workspace_export" in tool_names
assert "workspace_disk_export" in tool_names
assert "workspace_disk_list" in tool_names
assert "workspace_disk_read" in tool_names
assert "workspace_logs" in tool_names
assert "workspace_sync_push" in tool_names
assert "shell_open" in tool_names
assert "shell_read" in tool_names
assert "shell_write" in tool_names
assert "shell_signal" in tool_names
assert "shell_close" in tool_names
assert "service_start" in tool_names
assert "service_list" in tool_names
assert "service_status" in tool_names
assert "service_logs" in tool_names
assert "service_stop" in tool_names
assert "snapshot_create" in tool_names
assert "snapshot_delete" in tool_names
assert "snapshot_list" in tool_names
assert "workspace_reset" in tool_names
def test_vm_run_round_trip(tmp_path: Path) -> None:
manager = VmManager(
backend_name="mock",
base_dir=tmp_path / "vms",
network_manager=TapNetworkManager(enabled=False),
)
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)
executed = _extract_structured(
await server.call_tool(
"vm_run",
{
"environment": "debian:12",
"command": "printf 'git version 2.0\\n'",
"ttl_seconds": 600,
"network": False,
"allow_host_compat": True,
},
)
)
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",
network_manager=TapNetworkManager(enabled=False),
)
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],
dict[str, Any],
list[dict[str, object]],
dict[str, Any],
]:
server = create_server(manager=manager)
environments_raw = await server.call_tool("vm_list_environments", {})
if not isinstance(environments_raw, tuple) or len(environments_raw) != 2:
raise TypeError("unexpected environments result")
_, environments_structured = environments_raw
if not isinstance(environments_structured, dict):
raise TypeError("environments tool should return a dictionary")
raw_environments = environments_structured.get("result")
if not isinstance(raw_environments, list):
raise TypeError("environments tool did not contain a result list")
created = _extract_structured(
await server.call_tool(
"vm_create",
{
"environment": "debian:12-base",
"ttl_seconds": 600,
"allow_host_compat": True,
},
)
)
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}))
network = _extract_structured(await server.call_tool("vm_network_info", {"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",
{
"environment": "debian:12-base",
"ttl_seconds": 1,
"allow_host_compat": True,
},
)
)
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,
network,
stopped,
deleted,
cast(list[dict[str, object]], raw_environments),
reaped,
)
status, network, stopped, deleted, environments, reaped = asyncio.run(_run())
assert status["state"] == "started"
assert network["network_enabled"] is False
assert stopped["state"] == "stopped"
assert bool(deleted["deleted"]) is True
assert environments[0]["name"] == "debian:12"
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"}
def test_workspace_tools_round_trip(tmp_path: Path) -> None:
manager = VmManager(
backend_name="mock",
base_dir=tmp_path / "vms",
network_manager=TapNetworkManager(enabled=False),
)
source_dir = tmp_path / "seed"
source_dir.mkdir()
(source_dir / "note.txt").write_text("ok\n", encoding="utf-8")
secret_file = tmp_path / "token.txt"
secret_file.write_text("from-file\n", encoding="utf-8")
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], ...]:
server = create_server(manager=manager)
created = _extract_structured(
await server.call_tool(
"workspace_create",
{
"environment": "debian:12-base",
"allow_host_compat": True,
"seed_path": str(source_dir),
"secrets": [
{"name": "API_TOKEN", "value": "expected"},
{"name": "FILE_TOKEN", "file_path": str(secret_file)},
],
},
)
)
workspace_id = str(created["workspace_id"])
update_dir = tmp_path / "update"
update_dir.mkdir()
(update_dir / "more.txt").write_text("more\n", encoding="utf-8")
synced = _extract_structured(
await server.call_tool(
"workspace_sync_push",
{
"workspace_id": workspace_id,
"source_path": str(update_dir),
"dest": "subdir",
},
)
)
executed = _extract_structured(
await server.call_tool(
"workspace_exec",
{
"workspace_id": workspace_id,
"command": 'sh -lc \'printf "%s\\n" "$API_TOKEN"\'',
"secret_env": {"API_TOKEN": "API_TOKEN"},
},
)
)
diffed = _extract_structured(
await server.call_tool("workspace_diff", {"workspace_id": workspace_id})
)
snapshot = _extract_structured(
await server.call_tool(
"snapshot_create",
{"workspace_id": workspace_id, "snapshot_name": "checkpoint"},
)
)
snapshots = _extract_structured(
await server.call_tool("snapshot_list", {"workspace_id": workspace_id})
)
export_path = tmp_path / "exported-more.txt"
exported = _extract_structured(
await server.call_tool(
"workspace_export",
{
"workspace_id": workspace_id,
"path": "subdir/more.txt",
"output_path": str(export_path),
},
)
)
service = _extract_structured(
await server.call_tool(
"service_start",
{
"workspace_id": workspace_id,
"service_name": "app",
"command": (
'sh -lc \'trap "exit 0" TERM; printf "%s\\n" "$API_TOKEN" >&2; '
'touch .ready; while true; do sleep 60; done\''
),
"ready_file": ".ready",
"secret_env": {"API_TOKEN": "API_TOKEN"},
},
)
)
services = _extract_structured(
await server.call_tool("service_list", {"workspace_id": workspace_id})
)
service_status = _extract_structured(
await server.call_tool(
"service_status",
{
"workspace_id": workspace_id,
"service_name": "app",
},
)
)
service_logs = _extract_structured(
await server.call_tool(
"service_logs",
{
"workspace_id": workspace_id,
"service_name": "app",
"all": True,
},
)
)
service_stopped = _extract_structured(
await server.call_tool(
"service_stop",
{
"workspace_id": workspace_id,
"service_name": "app",
},
)
)
reset = _extract_structured(
await server.call_tool(
"workspace_reset",
{"workspace_id": workspace_id, "snapshot": "checkpoint"},
)
)
deleted_snapshot = _extract_structured(
await server.call_tool(
"snapshot_delete",
{"workspace_id": workspace_id, "snapshot_name": "checkpoint"},
)
)
logs = _extract_structured(
await server.call_tool("workspace_logs", {"workspace_id": workspace_id})
)
deleted = _extract_structured(
await server.call_tool("workspace_delete", {"workspace_id": workspace_id})
)
return (
created,
synced,
executed,
diffed,
snapshot,
snapshots,
exported,
service,
services,
service_status,
service_logs,
service_stopped,
reset,
deleted_snapshot,
logs,
deleted,
)
(
created,
synced,
executed,
diffed,
snapshot,
snapshots,
exported,
service,
services,
service_status,
service_logs,
service_stopped,
reset,
deleted_snapshot,
logs,
deleted,
) = asyncio.run(_run())
assert created["state"] == "started"
assert created["workspace_seed"]["mode"] == "directory"
assert created["secrets"] == [
{"name": "API_TOKEN", "source_kind": "literal"},
{"name": "FILE_TOKEN", "source_kind": "file"},
]
assert synced["workspace_sync"]["destination"] == "/workspace/subdir"
assert executed["stdout"] == "[REDACTED]\n"
assert diffed["changed"] is True
assert snapshot["snapshot"]["snapshot_name"] == "checkpoint"
assert [entry["snapshot_name"] for entry in snapshots["snapshots"]] == [
"baseline",
"checkpoint",
]
assert exported["artifact_type"] == "file"
assert Path(str(exported["output_path"])).read_text(encoding="utf-8") == "more\n"
assert service["state"] == "running"
assert services["count"] == 1
assert service_status["state"] == "running"
assert service_logs["stderr"].count("[REDACTED]") >= 1
assert service_logs["tail_lines"] is None
assert service_stopped["state"] == "stopped"
assert reset["workspace_reset"]["snapshot_name"] == "checkpoint"
assert reset["secrets"] == created["secrets"]
assert reset["command_count"] == 0
assert reset["service_count"] == 0
assert deleted_snapshot["deleted"] is True
assert logs["count"] == 0
assert deleted["deleted"] is True