Bundle firecracker runtime and switch ollama demo to live logs

tests/test_server.py

@@ -1,36 +1,118 @@
 from __future__ import annotations
 
 import asyncio
-from typing import Any
+from pathlib import Path
+from typing import Any, cast
 
 import pytest
-from mcp.types import TextContent
 
 import pyro_mcp.server as server_module
-from pyro_mcp.server import HELLO_STATIC_PAYLOAD, create_server
+from pyro_mcp.server import create_server
+from pyro_mcp.vm_manager import VmManager
 
 
-def test_create_server_registers_static_tool() -> None:
+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()
+        server = create_server(manager=manager)
         tools = await server.list_tools()
-        return [tool.name for tool in tools]
+        return sorted(tool.name for tool in tools)
 
     tool_names = asyncio.run(_run())
-    assert "hello_static" in tool_names
+    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_hello_static_returns_expected_payload() -> None:
-    async def _run() -> tuple[list[TextContent], dict[str, Any]]:
-        server = create_server()
-        blocks, structured = await server.call_tool("hello_static", {})
-        assert isinstance(blocks, list)
-        assert all(isinstance(block, TextContent) for block in blocks)
-        assert isinstance(structured, dict)
-        return blocks, structured
+def test_vm_tools_lifecycle_round_trip(tmp_path: Path) -> None:
+    manager = VmManager(backend_name="mock", base_dir=tmp_path / "vms")
 
-    _, structured_output = asyncio.run(_run())
-    assert structured_output == HELLO_STATIC_PAYLOAD
+    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:
@@ -42,5 +124,4 @@ def test_server_main_runs_stdio_transport(monkeypatch: pytest.MonkeyPatch) -> No
 
     monkeypatch.setattr(server_module, "create_server", lambda: StubServer())
     server_module.main()
-
     assert called == {"transport": "stdio"}