164 lines
6.2 KiB
Python
164 lines
6.2 KiB
Python
from __future__ import annotations
|
|
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import pytest
|
|
|
|
import pyro_mcp.vm_manager as vm_manager_module
|
|
from pyro_mcp.runtime import resolve_runtime_paths
|
|
from pyro_mcp.vm_manager import VmManager
|
|
|
|
|
|
def test_vm_manager_lifecycle_and_auto_cleanup(tmp_path: Path) -> None:
|
|
manager = VmManager(backend_name="mock", base_dir=tmp_path / "vms")
|
|
created = manager.create_vm(profile="debian-git", vcpu_count=1, mem_mib=512, ttl_seconds=600)
|
|
vm_id = str(created["vm_id"])
|
|
started = manager.start_vm(vm_id)
|
|
assert started["state"] == "started"
|
|
|
|
executed = manager.exec_vm(vm_id, command="printf 'git version 2.43.0\\n'", timeout_seconds=30)
|
|
assert executed["exit_code"] == 0
|
|
assert "git version" in str(executed["stdout"])
|
|
with pytest.raises(ValueError, match="does not exist"):
|
|
manager.status_vm(vm_id)
|
|
|
|
|
|
def test_vm_manager_exec_timeout(tmp_path: Path) -> None:
|
|
manager = VmManager(backend_name="mock", base_dir=tmp_path / "vms")
|
|
vm_id = str(
|
|
manager.create_vm(profile="debian-base", vcpu_count=1, mem_mib=512, ttl_seconds=600)[
|
|
"vm_id"
|
|
]
|
|
)
|
|
manager.start_vm(vm_id)
|
|
result = manager.exec_vm(vm_id, command="sleep 2", timeout_seconds=1)
|
|
assert result["exit_code"] == 124
|
|
assert "timed out" in str(result["stderr"])
|
|
|
|
|
|
def test_vm_manager_stop_and_delete(tmp_path: Path) -> None:
|
|
manager = VmManager(backend_name="mock", base_dir=tmp_path / "vms")
|
|
vm_id = str(
|
|
manager.create_vm(profile="debian-base", vcpu_count=1, mem_mib=512, ttl_seconds=600)[
|
|
"vm_id"
|
|
]
|
|
)
|
|
manager.start_vm(vm_id)
|
|
stopped = manager.stop_vm(vm_id)
|
|
assert stopped["state"] == "stopped"
|
|
deleted = manager.delete_vm(vm_id)
|
|
assert deleted["deleted"] is True
|
|
|
|
|
|
def test_vm_manager_reaps_expired(tmp_path: Path) -> None:
|
|
manager = VmManager(backend_name="mock", base_dir=tmp_path / "vms")
|
|
manager.MIN_TTL_SECONDS = 1
|
|
vm_id = str(
|
|
manager.create_vm(profile="debian-base", vcpu_count=1, mem_mib=512, ttl_seconds=1)["vm_id"]
|
|
)
|
|
instance = manager._instances[vm_id] # noqa: SLF001
|
|
instance.expires_at = 0.0
|
|
result = manager.reap_expired()
|
|
assert result["count"] == 1
|
|
with pytest.raises(ValueError):
|
|
manager.status_vm(vm_id)
|
|
|
|
|
|
def test_vm_manager_reaps_started_vm(tmp_path: Path) -> None:
|
|
manager = VmManager(backend_name="mock", base_dir=tmp_path / "vms")
|
|
manager.MIN_TTL_SECONDS = 1
|
|
vm_id = str(
|
|
manager.create_vm(profile="debian-base", vcpu_count=1, mem_mib=512, ttl_seconds=1)["vm_id"]
|
|
)
|
|
manager.start_vm(vm_id)
|
|
manager._instances[vm_id].expires_at = 0.0 # noqa: SLF001
|
|
result = manager.reap_expired()
|
|
assert result["count"] == 1
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("kwargs", "msg"),
|
|
[
|
|
({"vcpu_count": 0, "mem_mib": 512, "ttl_seconds": 600}, "vcpu_count must be between"),
|
|
({"vcpu_count": 1, "mem_mib": 64, "ttl_seconds": 600}, "mem_mib must be between"),
|
|
({"vcpu_count": 1, "mem_mib": 512, "ttl_seconds": 30}, "ttl_seconds must be between"),
|
|
],
|
|
)
|
|
def test_vm_manager_validates_limits(tmp_path: Path, kwargs: dict[str, int], msg: str) -> None:
|
|
manager = VmManager(backend_name="mock", base_dir=tmp_path / "vms")
|
|
with pytest.raises(ValueError, match=msg):
|
|
manager.create_vm(profile="debian-base", **kwargs)
|
|
|
|
|
|
def test_vm_manager_max_active_limit(tmp_path: Path) -> None:
|
|
manager = VmManager(backend_name="mock", base_dir=tmp_path / "vms", max_active_vms=1)
|
|
manager.create_vm(profile="debian-base", vcpu_count=1, mem_mib=512, ttl_seconds=600)
|
|
with pytest.raises(RuntimeError, match="max active VMs reached"):
|
|
manager.create_vm(profile="debian-base", vcpu_count=1, mem_mib=512, ttl_seconds=600)
|
|
|
|
|
|
def test_vm_manager_state_validation(tmp_path: Path) -> None:
|
|
manager = VmManager(backend_name="mock", base_dir=tmp_path / "vms")
|
|
vm_id = str(
|
|
manager.create_vm(profile="debian-base", vcpu_count=1, mem_mib=512, ttl_seconds=600)[
|
|
"vm_id"
|
|
]
|
|
)
|
|
with pytest.raises(RuntimeError, match="must be in 'started' state"):
|
|
manager.exec_vm(vm_id, command="echo hi", timeout_seconds=30)
|
|
with pytest.raises(ValueError, match="must be positive"):
|
|
manager.exec_vm(vm_id, command="echo hi", timeout_seconds=0)
|
|
manager.start_vm(vm_id)
|
|
with pytest.raises(RuntimeError, match="cannot be started from state"):
|
|
manager.start_vm(vm_id)
|
|
|
|
|
|
def test_vm_manager_status_expired_raises(tmp_path: Path) -> None:
|
|
manager = VmManager(backend_name="mock", base_dir=tmp_path / "vms")
|
|
manager.MIN_TTL_SECONDS = 1
|
|
vm_id = str(
|
|
manager.create_vm(profile="debian-base", vcpu_count=1, mem_mib=512, ttl_seconds=1)["vm_id"]
|
|
)
|
|
manager._instances[vm_id].expires_at = 0.0 # noqa: SLF001
|
|
with pytest.raises(RuntimeError, match="expired and was automatically deleted"):
|
|
manager.status_vm(vm_id)
|
|
|
|
|
|
def test_vm_manager_invalid_backend(tmp_path: Path) -> None:
|
|
with pytest.raises(ValueError, match="invalid backend"):
|
|
VmManager(backend_name="nope", base_dir=tmp_path / "vms")
|
|
|
|
|
|
def test_vm_manager_firecracker_backend_path(
|
|
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
|
) -> None:
|
|
class StubFirecrackerBackend:
|
|
def __init__(self, artifacts_dir: Path, firecracker_bin: Path, jailer_bin: Path) -> None:
|
|
self.artifacts_dir = artifacts_dir
|
|
self.firecracker_bin = firecracker_bin
|
|
self.jailer_bin = jailer_bin
|
|
|
|
def create(self, instance: Any) -> None:
|
|
del instance
|
|
|
|
def start(self, instance: Any) -> None:
|
|
del instance
|
|
|
|
def exec(self, instance: Any, command: str, timeout_seconds: int) -> Any:
|
|
del instance, command, timeout_seconds
|
|
return None
|
|
|
|
def stop(self, instance: Any) -> None:
|
|
del instance
|
|
|
|
def delete(self, instance: Any) -> None:
|
|
del instance
|
|
|
|
monkeypatch.setattr(vm_manager_module, "FirecrackerBackend", StubFirecrackerBackend)
|
|
manager = VmManager(
|
|
backend_name="firecracker",
|
|
base_dir=tmp_path / "vms",
|
|
runtime_paths=resolve_runtime_paths(),
|
|
)
|
|
assert manager._backend_name == "firecracker" # noqa: SLF001
|