"""Unit tests for the Firecracker infra VM (control-plane VM boot). The KVM boot / HTTP reachability is integration-tested on a KVM host; here we cover the URL shape, the rootfs-variant wiring, and the health-poll decisions that must hold without a VM. """ from __future__ import annotations import os import unittest from pathlib import Path from unittest.mock import MagicMock, patch from bot_bottle.backend.firecracker import infra_vm class TestControlPlaneUrl(unittest.TestCase): def test_url_uses_guest_ip_and_port(self): infra = infra_vm.InfraVm( vm=MagicMock(), guest_ip="10.243.255.1", private_key=Path("/k")) self.assertEqual( f"http://10.243.255.1:{infra_vm.CONTROL_PLANE_PORT}", infra.control_plane_url, ) class TestBuildInfraRootfs(unittest.TestCase): def test_uses_infra_variant_and_init(self): with patch.object(infra_vm.util, "build_base_rootfs_dir") as build: build.return_value = Path("/cache/rootfs/x-infra") infra_vm.build_infra_rootfs_dir() build.assert_called_once() self.assertEqual(infra_vm._INFRA_IMAGE, build.call_args.args[0]) # variant is "-infra-" so an init change rebuilds the rootfs. self.assertTrue(build.call_args.kwargs["variant"].startswith("-infra-")) # The init runs BOTH the control plane and the gateway data plane, # and exports PATH so gateway_init's subprocess daemons find python3. init = build.call_args.kwargs["init_script"] self.assertIn("bot_bottle.orchestrator", init) self.assertIn("gateway_init.py", init) self.assertIn("export PATH=", init) # Persistent registry volume mounted at the DB dir before the CP starts. self.assertIn("/dev/vdb", init) # VM backend uses git-http (9420); the git:// daemon is left out. self.assertIn("BOT_BOTTLE_GATEWAY_DAEMONS=egress,git-http,supervise", init) class TestSshGatewayTransport(unittest.TestCase): def test_cp_into_preserves_source_mode(self): import os import tempfile from subprocess import CompletedProcess with tempfile.NamedTemporaryFile() as f: os.chmod(f.name, 0o700) # like the staged access-hook t = infra_vm.SshGatewayTransport(Path("/k"), "10.0.0.1") with patch.object(infra_vm.subprocess, "run", return_value=CompletedProcess([], 0)) as run: t.cp_into(f.name, "/etc/git-gate/access-hook") remote_cmd = run.call_args.args[0][-1] self.assertIn("chmod 700", remote_cmd) # exec bit preserved over SSH def test_exec_raises_on_failure(self): from subprocess import CompletedProcess t = infra_vm.SshGatewayTransport(Path("/k"), "10.0.0.1") with patch.object(infra_vm.subprocess, "run", return_value=CompletedProcess([], 1, stderr="nope")), \ self.assertRaises(infra_vm.GatewayProvisionError): t.exec(["mkdir", "-p", "/git-gate"]) class TestRegistryVolume(unittest.TestCase): def test_reuses_existing_volume(self): import tempfile with tempfile.TemporaryDirectory() as td: vol = Path(td) / "registry.ext4" vol.write_bytes(b"") # already present with patch.object(infra_vm, "registry_volume_path", return_value=vol), \ patch.object(infra_vm.subprocess, "run") as run: out = infra_vm._ensure_registry_volume() run.assert_not_called() # no mke2fs when it exists self.assertEqual(vol, out) def test_creates_volume_when_missing(self): import tempfile from subprocess import CompletedProcess with tempfile.TemporaryDirectory() as td: vol = Path(td) / "registry.ext4" with patch.object(infra_vm, "registry_volume_path", return_value=vol), \ patch.object(infra_vm.subprocess, "run", return_value=CompletedProcess([], 0)) as run: infra_vm._ensure_registry_volume() argv = run.call_args.args[0] self.assertIn("mke2fs", argv) self.assertIn(str(vol), argv) class TestEnsureBuilt(unittest.TestCase): def test_default_pulls_artifact_without_docker(self): # PRD 0069 Stage 2: the launch host pulls the prebuilt rootfs; no Docker. with patch.object(infra_vm.docker_mod, "build_image") as build, \ patch.object(infra_vm.infra_artifact, "ensure_artifact_gz") as pull: infra_vm.ensure_built() build.assert_not_called() pull.assert_called_once() def test_local_mode_builds_deps_before_infra(self): with patch.dict(os.environ, {"BOT_BOTTLE_INFRA_BUILD": "local"}), \ patch.object(infra_vm.docker_mod, "build_image") as build: infra_vm.ensure_built() tags = [c.args[0] for c in build.call_args_list] # infra is FROM gateway and COPY --from orchestrator, so both first. self.assertEqual(infra_vm._INFRA_IMAGE, tags[-1]) self.assertIn(infra_vm._ORCHESTRATOR_IMAGE, tags[:-1]) self.assertIn(infra_vm._GATEWAY_IMAGE, tags[:-1]) class TestWaitForHealth(unittest.TestCase): def _infra(self, alive: bool = True) -> infra_vm.InfraVm: vm = MagicMock() vm.is_alive.return_value = alive return infra_vm.InfraVm(vm=vm, guest_ip="10.0.0.1", private_key=Path("/k")) def test_returns_on_200(self): infra = self._infra() cm = MagicMock() cm.__enter__.return_value.status = 200 with patch.object(infra_vm.urllib.request, "urlopen", return_value=cm): infra_vm.wait_for_health(infra, timeout=5) # must not raise def test_dies_when_vm_exits(self): infra = self._infra(alive=False) assert infra.vm is not None # narrow for the type checker (it's a mock) infra.vm.process.returncode = 1 with patch.object(infra_vm.firecracker_vm, "_console_tail", return_value=""), \ self.assertRaises(SystemExit): infra_vm.wait_for_health(infra, timeout=5) class TestEnsureRunningSingleton(unittest.TestCase): def test_adopts_when_healthy(self): # A healthy control plane + existing key -> adopt (no boot), vm=None. with patch.object(infra_vm, "_health_ok", return_value=True), \ patch.object(infra_vm, "_infra_dir") as d, \ patch.object(infra_vm, "boot") as boot: keydir = MagicMock() (keydir / "id_ed25519").exists.return_value = True d.return_value = keydir infra = infra_vm.ensure_running() boot.assert_not_called() self.assertIsNone(infra.vm) def test_boots_when_unhealthy(self): with patch.object(infra_vm, "_health_ok", return_value=False), \ patch.object(infra_vm, "stop") as stop, \ patch.object(infra_vm, "ensure_built") as built, \ patch.object(infra_vm, "boot") as boot, \ patch.object(infra_vm, "wait_for_health") as wait: boot.return_value = infra_vm.InfraVm( guest_ip="10.243.255.1", private_key=Path("/k"), vm=MagicMock()) infra_vm.ensure_running() stop.assert_called_once() # clear a stale VM first built.assert_called_once() boot.assert_called_once() wait.assert_called_once() class TestKillPidfile(unittest.TestCase): def test_noop_when_no_pidfile(self): import tempfile with tempfile.TemporaryDirectory() as td: with patch.object(infra_vm, "_pid_file", return_value=Path(td) / "vm.pid"), \ patch.object(infra_vm.os, "kill") as kill: infra_vm._kill_pidfile() # must not raise kill.assert_not_called() def test_skips_dead_or_recycled_pid(self): import tempfile with tempfile.TemporaryDirectory() as td: pidf = Path(td) / "vm.pid" pidf.write_text("999999") # a PID that isn't a live firecracker with patch.object(infra_vm, "_pid_file", return_value=pidf), \ patch.object(infra_vm.os, "kill") as kill: infra_vm._kill_pidfile() kill.assert_not_called() if __name__ == "__main__": unittest.main()