bot-bottle/tests/unit/test_firecracker_infra_vm.py

"""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 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])
        self.assertEqual("-infra", build.call_args.kwargs["variant"])
        # 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)


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_builds_gateway_then_infra(self):
        with 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]
        self.assertEqual([infra_vm._GATEWAY_IMAGE, infra_vm._INFRA_IMAGE], tags)


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()