"""Unit: the prebuilt infra-rootfs artifact pull (PRD 0069 Stage 2). The launch-host half — version hashing and download/verify/decompress — is what keeps a docker-free host from booting a stale or corrupted rootfs, so the checksum + fail-closed paths are locked here. Network is mocked; no Docker. """ from __future__ import annotations import gzip import hashlib import io import os import tempfile import unittest import urllib.error import urllib.request from pathlib import Path from unittest import mock from bot_bottle.backend.firecracker import infra_artifact as ia from bot_bottle.log import Die def _gz(data: bytes) -> bytes: return gzip.compress(data) class _FakeNet: """Map artifact URLs to bytes (or an HTTPError) for urlopen.""" def __init__(self, responses: "dict[str, bytes | Exception]") -> None: self._responses = responses self.calls: list[str] = [] def urlopen(self, req: urllib.request.Request, *a: object, **k: object) -> io.BytesIO: url = req.full_url self.calls.append(url) val = self._responses.get(url) if isinstance(val, Exception): raise val if val is None: raise urllib.error.HTTPError(url, 404, "not found", {}, None) # type: ignore[arg-type] return io.BytesIO(val) class _CacheMixin(unittest.TestCase): def setUp(self) -> None: self._tmp = tempfile.TemporaryDirectory() self._env = mock.patch.dict( os.environ, {"BOT_BOTTLE_FC_CACHE": self._tmp.name, "BOT_BOTTLE_INFRA_ARTIFACT_TOKEN": ""}, clear=False, ) self._env.start() self.addCleanup(self._env.stop) self.addCleanup(self._tmp.cleanup) def _serve(self, version: str, gz_bytes: bytes, sha_text: str | None = None): if sha_text is None: sha_text = f"{hashlib.sha256(gz_bytes).hexdigest()} rootfs.ext4.gz\n" net = _FakeNet({ ia.artifact_url(version, "rootfs.ext4.gz"): gz_bytes, ia.artifact_url(version, "rootfs.ext4.gz.sha256"): sha_text.encode(), }) return mock.patch.object(ia.urllib.request, "urlopen", net.urlopen), net class TestVersion(unittest.TestCase): def test_deterministic_16_hex(self) -> None: v = ia.infra_artifact_version("#!/bin/sh\ntrue\n") self.assertEqual(v, ia.infra_artifact_version("#!/bin/sh\ntrue\n")) self.assertEqual(16, len(v)) int(v, 16) # hex def test_init_change_bumps_version(self) -> None: self.assertNotEqual( ia.infra_artifact_version("a"), ia.infra_artifact_version("b")) class TestEnsureArtifact(_CacheMixin): def test_downloads_verifies_and_caches(self) -> None: version = "deadbeef00000000" gz = _gz(b"fake ext4 bytes") patcher, net = self._serve(version, gz) with patcher: path = ia.ensure_artifact_gz(version) self.assertTrue(path.is_file()) self.assertEqual(gz, path.read_bytes()) first_calls = len(net.calls) # Second call is a cache hit — no further network. ia.ensure_artifact_gz(version) self.assertEqual(first_calls, len(net.calls)) def test_checksum_mismatch_fails_closed(self) -> None: version = "beefbeefbeefbeef" gz = _gz(b"payload") patcher, _ = self._serve(version, gz, sha_text="0" * 64 + " rootfs.ext4.gz\n") with patcher: with self.assertRaises(Die) as ctx: ia.ensure_artifact_gz(version) self.assertIn("checksum mismatch", str(ctx.exception.message)) # nothing left cached to accidentally boot self.assertFalse((ia._cache_root(version) / "rootfs.ext4.gz").exists()) def test_missing_artifact_points_at_publish(self) -> None: version = "0000000000000000" net = _FakeNet({}) # everything 404s with mock.patch.object(ia.urllib.request, "urlopen", net.urlopen): with self.assertRaises(Die) as ctx: ia.ensure_artifact_gz(version) self.assertIn("publish_infra", str(ctx.exception.message)) def test_materialize_gunzips_to_dest(self) -> None: version = "1234123412341234" raw = b"the real rootfs contents" * 100 patcher, _ = self._serve(version, _gz(raw)) with patcher, tempfile.TemporaryDirectory() as d: dest = Path(d) / "rootfs.ext4" ia.materialize_ext4(version, dest) self.assertEqual(raw, dest.read_bytes()) class TestConfig(unittest.TestCase): def test_base_and_owner_overridable(self) -> None: with mock.patch.dict(os.environ, { "BOT_BOTTLE_INFRA_ARTIFACT_BASE": "https://mirror.example/", "BOT_BOTTLE_INFRA_ARTIFACT_OWNER": "acme", }): url = ia.artifact_url("v1", "rootfs.ext4.gz") self.assertEqual( "https://mirror.example/api/packages/acme/generic/" "bot-bottle-infra/v1/rootfs.ext4.gz", url) def test_local_build_flag(self) -> None: with mock.patch.dict(os.environ, {"BOT_BOTTLE_INFRA_BUILD": "local"}): self.assertTrue(ia.local_build_requested()) with mock.patch.dict(os.environ, {"BOT_BOTTLE_INFRA_BUILD": ""}): self.assertFalse(ia.local_build_requested()) if __name__ == "__main__": unittest.main()