Bundle firecracker runtime and switch ollama demo to live logs

tests/test_demo.py

@@ -1,86 +1,71 @@
 from __future__ import annotations
 
-import asyncio
-from collections.abc import Sequence
+import json
 from typing import Any
 
 import pytest
-from mcp.types import TextContent
 
 import pyro_mcp.demo as demo_module
-from pyro_mcp.demo import run_demo
-from pyro_mcp.server import HELLO_STATIC_PAYLOAD
 
 
-def test_run_demo_returns_static_payload() -> None:
-    payload = asyncio.run(run_demo())
-    assert payload == HELLO_STATIC_PAYLOAD
+def test_run_demo_happy_path(monkeypatch: pytest.MonkeyPatch) -> None:
+    calls: list[tuple[str, dict[str, Any]]] = []
 
+    class StubManager:
+        def __init__(self) -> None:
+            pass
 
-def test_run_demo_raises_for_non_text_blocks(monkeypatch: pytest.MonkeyPatch) -> None:
-    class StubServer:
-        async def call_tool(
+        def create_vm(
             self,
-            name: str,
-            arguments: dict[str, Any],
-        ) -> tuple[Sequence[int], dict[str, str]]:
-            assert name == "hello_static"
-            assert arguments == {}
-            return [123], HELLO_STATIC_PAYLOAD
+            *,
+            profile: str,
+            vcpu_count: int,
+            mem_mib: int,
+            ttl_seconds: int,
+        ) -> dict[str, str]:
+            calls.append(
+                (
+                    "create_vm",
+                    {
+                        "profile": profile,
+                        "vcpu_count": vcpu_count,
+                        "mem_mib": mem_mib,
+                        "ttl_seconds": ttl_seconds,
+                    },
+                )
+            )
+            return {"vm_id": "vm-1"}
 
-    monkeypatch.setattr(demo_module, "create_server", lambda: StubServer())
+        def start_vm(self, vm_id: str) -> dict[str, str]:
+            calls.append(("start_vm", {"vm_id": vm_id}))
+            return {"vm_id": vm_id}
 
-    with pytest.raises(TypeError, match="unexpected MCP content block output"):
-        asyncio.run(demo_module.run_demo())
+        def exec_vm(self, vm_id: str, *, command: str, timeout_seconds: int) -> dict[str, Any]:
+            calls.append(
+                (
+                    "exec_vm",
+                    {"vm_id": vm_id, "command": command, "timeout_seconds": timeout_seconds},
+                )
+            )
+            return {"vm_id": vm_id, "stdout": "git version 2.x", "exit_code": 0}
+
+    monkeypatch.setattr(demo_module, "VmManager", StubManager)
+    result = demo_module.run_demo()
+
+    assert result["exit_code"] == 0
+    assert calls[0][0] == "create_vm"
+    assert calls[1] == ("start_vm", {"vm_id": "vm-1"})
+    assert calls[2][0] == "exec_vm"
 
 
-def test_run_demo_raises_for_non_dict_payload(monkeypatch: pytest.MonkeyPatch) -> None:
-    class StubServer:
-        async def call_tool(
-            self,
-            name: str,
-            arguments: dict[str, Any],
-        ) -> tuple[list[TextContent], str]:
-            assert name == "hello_static"
-            assert arguments == {}
-            return [TextContent(type="text", text="x")], "bad"
-
-    monkeypatch.setattr(demo_module, "create_server", lambda: StubServer())
-
-    with pytest.raises(TypeError, match="expected a structured dictionary payload"):
-        asyncio.run(demo_module.run_demo())
-
-
-def test_run_demo_raises_for_unexpected_payload(monkeypatch: pytest.MonkeyPatch) -> None:
-    class StubServer:
-        async def call_tool(
-            self,
-            name: str,
-            arguments: dict[str, Any],
-        ) -> tuple[list[TextContent], dict[str, str]]:
-            assert name == "hello_static"
-            assert arguments == {}
-            return [TextContent(type="text", text="x")], {
-                "message": "different",
-                "status": "ok",
-                "version": "0.0.1",
-            }
-
-    monkeypatch.setattr(demo_module, "create_server", lambda: StubServer())
-
-    with pytest.raises(ValueError, match="static payload did not match expected value"):
-        asyncio.run(demo_module.run_demo())
-
-
-def test_demo_main_prints_json(
-    monkeypatch: pytest.MonkeyPatch,
-    capsys: pytest.CaptureFixture[str],
+def test_main_prints_json(
+    monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str]
 ) -> None:
-    async def fake_run_demo() -> dict[str, str]:
-        return HELLO_STATIC_PAYLOAD
-
-    monkeypatch.setattr(demo_module, "run_demo", fake_run_demo)
+    monkeypatch.setattr(
+        demo_module,
+        "run_demo",
+        lambda: {"stdout": "git version 2.x", "exit_code": 0},
+    )
     demo_module.main()
-
-    output = capsys.readouterr().out
-    assert '"message": "hello from pyro_mcp"' in output
+    rendered = json.loads(capsys.readouterr().out)
+    assert rendered["exit_code"] == 0