pyro-mcp/tests/test_vm_environments.py

from __future__ import annotations

import tarfile
from pathlib import Path

import pytest

from pyro_mcp.runtime import resolve_runtime_paths
from pyro_mcp.vm_environments import EnvironmentStore, get_environment, list_environments


def test_list_environments_includes_expected_entries() -> None:
    environments = list_environments(runtime_paths=resolve_runtime_paths())
    names = {str(entry["name"]) for entry in environments}
    assert {"debian:12", "debian:12-base", "debian:12-build"} <= names


def test_get_environment_rejects_unknown() -> None:
    with pytest.raises(ValueError, match="unknown environment"):
        get_environment("does-not-exist")


def test_environment_store_installs_from_local_runtime_source(tmp_path: Path) -> None:
    store = EnvironmentStore(runtime_paths=resolve_runtime_paths(), cache_dir=tmp_path / "cache")
    installed = store.ensure_installed("debian:12")

    assert installed.kernel_image.exists()
    assert installed.rootfs_image.exists()
    assert (installed.install_dir / "environment.json").exists()


def test_environment_store_pull_and_cached_inspect(tmp_path: Path) -> None:
    store = EnvironmentStore(runtime_paths=resolve_runtime_paths(), cache_dir=tmp_path / "cache")

    before = store.inspect_environment("debian:12")
    assert before["installed"] is False

    pulled = store.pull_environment("debian:12")
    assert pulled["installed"] is True
    assert "install_manifest" in pulled

    cached = store.ensure_installed("debian:12")
    assert cached.installed is True

    after = store.inspect_environment("debian:12")
    assert after["installed"] is True
    assert "install_manifest" in after


def test_environment_store_uses_env_override_for_default_cache_dir(
    monkeypatch: pytest.MonkeyPatch, tmp_path: Path
) -> None:
    monkeypatch.setenv("PYRO_ENVIRONMENT_CACHE_DIR", str(tmp_path / "override-cache"))
    store = EnvironmentStore(runtime_paths=resolve_runtime_paths())
    assert store.cache_dir == tmp_path / "override-cache"


def test_environment_store_installs_from_archive_when_runtime_source_missing(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    runtime_paths = resolve_runtime_paths()
    source_environment = get_environment("debian:12-base", runtime_paths=runtime_paths)

    archive_dir = tmp_path / "archive"
    archive_dir.mkdir(parents=True, exist_ok=True)
    (archive_dir / "vmlinux").write_text("kernel\n", encoding="utf-8")
    (archive_dir / "rootfs.ext4").write_text("rootfs\n", encoding="utf-8")
    archive_path = tmp_path / "environment.tgz"
    with tarfile.open(archive_path, "w:gz") as archive:
        archive.add(archive_dir / "vmlinux", arcname="vmlinux")
        archive.add(archive_dir / "rootfs.ext4", arcname="rootfs.ext4")

    missing_bundle = tmp_path / "bundle"
    platform_root = missing_bundle / "linux-x86_64"
    platform_root.mkdir(parents=True, exist_ok=True)
    (missing_bundle / "NOTICE").write_text(
        runtime_paths.notice_path.read_text(encoding="utf-8"),
        encoding="utf-8",
    )
    (platform_root / "manifest.json").write_text(
        runtime_paths.manifest_path.read_text(encoding="utf-8"),
        encoding="utf-8",
    )
    (platform_root / "bin").mkdir(parents=True, exist_ok=True)
    (platform_root / "bin" / "firecracker").write_bytes(runtime_paths.firecracker_bin.read_bytes())
    (platform_root / "bin" / "jailer").write_bytes(runtime_paths.jailer_bin.read_bytes())
    guest_agent_path = runtime_paths.guest_agent_path
    if guest_agent_path is None:
        raise AssertionError("expected guest agent path")
    (platform_root / "guest").mkdir(parents=True, exist_ok=True)
    (platform_root / "guest" / "pyro_guest_agent.py").write_text(
        guest_agent_path.read_text(encoding="utf-8"),
        encoding="utf-8",
    )

    monkeypatch.setenv("PYRO_RUNTIME_BUNDLE_DIR", str(missing_bundle))
    monkeypatch.setattr(
        "pyro_mcp.vm_environments.CATALOG",
        {
            "debian:12-base": source_environment.__class__(
                name=source_environment.name,
                version=source_environment.version,
                description=source_environment.description,
                default_packages=source_environment.default_packages,
                distribution=source_environment.distribution,
                distribution_version=source_environment.distribution_version,
                source_profile=source_environment.source_profile,
                platform=source_environment.platform,
                source_url=archive_path.resolve().as_uri(),
                source_digest=source_environment.source_digest,
                compatibility=source_environment.compatibility,
            )
        },
    )
    store = EnvironmentStore(
        runtime_paths=resolve_runtime_paths(verify_checksums=False),
        cache_dir=tmp_path / "cache",
    )
    installed = store.ensure_installed("debian:12-base")

    assert installed.kernel_image.read_text(encoding="utf-8") == "kernel\n"
    assert installed.rootfs_image.read_text(encoding="utf-8") == "rootfs\n"


def test_environment_store_prunes_stale_entries(tmp_path: Path) -> None:
    store = EnvironmentStore(runtime_paths=resolve_runtime_paths(), cache_dir=tmp_path / "cache")
    platform_dir = store.cache_dir / "linux-x86_64"
    platform_dir.mkdir(parents=True, exist_ok=True)

    (platform_dir / ".partial-download").mkdir()
    (platform_dir / "missing-marker").mkdir()

    invalid = platform_dir / "invalid"
    invalid.mkdir()
    (invalid / "environment.json").write_text('{"name": 1, "version": 2}', encoding="utf-8")

    unknown = platform_dir / "unknown"
    unknown.mkdir()
    (unknown / "environment.json").write_text(
        '{"name": "unknown:1", "version": "1.0.0"}',
        encoding="utf-8",
    )

    stale = platform_dir / "stale"
    stale.mkdir()
    (stale / "environment.json").write_text(
        '{"name": "debian:12", "version": "0.9.0"}',
        encoding="utf-8",
    )

    result = store.prune_environments()

    assert result["count"] == 5