sand-boxer/tests/test_snapshots.py

"""Snapshot store and manager checkpoint tests."""

from __future__ import annotations

from datetime import UTC, datetime
from pathlib import Path
from unittest.mock import patch

import pytest

from sandboxer.core.manager import SandboxManager
from sandboxer.lifecycle.store import SandboxStore
from sandboxer.models import (
    ActorType,
    Consumer,
    Reachability,
    SandboxCreateRequest,
    SandboxState,
    SandboxStatus,
    SnapshotRecord,
)
from sandboxer.snapshots.store import SnapshotStore


class SnapshotBackend:
    def supports_snapshots(self) -> bool:
        return True

    def provision(self, profile, inputs, host):
        return {
            "sandbox_id": "test1234",
            "host": host,
            "remote_dir": "/tmp/sandboxer/test1234",
            "compose_project": "sbx-e2e-test1234",
            "compose_file": "docker-compose.yml",
            "ssh_user": "root",
        }

    def wait_ready(self, handle):
        return {
            "ssh": f"root@{handle['host']}",
            "remote_dir": handle["remote_dir"],
            "compose_project": handle["compose_project"],
            "host": handle["host"],
        }

    def teardown(self, handle):
        return {"compose_removed": "True", "remote_dir_removed": "True"}

    def snapshot(self, handle):
        return {
            "snapshot_id": "snap12345678",
            "artifact_path": "/tmp/sandboxer/snapshots/snap12345678.tar.gz",
            "host": handle["host"],
            "size_bytes": "4096",
        }

    def restore_from_snapshot(self, profile, snapshot_meta, inputs, host):
        return {
            "sandbox_id": "restored1",
            "host": host,
            "remote_dir": "/tmp/sandboxer/restored1",
            "compose_project": "sbx-e2e-restored1",
            "compose_file": "docker-compose.yml",
            "ssh_user": "root",
        }


@pytest.fixture
def store(tmp_path: Path) -> SandboxStore:
    return SandboxStore(path=tmp_path / "sandboxes.json")


@pytest.fixture
def snapshots(tmp_path: Path) -> SnapshotStore:
    return SnapshotStore(path=tmp_path / "snapshots.json")


def _ready_status(sandbox_id: str = "test1234") -> SandboxStatus:
    now = datetime.now(UTC)
    return SandboxStatus(
        sandbox_id=sandbox_id,
        profile_id="profile.compose-checkpoint",
        extension_id="ext.compose-ssh",
        state=SandboxState.READY,
        consumer=Consumer(actor=ActorType.ADM, project="sand-boxer"),
        host="coulombcore",
        reachability=Reachability(
            ssh="root@coulombcore",
            remote_dir="/tmp/sandboxer/test1234",
            compose_project="sbx-e2e-test1234",
            host="coulombcore",
        ),
        inputs={
            "repo": "/tmp/repo",
            "compose_file": "docker-compose.yml",
            "ssh_user": "root",
        },
        created_at=now,
        updated_at=now,
        ready_at=now,
    )


def test_snapshot_store_roundtrip(snapshots: SnapshotStore) -> None:
    now = datetime.now(UTC)
    record = SnapshotRecord(
        snapshot_id="snap12345678",
        sandbox_id="test1234",
        profile_id="profile.compose-checkpoint",
        extension_id="ext.compose-ssh",
        host="coulombcore",
        artifact_path="/tmp/snap.tar.gz",
        created_at=now,
    )
    snapshots.save(record)
    loaded = snapshots.get("snap12345678")
    assert loaded is not None
    assert loaded.sandbox_id == "test1234"


def test_manager_snapshot_and_restore(store: SandboxStore, snapshots: SnapshotStore) -> None:
    store.save(_ready_status())
    manager = SandboxManager(store=store, snapshots=snapshots)
    backend = SnapshotBackend()

    with (
        patch("sandboxer.core.manager.resolve_backend", return_value=backend),
        patch("sandboxer.core.manager.load_extension"),
        patch("sandboxer.core.manager.emit_lifecycle_event", return_value=None),
        patch("sandboxer.core.manager.load_profile"),
        patch("sandboxer.core.manager.resolve_host", return_value="coulombcore"),
    ):
        record = manager.snapshot("test1234", name="pre-test")
        assert record.snapshot_id == "snap12345678"
        assert record.name == "pre-test"
        assert record.size_bytes == 4096

        status = manager.restore("snap12345678")
        assert status.state == SandboxState.READY
        assert status.sandbox_id == "restored1"
        assert status.inputs.get("restored_from") == "snap12345678"


def test_snapshot_requires_ready(store: SandboxStore, snapshots: SnapshotStore) -> None:
    status = _ready_status()
    status.state = SandboxState.PROVISIONING
    store.save(status)
    manager = SandboxManager(store=store, snapshots=snapshots)

    with pytest.raises(RuntimeError, match="ready"):
        manager.snapshot("test1234")


def test_create_snapshot_restore_flow(store: SandboxStore, snapshots: SnapshotStore) -> None:
    manager = SandboxManager(store=store, snapshots=snapshots)
    backend = SnapshotBackend()
    request = SandboxCreateRequest(
        profile="profile.compose-checkpoint",
        inputs={"repo": "/tmp/repo"},
        consumer=Consumer(actor=ActorType.ADM, project="sand-boxer"),
    )

    with (
        patch("sandboxer.core.manager.resolve_backend", return_value=backend),
        patch("sandboxer.core.manager.emit_lifecycle_event", return_value=None),
        patch("sandboxer.core.manager.resolve_host", return_value="coulombcore"),
    ):
        created = manager.create(request)
        record = manager.snapshot(created.sandbox_id)
        restored = manager.restore(record.snapshot_id)
        assert restored.sandbox_id == "restored1"