"""Unit tests for the restore-from-snapshot surface. Coverage: - ``_rehydrate_config_from_snapshot``: annotations supply defaults; overrides win; missing image without override is a loud error. - ``_collect_source_annotations``: reads the source Deployment + PVC and returns the ``k7.io/source-*`` annotation dict that gets stamped onto pause/named snapshots so restore can rehydrate later. - API restore endpoint: maps ``new_sandbox_name`` / overrides into ``K7Core.restore_sandbox`` and the response status codes (``201`` / ``400`` / ``404`` / ``409``). - SDK ``Client.restore``: builds the right body and returns a ``SandboxProxy`` for the new sandbox. """ import hashlib import json import time from pathlib import Path from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest from k7.api.main import app from k7.core.core import K7Core from k7.core.models import ( SNAPSHOT_KIND_PAUSE, OperationResult, SandboxConfig, SandboxConfigOverrides, ) from k7_sdk.client import Client, SandboxProxy TEST_KEY = "k7-test-secret-key-restore" TEST_KEY_HASH = hashlib.sha256(TEST_KEY.encode()).hexdigest() # --- _rehydrate_config_from_snapshot ---------------------------------------- class TestRehydrateConfigFromSnapshot: def test_uses_annotations_when_no_override(self, core: K7Core): ann = { "k7.io/source-image": "python:3.12-slim", "k7.io/source-backend": "kata-qemu-longhorn", "k7.io/source-root-disk-size": "20Gi", "k7.io/source-sidecar": "docker", "k7.io/source-limits": json.dumps({"cpu": "2", "memory": "4Gi"}), } result = core._rehydrate_config_from_snapshot( annotations=ann, new_name="restored", namespace="default", overrides=None ) assert result.success config: SandboxConfig = result.data assert config.image == "python:3.12-slim" assert config.backend == "kata-qemu-longhorn" assert config.root_disk_size == "20Gi" assert config.sidecar == "docker" assert config.limits == {"cpu": "2", "memory": "4Gi"} assert config.name == "restored" assert config.namespace == "default" def test_overrides_win_over_annotations(self, core: K7Core): ann = { "k7.io/source-image": "python:3.12-slim", "k7.io/source-backend": "kata-qemu-longhorn", "k7.io/source-root-disk-size": "10Gi", } overrides = SandboxConfigOverrides( image="alpine:3.20", root_disk_size="50Gi", limits={"memory": "16Gi"}, cmd=["sleep", "infinity"], ) result = core._rehydrate_config_from_snapshot( annotations=ann, new_name="r", namespace="default", overrides=overrides ) assert result.success config = result.data assert config.image == "alpine:3.20" assert config.root_disk_size == "50Gi" assert config.limits == {"memory": "16Gi"} assert config.cmd == ["sleep", "infinity"] # Backend not overridden → from annotation. assert config.backend == "kata-qemu-longhorn" def test_missing_image_and_no_override_errors_helpfully(self, core: K7Core): ann = {"k7.io/source-backend": "kata-qemu-longhorn"} # no image result = core._rehydrate_config_from_snapshot( annotations=ann, new_name="r", namespace="default", overrides=None ) assert not result.success assert "k7.io/source-image" in result.error assert "--image" in result.error def test_image_override_when_annotation_missing(self, core: K7Core): ann = {} overrides = SandboxConfigOverrides(image="alpine:3.20") result = core._rehydrate_config_from_snapshot( annotations=ann, new_name="r", namespace="default", overrides=overrides ) assert result.success assert result.data.image == "alpine:3.20" # Falls back to default backend when neither annotation nor override set. assert result.data.backend == "kata-qemu-longhorn" def test_docker_annotation_sets_docker_true(self, core: K7Core): ann = { "k7.io/source-image": "ubuntu:24.04", "k7.io/source-backend": "kata-qemu-longhorn", "k7.io/source-docker": "true", "k7.io/source-docker-disk": "20Gi", } result = core._rehydrate_config_from_snapshot( annotations=ann, new_name="r", namespace="default", overrides=None ) assert result.success assert result.data.docker is True assert result.data.docker_disk == "20Gi" def test_malformed_limits_annotation_falls_back_to_none(self, core: K7Core): ann = {"k7.io/source-image": "alpine:3.20", "k7.io/source-limits": "not json"} result = core._rehydrate_config_from_snapshot( annotations=ann, new_name="r", namespace="default", overrides=None ) assert result.success # SandboxConfig.__post_init__ normalises None → {} so we just check it didn't crash # and that the bad annotation didn't poison the limits dict. assert result.data.limits in (None, {}) # --- _collect_source_annotations ------------------------------------------- class TestCollectSourceAnnotations: async def test_reads_image_backend_sidecar_limits_pvc_size(self, core: K7Core): # Build a fake Deployment + PVC reachable through the apps_v1 / v1 clients. deployment = MagicMock() deployment.metadata = MagicMock() deployment.metadata.annotations = { "k7.katakate.org/backend": "kata-qemu-longhorn", "k7.katakate.org/sidecar": "docker", "k7.katakate.org/root-pvc-name": "demo-root-lh", } sandbox_container = MagicMock() sandbox_container.name = "sandbox" sandbox_container.image = "python:3.12-slim" sandbox_container.resources = MagicMock() sandbox_container.resources.limits = {"cpu": "2", "memory": "4Gi", "ephemeral-storage": "10Gi"} deployment.spec = MagicMock() deployment.spec.template = MagicMock() deployment.spec.template.spec = MagicMock() deployment.spec.template.spec.containers = [sandbox_container] mock_apps = AsyncMock() mock_apps.read_namespaced_deployment.return_value = deployment pvc = MagicMock() pvc.spec = MagicMock() pvc.spec.resources = MagicMock() pvc.spec.resources.requests = {"storage": "20Gi"} mock_v1 = AsyncMock() mock_v1.read_namespaced_persistent_volume_claim.return_value = pvc core._apps_v1_client = mock_apps core._core_v1_client = mock_v1 core._config_loaded = True ann = await core._collect_source_annotations("demo", "default") assert ann["k7.io/source-image"] == "python:3.12-slim" assert ann["k7.io/source-backend"] == "kata-qemu-longhorn" assert ann["k7.io/source-sidecar"] == "docker" assert ann["k7.io/source-root-disk-size"] == "20Gi" limits = json.loads(ann["k7.io/source-limits"]) assert limits == {"cpu": "2", "memory": "4Gi", "ephemeral-storage": "10Gi"} async def test_missing_deployment_returns_empty_dict(self, core: K7Core): mock_apps = AsyncMock() mock_apps.read_namespaced_deployment.side_effect = Exception("404") core._apps_v1_client = mock_apps core._config_loaded = True ann = await core._collect_source_annotations("ghost", "default") assert ann == {} # --- API: POST /api/v1/snapshots/{name}/restore ----------------------------- @pytest.fixture() def keys_file(tmp_path: Path): p = tmp_path / "api_keys.json" future_ts = int(time.time()) + 86400 p.write_text(json.dumps({TEST_KEY_HASH: {"name": "test-key", "expires": future_ts, "last_used": None}})) with patch("k7.api.main.API_KEYS_FILE", p): yield p async def _post(client, path, body): return await client.post(path, json=body, headers={"X-API-Key": TEST_KEY}) class TestRestoreApiEndpoint: async def test_success_returns_201_with_location(self, keys_file): fake = AsyncMock( return_value=OperationResult( success=True, message="Sandbox demo-r restored from snapshot demo-paused-1", data={"source_snapshot": "demo-paused-1"}, ) ) with patch("k7.api.main.K7Core") as core_cls: core_cls.return_value.restore_sandbox = fake transport = httpx.ASGITransport(app=app) async with httpx.AsyncClient(transport=transport, base_url="http://test") as client: resp = await _post( client, "/api/v1/snapshots/demo-paused-1/restore", { "new_sandbox_name": "demo-r", "namespace": "ns", "overrides": {"image": "alpine:3.20", "root_disk_size": "20Gi"}, "keep_snapshot": False, }, ) assert resp.status_code == 201 body = resp.json()["data"] assert body["name"] == "demo-r" assert body["namespace"] == "ns" assert body["source_snapshot"] == "demo-paused-1" assert resp.headers["location"] == "/api/v1/sandboxes/demo-r?namespace=ns" fake.assert_awaited_once() kwargs = fake.await_args.kwargs assert kwargs["snapshot_name"] == "demo-paused-1" assert kwargs["new_sandbox_name"] == "demo-r" assert kwargs["namespace"] == "ns" assert kwargs["keep_snapshot"] is False overrides: SandboxConfigOverrides = kwargs["overrides"] assert overrides.image == "alpine:3.20" assert overrides.root_disk_size == "20Gi" async def test_missing_new_sandbox_name_returns_400(self, keys_file): transport = httpx.ASGITransport(app=app) async with httpx.AsyncClient(transport=transport, base_url="http://test") as client: resp = await _post(client, "/api/v1/snapshots/snap/restore", {}) assert resp.status_code == 400 assert "new_sandbox_name is required" in resp.json()["error"]["message"] async def test_snapshot_not_found_returns_404(self, keys_file): fake = AsyncMock( return_value=OperationResult(success=False, error="VolumeSnapshot ghost not found in namespace default") ) with patch("k7.api.main.K7Core") as core_cls: core_cls.return_value.restore_sandbox = fake transport = httpx.ASGITransport(app=app) async with httpx.AsyncClient(transport=transport, base_url="http://test") as client: resp = await _post(client, "/api/v1/snapshots/ghost/restore", {"new_sandbox_name": "r"}) assert resp.status_code == 404 async def test_target_already_exists_returns_409(self, keys_file): fake = AsyncMock(return_value=OperationResult(success=False, error="Sandbox demo-r already exists")) with patch("k7.api.main.K7Core") as core_cls: core_cls.return_value.restore_sandbox = fake transport = httpx.ASGITransport(app=app) async with httpx.AsyncClient(transport=transport, base_url="http://test") as client: resp = await _post(client, "/api/v1/snapshots/s/restore", {"new_sandbox_name": "demo-r"}) assert resp.status_code == 409 async def test_overrides_filtered_to_allowed_keys(self, keys_file): fake = AsyncMock(return_value=OperationResult(success=True, message="ok")) with patch("k7.api.main.K7Core") as core_cls: core_cls.return_value.restore_sandbox = fake transport = httpx.ASGITransport(app=app) async with httpx.AsyncClient(transport=transport, base_url="http://test") as client: resp = await _post( client, "/api/v1/snapshots/s/restore", { "new_sandbox_name": "r", "overrides": { "image": "alpine:3.20", "bogus_field": "ignored", }, }, ) assert resp.status_code == 201 overrides: SandboxConfigOverrides = fake.await_args.kwargs["overrides"] assert overrides.image == "alpine:3.20" # bogus_field must not have leaked into the dataclass. assert not hasattr(overrides, "bogus_field") # --- SDK Client.restore ----------------------------------------------------- def _mock_response(payload: dict, status_code: int = 200) -> MagicMock: r = MagicMock() r.status_code = status_code r.json.return_value = payload r.raise_for_status.return_value = None return r class TestSdkClientRestore: def test_returns_proxy_and_posts_correct_body(self): c = Client(endpoint="http://api.test", api_key="k") with patch.object( c.session, "post", return_value=_mock_response({"data": {"name": "child", "source_snapshot": "snap"}}, 201), ) as post: proxy = c.restore( "snap", "child", namespace="ns", overrides={"image": "alpine:3.20"}, keep_snapshot=False, ) assert isinstance(proxy, SandboxProxy) assert proxy.name == "child" assert proxy.namespace == "ns" assert proxy._client is c assert post.call_args[0][0] == "http://api.test/api/v1/snapshots/snap/restore" body = post.call_args[1]["json"] assert body == { "new_sandbox_name": "child", "namespace": "ns", "keep_snapshot": False, "overrides": {"image": "alpine:3.20"}, } def test_omits_overrides_when_none(self): c = Client(endpoint="http://api.test", api_key="k") with patch.object(c.session, "post", return_value=_mock_response({"data": {"name": "child"}}, 201)) as post: c.restore("snap", "child") body = post.call_args[1]["json"] assert "overrides" not in body assert body["keep_snapshot"] is True # --- Snapshot creation stamps source annotations --------------------------- class TestCreateVolumeSnapshotStampsAnnotations: async def test_named_snapshot_calls_collect_source_annotations(self, core: K7Core): mock_custom = AsyncMock() core._custom_objects_client = mock_custom core._config_loaded = True collected = { "k7.io/source-image": "python:3.12-slim", "k7.io/source-backend": "kata-qemu-longhorn", } with patch.object( core, "_collect_source_annotations", new_callable=AsyncMock, return_value=collected, ) as collect: result = await core._create_volume_snapshot( pvc_name="demo-root-lh", snapshot_name="demo-v1", snapshot_class="longhorn", namespace="default", kind=SNAPSHOT_KIND_PAUSE, source_sandbox="demo", ) assert result.success collect.assert_awaited_once_with("demo", "default") mock_custom.create_namespaced_custom_object.assert_called_once() body = mock_custom.create_namespaced_custom_object.call_args.kwargs["body"] ann = body["metadata"]["annotations"] assert ann["k7.io/kind"] == SNAPSHOT_KIND_PAUSE assert ann["k7.io/source-sandbox"] == "demo" assert ann["k7.io/source-image"] == "python:3.12-slim" assert ann["k7.io/source-backend"] == "kata-qemu-longhorn" async def test_fork_snapshot_does_not_collect_source_annotations(self, core: K7Core): """Fork-temp snapshots are auto-deleted; no point in the extra K8s read.""" from k7.core.models import SNAPSHOT_KIND_FORK mock_custom = AsyncMock() core._custom_objects_client = mock_custom core._config_loaded = True with patch.object(core, "_collect_source_annotations", new_callable=AsyncMock, return_value={}) as collect: result = await core._create_volume_snapshot( pvc_name="demo-root-lh", snapshot_name="demo-fork-1", snapshot_class="longhorn", namespace="default", kind=SNAPSHOT_KIND_FORK, source_sandbox="demo", ) assert result.success collect.assert_not_awaited()