pyro-mcp/tests/test_ollama_demo.py

716 lines
26 KiB
Python

from __future__ import annotations
import argparse
import json
import urllib.error
import urllib.request
from pathlib import Path
from typing import Any
import pytest
import pyro_mcp.ollama_demo as ollama_demo
from pyro_mcp.api import Pyro as RealPyro
from pyro_mcp.vm_manager import VmManager as RealVmManager
@pytest.fixture(autouse=True)
def _mock_vm_manager_for_tests(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None:
class TestPyro(RealPyro):
def __init__(self) -> None:
super().__init__(manager=RealVmManager(backend_name="mock", base_dir=tmp_path / "vms"))
monkeypatch.setattr(ollama_demo, "Pyro", TestPyro)
def _stepwise_model_response(payload: dict[str, Any], step: int) -> dict[str, Any]:
if step == 1:
return {
"choices": [
{
"message": {
"role": "assistant",
"content": "",
"tool_calls": [{"id": "1", "function": {"name": "vm_list_environments"}}],
}
}
]
}
if step == 2:
return {
"choices": [
{
"message": {
"role": "assistant",
"content": "",
"tool_calls": [
{
"id": "2",
"function": {
"name": "vm_run",
"arguments": json.dumps(
{
"environment": "debian:12",
"command": "printf 'true\\n'",
"network": True,
"allow_host_compat": True,
}
),
},
}
],
}
}
]
}
return {
"choices": [
{"message": {"role": "assistant", "content": "Executed git command in ephemeral VM."}}
]
}
def test_run_ollama_tool_demo_happy_path(monkeypatch: pytest.MonkeyPatch) -> None:
requests: list[dict[str, Any]] = []
def fake_post_chat_completion(base_url: str, payload: dict[str, Any]) -> dict[str, Any]:
assert base_url == "http://localhost:11434/v1"
requests.append(payload)
return _stepwise_model_response(payload, len(requests))
monkeypatch.setattr(ollama_demo, "_post_chat_completion", fake_post_chat_completion)
logs: list[str] = []
result = ollama_demo.run_ollama_tool_demo(log=logs.append)
assert result["fallback_used"] is False
assert str(result["exec_result"]["stdout"]).strip() == "true"
assert result["final_response"] == "Executed git command in ephemeral VM."
assert len(result["tool_events"]) == 2
assert any(line == "[model] input user" for line in logs)
assert any(line == "[model] output assistant" for line in logs)
assert any("[model] tool_call vm_run" in line for line in logs)
assert any(line == "[tool] calling vm_run" for line in logs)
assert any(line == "[tool] result vm_run" for line in logs)
def test_run_ollama_tool_demo_accepts_legacy_profile_and_string_network(
monkeypatch: pytest.MonkeyPatch,
) -> None:
requests: list[dict[str, Any]] = []
def fake_post_chat_completion(base_url: str, payload: dict[str, Any]) -> dict[str, Any]:
assert base_url == "http://localhost:11434/v1"
requests.append(payload)
if len(requests) == 1:
return {
"choices": [
{
"message": {
"role": "assistant",
"content": "",
"tool_calls": [
{
"id": "1",
"function": {
"name": "vm_run",
"arguments": json.dumps(
{
"profile": "debian:12",
"command": "printf 'true\\n'",
"network": "true",
"allow_host_compat": True,
}
),
},
}
],
}
}
]
}
return {
"choices": [
{
"message": {
"role": "assistant",
"content": "Executed git command in ephemeral VM.",
}
}
]
}
monkeypatch.setattr(ollama_demo, "_post_chat_completion", fake_post_chat_completion)
result = ollama_demo.run_ollama_tool_demo()
assert result["fallback_used"] is False
assert int(result["exec_result"]["exit_code"]) == 0
assert result["tool_events"][0]["success"] is True
assert result["tool_events"][0]["arguments"]["environment"] == "debian:12"
assert "profile" not in result["tool_events"][0]["arguments"]
def test_run_ollama_tool_demo_recovers_from_bad_vm_id(
monkeypatch: pytest.MonkeyPatch,
) -> None:
requests: list[dict[str, Any]] = []
def fake_post_chat_completion(base_url: str, payload: dict[str, Any]) -> dict[str, Any]:
assert base_url == "http://localhost:11434/v1"
requests.append(payload)
step = len(requests)
if step == 1:
return {
"choices": [
{
"message": {
"role": "assistant",
"tool_calls": [
{
"id": "1",
"function": {
"name": "vm_exec",
"arguments": json.dumps(
{
"vm_id": "vm_list_environments",
"command": ollama_demo.NETWORK_PROOF_COMMAND,
}
),
},
}
],
}
}
]
}
return _stepwise_model_response(payload, step - 1)
monkeypatch.setattr(ollama_demo, "_post_chat_completion", fake_post_chat_completion)
result = ollama_demo.run_ollama_tool_demo()
first_event = result["tool_events"][0]
assert first_event["tool_name"] == "vm_exec"
assert first_event["success"] is False
assert "does not exist" in str(first_event["result"]["error"])
assert int(result["exec_result"]["exit_code"]) == 0
def test_run_ollama_tool_demo_resolves_vm_id_placeholder(
monkeypatch: pytest.MonkeyPatch,
) -> None:
requests: list[dict[str, Any]] = []
def fake_post_chat_completion(base_url: str, payload: dict[str, Any]) -> dict[str, Any]:
assert base_url == "http://localhost:11434/v1"
requests.append(payload)
step = len(requests)
if step == 1:
return {
"choices": [
{
"message": {
"role": "assistant",
"content": "",
"tool_calls": [
{"id": "1", "function": {"name": "vm_list_environments"}},
{"id": "2", "function": {"name": "vm_list_environments"}},
{
"id": "3",
"function": {
"name": "vm_create",
"arguments": json.dumps(
{
"environment": "debian:12",
"allow_host_compat": True,
}
),
},
},
{
"id": "4",
"function": {
"name": "vm_exec",
"arguments": json.dumps(
{
"vm_id": "<vm_id_returned_by_vm_create>",
"command": "printf 'true\\n'",
"timeout_seconds": "300",
}
),
},
},
],
}
}
]
}
return {
"choices": [
{
"message": {
"role": "assistant",
"content": "Executed git command in ephemeral VM.",
}
}
]
}
monkeypatch.setattr(ollama_demo, "_post_chat_completion", fake_post_chat_completion)
logs: list[str] = []
result = ollama_demo.run_ollama_tool_demo(log=logs.append)
assert result["fallback_used"] is False
assert int(result["exec_result"]["exit_code"]) == 0
assert any("resolved vm_id placeholder" in line for line in logs)
exec_event = next(
event for event in result["tool_events"] if event["tool_name"] == "vm_exec"
)
assert exec_event["success"] is True
def test_dispatch_tool_call_vm_exec_autostarts_created_vm(tmp_path: Path) -> None:
pyro = RealPyro(manager=RealVmManager(backend_name="mock", base_dir=tmp_path / "vms"))
created = pyro.create_vm(
environment="debian:12-base",
vcpu_count=1,
mem_mib=512,
ttl_seconds=60,
allow_host_compat=True,
)
vm_id = str(created["vm_id"])
executed = ollama_demo._dispatch_tool_call(
pyro,
"vm_exec",
{"vm_id": vm_id, "command": "printf 'git version\\n'", "timeout_seconds": "30"},
)
assert int(executed["exit_code"]) == 0
def test_run_ollama_tool_demo_raises_without_vm_exec(monkeypatch: pytest.MonkeyPatch) -> None:
def fake_post_chat_completion(base_url: str, payload: dict[str, Any]) -> dict[str, Any]:
del base_url, payload
return {"choices": [{"message": {"role": "assistant", "content": "No tools"}}]}
monkeypatch.setattr(ollama_demo, "_post_chat_completion", fake_post_chat_completion)
with pytest.raises(RuntimeError, match="did not execute a successful vm_run or vm_exec"):
ollama_demo.run_ollama_tool_demo()
def test_run_ollama_tool_demo_uses_fallback_when_not_strict(
monkeypatch: pytest.MonkeyPatch, tmp_path: Path
) -> None:
def fake_post_chat_completion(base_url: str, payload: dict[str, Any]) -> dict[str, Any]:
del base_url, payload
return {"choices": [{"message": {"role": "assistant", "content": "No tools"}}]}
class TestPyro(RealPyro):
def __init__(self) -> None:
super().__init__(manager=RealVmManager(backend_name="mock", base_dir=tmp_path / "vms"))
monkeypatch.setattr(ollama_demo, "_post_chat_completion", fake_post_chat_completion)
monkeypatch.setattr(ollama_demo, "Pyro", TestPyro)
monkeypatch.setattr(
ollama_demo,
"_run_direct_lifecycle_fallback",
lambda pyro: {
"vm_id": "vm-1",
"command": ollama_demo.NETWORK_PROOF_COMMAND,
"stdout": "true\n",
"stderr": "",
"exit_code": 0,
"duration_ms": 5,
"execution_mode": "host_compat",
"cleanup": {"deleted": True, "reason": "post_exec_cleanup", "vm_id": "vm-1"},
},
)
logs: list[str] = []
result = ollama_demo.run_ollama_tool_demo(strict=False, log=logs.append)
assert result["fallback_used"] is True
assert int(result["exec_result"]["exit_code"]) == 0
assert any(line == "[model] output assistant" for line in logs)
assert any("[fallback]" in line for line in logs)
def test_run_ollama_tool_demo_verbose_logs_values(monkeypatch: pytest.MonkeyPatch) -> None:
requests = 0
def fake_post_chat_completion(base_url: str, payload: dict[str, Any]) -> dict[str, Any]:
del base_url
nonlocal requests
requests += 1
return _stepwise_model_response(payload, requests)
monkeypatch.setattr(ollama_demo, "_post_chat_completion", fake_post_chat_completion)
logs: list[str] = []
result = ollama_demo.run_ollama_tool_demo(verbose=True, log=logs.append)
assert result["fallback_used"] is False
assert str(result["exec_result"]["stdout"]).strip() == "true"
assert any("[model] input user:" in line for line in logs)
assert any("[model] tool_call vm_list_environments args={}" in line for line in logs)
assert any("[tool] result vm_run " in line for line in logs)
@pytest.mark.parametrize(
("tool_call", "error"),
[
(1, "invalid tool call entry"),
({"id": "", "function": {"name": "vm_list_environments"}}, "valid call id"),
({"id": "1"}, "function metadata"),
({"id": "1", "function": {"name": 3}}, "name is invalid"),
],
)
def test_run_ollama_tool_demo_tool_call_validation(
monkeypatch: pytest.MonkeyPatch,
tool_call: Any,
error: str,
) -> None:
def fake_post_chat_completion(base_url: str, payload: dict[str, Any]) -> dict[str, Any]:
del base_url, payload
return {"choices": [{"message": {"role": "assistant", "tool_calls": [tool_call]}}]}
monkeypatch.setattr(ollama_demo, "_post_chat_completion", fake_post_chat_completion)
with pytest.raises(RuntimeError, match=error):
ollama_demo.run_ollama_tool_demo()
def test_run_ollama_tool_demo_max_rounds(monkeypatch: pytest.MonkeyPatch) -> None:
def fake_post_chat_completion(base_url: str, payload: dict[str, Any]) -> dict[str, Any]:
del base_url, payload
return {
"choices": [
{
"message": {
"role": "assistant",
"tool_calls": [{"id": "1", "function": {"name": "vm_list_environments"}}],
}
}
]
}
monkeypatch.setattr(ollama_demo, "_post_chat_completion", fake_post_chat_completion)
with pytest.raises(RuntimeError, match="exceeded maximum rounds"):
ollama_demo.run_ollama_tool_demo()
@pytest.mark.parametrize(
("exec_result", "error"),
[
("bad", "result shape is invalid"),
({"exit_code": 1, "stdout": "true"}, "expected exit_code=0"),
({"exit_code": 0, "stdout": "false"}, "did not confirm repository clone success"),
],
)
def test_run_ollama_tool_demo_exec_result_validation(
monkeypatch: pytest.MonkeyPatch,
exec_result: Any,
error: str,
) -> None:
responses: list[dict[str, Any]] = [
{
"choices": [
{
"message": {
"role": "assistant",
"tool_calls": [
{"id": "1", "function": {"name": "vm_run", "arguments": "{}"}}
],
}
}
]
},
{"choices": [{"message": {"role": "assistant", "content": "done"}}]},
]
def fake_post_chat_completion(base_url: str, payload: dict[str, Any]) -> dict[str, Any]:
del base_url, payload
return responses.pop(0)
def fake_dispatch(pyro: Any, tool_name: str, arguments: dict[str, Any]) -> Any:
del pyro, arguments
if tool_name == "vm_run":
return exec_result
return {"ok": True}
monkeypatch.setattr(ollama_demo, "_post_chat_completion", fake_post_chat_completion)
monkeypatch.setattr(ollama_demo, "_dispatch_tool_call", fake_dispatch)
with pytest.raises(RuntimeError, match=error):
ollama_demo.run_ollama_tool_demo()
def test_dispatch_tool_call_coverage(tmp_path: Path) -> None:
pyro = RealPyro(manager=RealVmManager(backend_name="mock", base_dir=tmp_path / "vms"))
environments = ollama_demo._dispatch_tool_call(pyro, "vm_list_environments", {})
assert "environments" in environments
created = ollama_demo._dispatch_tool_call(
pyro,
"vm_create",
{
"environment": "debian:12-base",
"vcpu_count": "1",
"mem_mib": "512",
"ttl_seconds": "60",
"network": False,
"allow_host_compat": True,
},
)
vm_id = str(created["vm_id"])
started = ollama_demo._dispatch_tool_call(pyro, "vm_start", {"vm_id": vm_id})
assert started["state"] == "started"
status = ollama_demo._dispatch_tool_call(pyro, "vm_status", {"vm_id": vm_id})
assert status["vm_id"] == vm_id
executed = ollama_demo._dispatch_tool_call(
pyro,
"vm_exec",
{"vm_id": vm_id, "command": "printf 'true\\n'", "timeout_seconds": "30"},
)
assert int(executed["exit_code"]) == 0
executed_run = ollama_demo._dispatch_tool_call(
pyro,
"vm_run",
{
"environment": "debian:12-base",
"command": "printf 'true\\n'",
"timeout_seconds": "30",
"network": False,
"allow_host_compat": True,
},
)
assert int(executed_run["exit_code"]) == 0
with pytest.raises(RuntimeError, match="unexpected tool requested by model"):
ollama_demo._dispatch_tool_call(pyro, "nope", {})
def test_format_tool_error() -> None:
error = ValueError("bad args")
result = ollama_demo._format_tool_error("vm_exec", {"vm_id": "x"}, error)
assert result["ok"] is False
assert result["error_type"] == "ValueError"
@pytest.mark.parametrize(
("arguments", "error"),
[
({}, "must be a non-empty string"),
({"k": 3}, "must be a non-empty string"),
],
)
def test_require_str(arguments: dict[str, Any], error: str) -> None:
with pytest.raises(ValueError, match=error):
ollama_demo._require_str(arguments, "k")
@pytest.mark.parametrize(
("arguments", "expected"),
[
({"k": 1}, 1),
({"k": "1"}, 1),
({"k": " 42 "}, 42),
],
)
def test_require_int_accepts_numeric_strings(arguments: dict[str, Any], expected: int) -> None:
assert ollama_demo._require_int(arguments, "k") == expected
@pytest.mark.parametrize("value", ["", "abc", "2.5", True])
def test_require_int_validation(value: Any) -> None:
with pytest.raises(ValueError, match="must be an integer"):
ollama_demo._require_int({"k": value}, "k")
@pytest.mark.parametrize(
("arguments", "expected"),
[
({}, False),
({"k": True}, True),
({"k": "true"}, True),
({"k": " false "}, False),
({"k": 1}, True),
({"k": 0}, False),
],
)
def test_require_bool(arguments: dict[str, Any], expected: bool) -> None:
assert ollama_demo._require_bool(arguments, "k", default=False) is expected
@pytest.mark.parametrize("value", ["", "maybe", 2])
def test_require_bool_validation(value: Any) -> None:
with pytest.raises(ValueError, match="must be a boolean"):
ollama_demo._require_bool({"k": value}, "k")
def test_post_chat_completion_success(monkeypatch: pytest.MonkeyPatch) -> None:
class StubResponse:
def __enter__(self) -> StubResponse:
return self
def __exit__(self, exc_type: object, exc: object, tb: object) -> None:
del exc_type, exc, tb
def read(self) -> bytes:
return b'{"ok": true}'
def fake_urlopen(request: Any, timeout: int) -> StubResponse:
assert timeout == 90
assert request.full_url == "http://localhost:11434/v1/chat/completions"
return StubResponse()
monkeypatch.setattr(urllib.request, "urlopen", fake_urlopen)
assert ollama_demo._post_chat_completion("http://localhost:11434/v1", {"x": 1}) == {"ok": True}
def test_post_chat_completion_errors(monkeypatch: pytest.MonkeyPatch) -> None:
def fake_urlopen_error(request: Any, timeout: int) -> Any:
del request, timeout
raise urllib.error.URLError("boom")
monkeypatch.setattr(urllib.request, "urlopen", fake_urlopen_error)
with pytest.raises(RuntimeError, match="failed to call Ollama"):
ollama_demo._post_chat_completion("http://localhost:11434/v1", {"x": 1})
class StubResponse:
def __enter__(self) -> StubResponse:
return self
def __exit__(self, exc_type: object, exc: object, tb: object) -> None:
del exc_type, exc, tb
def read(self) -> bytes:
return b'["bad"]'
def fake_urlopen_non_object(request: Any, timeout: int) -> StubResponse:
del request, timeout
return StubResponse()
monkeypatch.setattr(urllib.request, "urlopen", fake_urlopen_non_object)
with pytest.raises(TypeError, match="unexpected Ollama response shape"):
ollama_demo._post_chat_completion("http://localhost:11434/v1", {"x": 1})
@pytest.mark.parametrize(
("raw", "expected"),
[(None, {}), ({}, {}), ("", {}), ('{"a":1}', {"a": 1})],
)
def test_parse_tool_arguments(raw: Any, expected: dict[str, Any]) -> None:
assert ollama_demo._parse_tool_arguments(raw) == expected
def test_parse_tool_arguments_invalid() -> None:
with pytest.raises(TypeError, match="decode to an object"):
ollama_demo._parse_tool_arguments("[]")
with pytest.raises(TypeError, match="dictionary or JSON object string"):
ollama_demo._parse_tool_arguments(3)
@pytest.mark.parametrize(
("response", "msg"),
[
({}, "did not contain completion choices"),
({"choices": [1]}, "unexpected completion choice format"),
({"choices": [{"message": "bad"}]}, "did not contain a message"),
],
)
def test_extract_message_validation(response: dict[str, Any], msg: str) -> None:
with pytest.raises(RuntimeError, match=msg):
ollama_demo._extract_message(response)
def test_build_parser_defaults() -> None:
parser = ollama_demo._build_parser()
args = parser.parse_args([])
assert args.model == ollama_demo.DEFAULT_OLLAMA_MODEL
assert args.base_url == ollama_demo.DEFAULT_OLLAMA_BASE_URL
assert args.verbose is False
def test_build_parser_verbose_flag() -> None:
parser = ollama_demo._build_parser()
args = parser.parse_args(["-v"])
assert args.verbose is True
def test_main_uses_parser_and_prints_logs(
monkeypatch: pytest.MonkeyPatch,
capsys: pytest.CaptureFixture[str],
) -> None:
class StubParser:
def parse_args(self) -> argparse.Namespace:
return argparse.Namespace(base_url="http://x", model="m", verbose=False)
monkeypatch.setattr(ollama_demo, "_build_parser", lambda: StubParser())
monkeypatch.setattr(
ollama_demo,
"run_ollama_tool_demo",
lambda base_url, model, strict=True, verbose=False, log=None: {
"exec_result": {
"exit_code": 0,
"stdout": "true\n",
"execution_mode": "host_compat",
},
"fallback_used": False,
},
)
ollama_demo.main()
output = capsys.readouterr().out
assert "[summary] exit_code=0 fallback_used=False execution_mode=host_compat" in output
assert "[summary] stdout=" not in output
def test_main_verbose_prints_stdout(
monkeypatch: pytest.MonkeyPatch,
capsys: pytest.CaptureFixture[str],
) -> None:
class StubParser:
def parse_args(self) -> argparse.Namespace:
return argparse.Namespace(base_url="http://x", model="m", verbose=True)
monkeypatch.setattr(ollama_demo, "_build_parser", lambda: StubParser())
monkeypatch.setattr(
ollama_demo,
"run_ollama_tool_demo",
lambda base_url, model, strict=True, verbose=False, log=None: {
"exec_result": {
"exit_code": 0,
"stdout": "true\n",
"execution_mode": "host_compat",
},
"fallback_used": False,
},
)
ollama_demo.main()
output = capsys.readouterr().out
assert "[summary] stdout=true" in output
def test_main_logs_error_and_exits_nonzero(
monkeypatch: pytest.MonkeyPatch,
capsys: pytest.CaptureFixture[str],
) -> None:
class StubParser:
def parse_args(self) -> argparse.Namespace:
return argparse.Namespace(base_url="http://x", model="m", verbose=False)
monkeypatch.setattr(ollama_demo, "_build_parser", lambda: StubParser())
def fake_run(
base_url: str,
model: str,
strict: bool = True,
verbose: bool = False,
log: Any = None,
) -> dict[str, Any]:
del base_url, model, strict, verbose, log
raise RuntimeError("demo did not execute a successful vm_exec")
monkeypatch.setattr(ollama_demo, "run_ollama_tool_demo", fake_run)
with pytest.raises(SystemExit, match="1"):
ollama_demo.main()
output = capsys.readouterr().out
assert "[error] demo did not execute a successful vm_exec" in output