pyro-mcp/tests/test_vm_environments.py
Thales Maciel 6e16e74fd5 Harden default environment pull behavior
Fix the default one-shot install path so empty bundled profile directories no longer win over OCI-backed environment pulls or leave broken cached symlinks behind.

Treat cached installs as valid only when the manifest and boot artifacts are all present, repair invalid installs on the next pull, and add human-mode phase markers for env pull and run without changing JSON output.

Align the Python lifecycle example and public docs with the current exec_vm/vm_exec auto-clean semantics, and validate the slice with focused pytest coverage, make check, make dist-check, and a real default-path pull/inspect/run smoke.
2026-03-11 19:27:09 -03:00

503 lines
19 KiB
Python

from __future__ import annotations
import hashlib
import io
import json
import tarfile
import urllib.error
import urllib.request
from email.message import Message
from pathlib import Path
import pytest
from pyro_mcp.runtime import RuntimePaths, resolve_runtime_paths
from pyro_mcp.vm_environments import (
EnvironmentStore,
VmEnvironment,
get_environment,
list_environments,
)
class FakeResponse:
def __init__(self, payload: bytes, *, headers: dict[str, str] | None = None) -> None:
self._buffer = io.BytesIO(payload)
self.headers = headers or {}
def read(self, size: int = -1) -> bytes:
return self._buffer.read(size)
def __enter__(self) -> FakeResponse:
return self
def __exit__(self, exc_type: object, exc: object, tb: object) -> None:
del exc_type, exc, tb
def _fake_runtime_paths(tmp_path: Path) -> RuntimePaths:
bundle_parent = tmp_path / "runtime"
bundle_root = bundle_parent / "linux-x86_64"
manifest_path = bundle_root / "manifest.json"
firecracker_bin = bundle_root / "bin" / "firecracker"
jailer_bin = bundle_root / "bin" / "jailer"
guest_agent_path = bundle_root / "guest" / "pyro_guest_agent.py"
artifacts_dir = bundle_root / "profiles"
notice_path = bundle_parent / "NOTICE"
artifacts_dir.mkdir(parents=True, exist_ok=True)
firecracker_bin.parent.mkdir(parents=True, exist_ok=True)
jailer_bin.parent.mkdir(parents=True, exist_ok=True)
guest_agent_path.parent.mkdir(parents=True, exist_ok=True)
manifest_path.write_text('{"platform": "linux-x86_64"}\n', encoding="utf-8")
firecracker_bin.write_text("firecracker\n", encoding="utf-8")
jailer_bin.write_text("jailer\n", encoding="utf-8")
guest_agent_path.write_text("print('guest')\n", encoding="utf-8")
notice_path.write_text("notice\n", encoding="utf-8")
return RuntimePaths(
bundle_root=bundle_root,
manifest_path=manifest_path,
firecracker_bin=firecracker_bin,
jailer_bin=jailer_bin,
guest_agent_path=guest_agent_path,
artifacts_dir=artifacts_dir,
notice_path=notice_path,
manifest={"platform": "linux-x86_64"},
)
def _write_local_profile(
runtime_paths: RuntimePaths,
profile_name: str,
*,
kernel: str = "kernel\n",
rootfs: str = "rootfs\n",
) -> None:
profile_dir = runtime_paths.artifacts_dir / profile_name
profile_dir.mkdir(parents=True, exist_ok=True)
(profile_dir / "vmlinux").write_text(kernel, encoding="utf-8")
(profile_dir / "rootfs.ext4").write_text(rootfs, encoding="utf-8")
def _sha256_digest(payload: bytes) -> str:
return f"sha256:{hashlib.sha256(payload).hexdigest()}"
def _layer_archive(filename: str, content: bytes) -> bytes:
archive_buffer = io.BytesIO()
with tarfile.open(fileobj=archive_buffer, mode="w:gz") as archive:
info = tarfile.TarInfo(name=filename)
info.size = len(content)
archive.addfile(info, io.BytesIO(content))
return archive_buffer.getvalue()
def _authorization_header(request: object) -> str | None:
if isinstance(request, urllib.request.Request):
for key, value in request.header_items():
if key.lower() == "authorization":
return value
return None
def _http_headers(headers: dict[str, str]) -> Message:
message = Message()
for key, value in headers.items():
message[key] = value
return message
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:
runtime_paths = _fake_runtime_paths(tmp_path)
_write_local_profile(runtime_paths, "debian-git")
store = EnvironmentStore(runtime_paths=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:
runtime_paths = _fake_runtime_paths(tmp_path)
_write_local_profile(runtime_paths, "debian-git")
store = EnvironmentStore(runtime_paths=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 = _fake_runtime_paths(tmp_path)
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")
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=runtime_paths,
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_skips_empty_local_source_dir_and_uses_archive(
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
runtime_paths = _fake_runtime_paths(tmp_path)
source_environment = get_environment("debian:12-base", runtime_paths=runtime_paths)
(runtime_paths.artifacts_dir / source_environment.source_profile).mkdir(
parents=True,
exist_ok=True,
)
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")
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=runtime_paths, 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_marks_broken_symlink_install_uninstalled_and_repairs_it(
tmp_path: Path,
) -> None:
runtime_paths = _fake_runtime_paths(tmp_path)
_write_local_profile(
runtime_paths,
"debian-git",
kernel="kernel-fixed\n",
rootfs="rootfs-fixed\n",
)
store = EnvironmentStore(runtime_paths=runtime_paths, cache_dir=tmp_path / "cache")
spec = get_environment("debian:12", runtime_paths=runtime_paths)
install_dir = store.cache_dir / "linux-x86_64" / "debian_12-1.0.0"
install_dir.mkdir(parents=True, exist_ok=True)
(install_dir / "environment.json").write_text(
json.dumps(
{
"catalog_version": "2.0.0",
"name": spec.name,
"version": spec.version,
"source": "bundled-runtime-source",
"source_digest": spec.source_digest,
"installed_at": 0,
}
),
encoding="utf-8",
)
(install_dir / "vmlinux").symlink_to("missing-vmlinux")
(install_dir / "rootfs.ext4").symlink_to("missing-rootfs.ext4")
inspected_before = store.inspect_environment("debian:12")
assert inspected_before["installed"] is False
pulled = store.pull_environment("debian:12")
assert pulled["installed"] is True
assert Path(str(pulled["kernel_image"])).read_text(encoding="utf-8") == "kernel-fixed\n"
assert Path(str(pulled["rootfs_image"])).read_text(encoding="utf-8") == "rootfs-fixed\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
def test_fetch_oci_manifest_resolves_linux_amd64_index_with_bearer_auth(
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
runtime_paths = _fake_runtime_paths(tmp_path)
store = EnvironmentStore(runtime_paths=runtime_paths, cache_dir=tmp_path / "cache")
spec = VmEnvironment(
name="debian:12-registry",
version="1.0.0",
description="OCI-backed environment",
default_packages=("bash", "git"),
distribution="debian",
distribution_version="12",
source_profile="missing-profile",
oci_registry="registry-1.docker.io",
oci_repository="thalesmaciel/pyro-environment-debian-12",
oci_reference="1.0.0",
)
child_manifest = {
"schemaVersion": 2,
"mediaType": "application/vnd.oci.image.manifest.v1+json",
"layers": [],
}
child_payload = json.dumps(child_manifest).encode("utf-8")
child_digest = _sha256_digest(child_payload)
index_manifest = {
"schemaVersion": 2,
"mediaType": "application/vnd.oci.image.index.v1+json",
"manifests": [
{
"digest": "sha256:arm64digest",
"mediaType": "application/vnd.oci.image.manifest.v1+json",
"platform": {"os": "linux", "architecture": "arm64"},
},
{
"digest": child_digest,
"mediaType": "application/vnd.oci.image.manifest.v1+json",
"platform": {"os": "linux", "architecture": "amd64"},
},
],
}
index_payload = json.dumps(index_manifest).encode("utf-8")
index_digest = _sha256_digest(index_payload)
authorized_urls: list[str] = []
def fake_urlopen(request: object, timeout: int = 90) -> FakeResponse:
del timeout
url = request.full_url if isinstance(request, urllib.request.Request) else str(request)
if url.startswith("https://auth.docker.io/token?"):
return FakeResponse(b'{"token":"secret-token"}')
authorization = _authorization_header(request)
if url.endswith("/manifests/1.0.0"):
if authorization is None:
raise urllib.error.HTTPError(
url,
401,
"Unauthorized",
_http_headers(
{
"WWW-Authenticate": (
'Bearer realm="https://auth.docker.io/token",'
'service="registry.docker.io",'
'scope="repository:thalesmaciel/pyro-environment-debian-12:pull"'
)
}
),
io.BytesIO(b""),
)
authorized_urls.append(url)
return FakeResponse(
index_payload,
headers={"Docker-Content-Digest": index_digest},
)
if url.endswith(f"/manifests/{child_digest}"):
if authorization is None:
raise urllib.error.HTTPError(
url,
401,
"Unauthorized",
_http_headers(
{
"WWW-Authenticate": (
'Bearer realm="https://auth.docker.io/token",'
'service="registry.docker.io",'
'scope="repository:thalesmaciel/pyro-environment-debian-12:pull"'
)
}
),
io.BytesIO(b""),
)
authorized_urls.append(url)
assert authorization == "Bearer secret-token"
return FakeResponse(
child_payload,
headers={"Docker-Content-Digest": child_digest},
)
raise AssertionError(f"unexpected OCI request: {url}")
monkeypatch.setattr(urllib.request, "urlopen", fake_urlopen)
manifest, resolved_digest = store._fetch_oci_manifest(spec) # noqa: SLF001
assert manifest == child_manifest
assert resolved_digest == child_digest
assert authorized_urls == [
"https://registry-1.docker.io/v2/thalesmaciel/pyro-environment-debian-12/manifests/1.0.0",
(
"https://registry-1.docker.io/v2/thalesmaciel/pyro-environment-debian-12/"
f"manifests/{child_digest}"
),
]
def test_environment_store_installs_from_oci_when_runtime_source_missing(
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
runtime_paths = _fake_runtime_paths(tmp_path)
kernel_layer = _layer_archive("vmlinux", b"kernel\n")
rootfs_layer = _layer_archive("rootfs.ext4", b"rootfs\n")
kernel_digest = _sha256_digest(kernel_layer)
rootfs_digest = _sha256_digest(rootfs_layer)
manifest = {
"schemaVersion": 2,
"mediaType": "application/vnd.oci.image.manifest.v1+json",
"layers": [
{
"mediaType": "application/vnd.oci.image.layer.v1.tar+gzip",
"digest": kernel_digest,
"size": len(kernel_layer),
},
{
"mediaType": "application/vnd.oci.image.layer.v1.tar+gzip",
"digest": rootfs_digest,
"size": len(rootfs_layer),
},
],
}
manifest_payload = json.dumps(manifest).encode("utf-8")
manifest_digest = _sha256_digest(manifest_payload)
environment = VmEnvironment(
name="debian:12-registry",
version="1.0.0",
description="OCI-backed environment",
default_packages=("bash", "git"),
distribution="debian",
distribution_version="12",
source_profile="missing-profile",
oci_registry="registry-1.docker.io",
oci_repository="thalesmaciel/pyro-environment-debian-12",
oci_reference="1.0.0",
)
def fake_urlopen(request: object, timeout: int = 90) -> FakeResponse:
del timeout
url = request.full_url if isinstance(request, urllib.request.Request) else str(request)
if url.endswith("/manifests/1.0.0"):
return FakeResponse(
manifest_payload,
headers={"Docker-Content-Digest": manifest_digest},
)
if url.endswith(f"/blobs/{kernel_digest}"):
return FakeResponse(kernel_layer)
if url.endswith(f"/blobs/{rootfs_digest}"):
return FakeResponse(rootfs_layer)
raise AssertionError(f"unexpected OCI request: {url}")
monkeypatch.setattr(urllib.request, "urlopen", fake_urlopen)
monkeypatch.setattr(
"pyro_mcp.vm_environments.CATALOG",
{environment.name: environment},
)
store = EnvironmentStore(runtime_paths=runtime_paths, cache_dir=tmp_path / "cache")
installed = store.ensure_installed(environment.name)
assert installed.kernel_image.read_text(encoding="utf-8") == "kernel\n"
assert installed.rootfs_image.read_text(encoding="utf-8") == "rootfs\n"
assert installed.source == (
"oci://registry-1.docker.io/thalesmaciel/pyro-environment-debian-12"
f"@{manifest_digest}"
)
metadata = json.loads((installed.install_dir / "environment.json").read_text(encoding="utf-8"))
assert metadata["source_digest"] == manifest_digest