bot-bottle/tests/unit/test_smolmachines_provision.py

"""Unit: smolmachines provisioning helpers (PRD 0023 chunks 4a + 4d).

Tests mock `bottle.exec` / `bottle.cp_in` and assert on the
dispatched script shape. The real round-trip lives in the chunk-4
integration smoke."""

from __future__ import annotations

import subprocess
import tempfile
import unittest
from dataclasses import replace
from pathlib import Path
from unittest.mock import MagicMock, patch

from bot_bottle.agent_provider import (
    AgentProvisionCommand,
    AgentProvisionDir,
    AgentProvisionFile,
    AgentProvisionPlan,
)
from bot_bottle.backend import Bottle, BottleSpec, ExecResult
from bot_bottle.backend.smolmachines.bottle_plan import (
    SmolmachinesBottlePlan,
)
from bot_bottle.backend.smolmachines.provision import (
    ca as _ca,
    git as _git,
    workspace as _workspace,
)
from bot_bottle.backend.smolmachines.launch import _bundle_launch_spec
from bot_bottle.egress import EgressPlan, EgressRoute
from bot_bottle.git_gate import GitGatePlan, GitGateUpstream
from bot_bottle.manifest import GitEntry, Manifest
from bot_bottle.pipelock import PipelockProxyPlan
from bot_bottle.supervise import SupervisePlan
from bot_bottle.workspace import workspace_plan


def _make_bottle(
    name: str = "bot-bottle-demo-abc12",
    exec_result: ExecResult | None = None,
) -> MagicMock:
    bottle = MagicMock(spec=Bottle)
    bottle.name = name
    bottle.exec.return_value = (
        exec_result if exec_result is not None
        else ExecResult(returncode=0, stdout="", stderr="")
    )
    return bottle


def _exec_scripts(bottle: MagicMock) -> list[str]:
    """All script strings passed to bottle.exec, in call order."""
    return [c.args[0] for c in bottle.exec.call_args_list]


def _exec_users(bottle: MagicMock) -> list[str]:  # type: ignore
    """user= kwarg from each bottle.exec call, in order."""
    return [c.kwargs.get("user", "node") for c in bottle.exec.call_args_list]


def _plan(
    *,
    agent_prompt: str = "",
    skills: list[str] | None = None,
    git: list[GitEntry] = (),  # type: ignore
    git_user: dict | None = None,  # type: ignore
    copy_cwd: bool = False,
    user_cwd: str = "/tmp/x",
    stage_dir: Path | None = None,
    egress_routes: tuple[EgressRoute, ...] = (),
    egress_ca_path: Path = Path(),
    pipelock_ca_path: Path = Path(),
    supervise: bool = False,
    bundle_ip: str = "192.168.50.2",
    agent_git_gate_host: str = "127.0.0.1:55555",
    agent_supervise_url: str = "http://127.0.0.1:55556/",
    codex_auth_file: Path | None = None,
    agent_provider_template: str = "claude",
    guest_env: dict[str, str] | None = None,
) -> SmolmachinesBottlePlan:
    bottle_json: dict = {}  # type: ignore
    git_gate_json: dict = {}  # type: ignore
    if git:
        git_gate_json["repos"] = {
            g.Name: {
                "url": g.Upstream,
                "identity": g.IdentityFile,
            }
            for g in git
        }
    if git_user is not None:
        git_gate_json["user"] = git_user
    if git_gate_json:
        bottle_json["git-gate"] = git_gate_json
    if supervise:
        bottle_json["supervise"] = True
    manifest = Manifest.from_json_obj({
        "bottles": {"dev": bottle_json},
        "agents": {
            "demo": {
                "skills": list(skills or []),
                "prompt": agent_prompt,
                "bottle": "dev",
            },
        },
    })
    spec = BottleSpec(
        manifest=manifest,
        agent_name="demo",
        copy_cwd=copy_cwd,
        user_cwd=user_cwd,
    )
    supervise_plan = None
    if supervise:
        supervise_plan = SupervisePlan(
            slug="demo-abc12",
            queue_dir=Path("/tmp/queue"),
            current_config_dir=Path("/tmp/current-config"),
        )
    return SmolmachinesBottlePlan(
        guest_home="/home/node",
        spec=spec,
        stage_dir=stage_dir or Path("/tmp/stage"),
        slug="demo-abc12",
        bundle_subnet="192.168.50.0/24",
        bundle_gateway="192.168.50.1",
        bundle_ip=bundle_ip,
        machine_name="bot-bottle-demo-abc12",
        agent_image_ref="bot-bottle-claude:latest",
        guest_env=dict(guest_env or {}),
        prompt_file=Path("/tmp/state/demo-abc12/agent/prompt.txt"),
        proxy_plan=PipelockProxyPlan(
            yaml_path=Path("/tmp/pipelock.yaml"),
            slug="demo-abc12",
            ca_cert_host_path=pipelock_ca_path,
        ),
        git_gate_plan=GitGatePlan(
            slug="demo-abc12",
            entrypoint_script=Path("/tmp/git-gate-entrypoint.sh"),
            hook_script=Path("/tmp/git-gate-hook"),
            access_hook_script=Path("/tmp/git-gate-access-hook"),
            upstreams=(),
        ),
        egress_plan=EgressPlan(
            slug="demo-abc12",
            routes_path=Path("/tmp/routes.yaml"),
            routes=egress_routes,
            token_env_map={},
            mitmproxy_ca_cert_only_host_path=egress_ca_path,
        ),
        supervise_plan=supervise_plan,
        agent_git_gate_host=agent_git_gate_host,
        agent_supervise_url=agent_supervise_url,
        agent_provision=_agent_provision(
            agent_provider_template,
            codex_auth_file=codex_auth_file,
            guest_env=dict(guest_env or {}),
        ),
        workspace_plan=workspace_plan(spec, guest_home="/home/node"),
    )


def _agent_provision(
    template: str,
    *,
    codex_auth_file: Path | None = None,
    guest_env: dict[str, str] | None = None,
) -> AgentProvisionPlan:
    if template != "codex":
        return AgentProvisionPlan(
            template=template,
            command=template,
            prompt_mode="append_file",
            image="",
            dockerfile="",
            guest_env=dict(guest_env or {}),
        )
    auth_dir = (guest_env or {}).get("CODEX_HOME", "/home/node/.codex")
    files = [
        AgentProvisionFile(
            Path("/tmp/codex-config.toml"),
            f"{auth_dir}/config.toml",
        ),
    ]
    pre_copy: tuple[AgentProvisionCommand, ...] = ()
    verify: tuple[AgentProvisionCommand, ...] = ()
    if codex_auth_file is not None:
        files.append(AgentProvisionFile(codex_auth_file, f"{auth_dir}/auth.json"))
        pre_copy = (AgentProvisionCommand((
            "find", auth_dir,
            "-maxdepth", "1",
            "-type", "f",
            "(",
            "-name", "*.sqlite",
            "-o", "-name", "*.sqlite-*",
            "-o", "-name", "*.codex-repair-*.bak",
            ")",
            "-delete",
        ), "codex host credentials: could not reset runtime db files"),)
        verify = (AgentProvisionCommand((
            "runuser", "-u", "node", "--",
            "env",
            "HOME=/home/node",
            f"CODEX_HOME={auth_dir}",
            "codex", "login", "status",
        ), "codex host credentials: dummy auth was copied into the guest"),)
    return AgentProvisionPlan(
        template="codex",
        command="codex",
        prompt_mode="read_prompt_file",
        image="bot-bottle-codex:latest",
        dockerfile="",
        guest_env=dict(guest_env or {}),
        dirs=(AgentProvisionDir(auth_dir),),
        files=tuple(files),
        pre_copy=pre_copy,
        verify=verify,
    )


def _write_self_signed_cert(path: Path) -> None:
    """Drop a real self-signed PEM at `path` so provision_ca's
    fingerprint computation (PEM_cert_to_DER_cert + sha256) has
    actual bytes to chew on. Generated once per test via openssl."""
    subprocess.run(
        ["openssl", "req", "-x509", "-newkey", "rsa:2048", "-nodes",
         "-keyout", "/dev/null",
         "-out", str(path),
         "-days", "1",
         "-subj", "/CN=test"],
        check=True, capture_output=True,
    )


class TestProvisionCA(unittest.TestCase):
    """provision_ca selects the right CA cert (egress when the
    bottle has routes, else pipelock) and dispatches
    cp_in + exec in the right order."""

    def setUp(self):
        self._tmp = tempfile.TemporaryDirectory(prefix="cb-prov-ca.")
        self.tmp = Path(self._tmp.name)
        self.pipelock_ca = self.tmp / "pipelock-ca.pem"
        self.egress_ca = self.tmp / "egress-ca.pem"
        _write_self_signed_cert(self.pipelock_ca)
        _write_self_signed_cert(self.egress_ca)

    def tearDown(self):
        self._tmp.cleanup()

    # provision_ca dies hard if update-ca-certificates' exit
    # is non-zero; supply a stock success return so the bulk of
    # the tests below exercise the happy path.
    _UPDATE_OK = ExecResult(
        returncode=0,
        stdout="Updating certificates in /etc/ssl/certs...\n1 added, 0 removed; done.\n",
        stderr="",
    )

    def test_pipelock_path_when_no_routes(self):
        plan = _plan(pipelock_ca_path=self.pipelock_ca)
        bottle = _make_bottle(exec_result=self._UPDATE_OK)
        _ca.provision_ca(plan, bottle)
        bottle.cp_in.assert_called_once_with(
            str(self.pipelock_ca),
            _ca.AGENT_CA_PATH,
        )
        # chmod + chown + update-ca-certificates are folded into
        # one exec invocation; look at the single exec's script
        # rather than expecting separate calls.
        bottle.exec.assert_called_once()
        script = bottle.exec.call_args.args[0]
        self.assertIn("chmod 644", script)
        self.assertIn("update-ca-certificates", script)
        self.assertEqual("root", bottle.exec.call_args.kwargs.get("user"))

    def test_egress_path_when_routes_declared(self):
        plan = _plan(
            egress_routes=(EgressRoute(host="api.anthropic.com"),),
            egress_ca_path=self.egress_ca,
            pipelock_ca_path=self.pipelock_ca,
        )
        bottle = _make_bottle(exec_result=self._UPDATE_OK)
        _ca.provision_ca(plan, bottle)
        # When routes are declared, egress is the agent's first hop,
        # so egress's CA is the one that gets installed.
        bottle.cp_in.assert_called_once_with(
            str(self.egress_ca),
            _ca.AGENT_CA_PATH,
        )

    def test_retries_smolvm_sigkill_during_update_ca(self):
        plan = _plan(pipelock_ca_path=self.pipelock_ca)
        killed = ExecResult(
            returncode=137,
            stdout="Updating certificates in /etc/ssl/certs...\n",
            stderr="",
        )
        bottle = _make_bottle()
        bottle.exec.side_effect = [killed, self._UPDATE_OK]
        with patch(
            "bot_bottle.backend.smolmachines.provision.ca.time.sleep"
        ) as sleep:
            _ca.provision_ca(plan, bottle)

        self.assertEqual(2, bottle.exec.call_count)
        sleep.assert_called_once_with(1.0)

    def test_dies_when_selected_cert_missing(self):
        # Plan claims a pipelock cert at a path that doesn't exist —
        # something went wrong in launch's pipelock_tls_init.
        plan = _plan(pipelock_ca_path=self.tmp / "does-not-exist.pem")
        bottle = _make_bottle()
        with self.assertRaises(SystemExit):
            _ca.provision_ca(plan, bottle)


class TestProvisionGit(unittest.TestCase):
    """provision_git dispatches two independent passes (cwd .git
    copy + gitconfig insteadOf write); each no-ops on its own
    when its condition doesn't hold."""

    def setUp(self):
        self._tmp = tempfile.TemporaryDirectory(prefix="cb-prov-git.")
        self.stage = Path(self._tmp.name)

    def tearDown(self):
        self._tmp.cleanup()

    def test_noop_when_no_cwd_and_no_git_entries(self):
        bottle = _make_bottle()
        _git.provision_git(_plan(stage_dir=self.stage), bottle)
        bottle.cp_in.assert_not_called()
        bottle.exec.assert_not_called()

    def test_copies_cwd_git_when_copy_cwd_and_git_present(self):
        # Stage a fake host .git dir under user_cwd so the path-
        # check in _provision_cwd_git fires.
        cwd = self.stage / "cwd"
        (cwd / ".git").mkdir(parents=True)
        plan = _plan(
            copy_cwd=True, user_cwd=str(cwd), stage_dir=self.stage,
        )
        bottle = _make_bottle()
        _git.provision_git(plan, bottle)
        bottle.cp_in.assert_called_once_with(
            f"{cwd}/.git",
            "/home/node/workspace/.git",
        )
        scripts = _exec_scripts(bottle)
        self.assertTrue(any("mkdir -p" in s and "/home/node/workspace" in s for s in scripts))
        # chown the workspace tree so the agent (node) owns it.
        self.assertTrue(
            any("chown -R" in s and "node:node" in s and "/home/node/workspace/.git" in s
                for s in scripts)
        )

    def test_skips_cwd_when_copy_cwd_false(self):
        plan = _plan(copy_cwd=False, stage_dir=self.stage)
        bottle = _make_bottle()
        _git.provision_git(plan, bottle)
        bottle.cp_in.assert_not_called()

    def test_writes_gitconfig_with_ip_port_form_for_smolmachines(self):
        # Smolmachines's TSI-allowlisted guest dials git-gate via
        # smart HTTP at `127.0.0.1:<host port>` — the bundle's
        # git HTTP port is published on host loopback at launch
        # time, and the plan carries the discovered host port.
        plan = _plan(
            git=[GitEntry(
                Name="bot-bottle",
                Upstream="ssh://git@host/repo.git",
                IdentityFile="~/.ssh/id_ed25519",
            )],
            stage_dir=self.stage,
            agent_git_gate_host="127.0.0.1:9418",
        )
        bottle = _make_bottle()
        _git.provision_git(plan, bottle)
        # The staged gitconfig path is whatever NamedTemporaryFile
        # picked; we read its contents.
        cp_call = bottle.cp_in.call_args
        staged_path = Path(cp_call.args[0])
        self.assertEqual(self.stage, staged_path.parent)
        content = staged_path.read_text()
        self.assertIn(
            '[url "http://127.0.0.1:9418/bot-bottle.git"]', content,
        )
        self.assertIn(
            "\tinsteadOf = ssh://git@host/repo.git", content,
        )


class TestBundleLaunchSpec(unittest.TestCase):
    def test_git_gate_uses_http_daemon_for_smolmachines(self):
        plan = _plan()
        plan = replace(
            plan,
            git_gate_plan=replace(
                plan.git_gate_plan,
                upstreams=(GitGateUpstream(
                    name="bot-bottle",
                    upstream_url="ssh://git@host/repo.git",
                    upstream_host="host",
                    upstream_port="22",
                    identity_file="/tmp/key",
                    known_host_key="",
                ),),
            ),
        )

        spec = _bundle_launch_spec(plan, "net", "127.0.0.16")

        self.assertEqual(
            "egress,pipelock,git-gate,git-http",
            spec.daemons_csv,
        )
        self.assertIn(9420, spec.ports_to_publish)
        self.assertNotIn(9418, spec.ports_to_publish)


class TestProvisionGitUser(unittest.TestCase):
    """`_provision_git_user` runs `git config --global` inside the
    guest as the node user. SmolmachinesBottle.exec sets HOME and
    USER automatically for the requested user, so --global lands
    in /home/node/.gitconfig. No-op when the bottle didn't declare
    git_user (issue #86)."""

    def _git_config_calls(self, bottle: MagicMock) -> list[tuple[str, str]]:
        """Filter bottle.exec calls down to git-config invocations,
        return list of (script, user) tuples."""
        out = []
        for c in bottle.exec.call_args_list:
            script = c.args[0] if c.args else ""
            user = c.kwargs.get("user", "node")
            if "git config" in script:
                out.append((script, user))
        return out

    def test_noop_when_no_git_user(self):
        bottle = _make_bottle()
        _git._provision_git_user(_plan(), bottle)
        self.assertEqual([], self._git_config_calls(bottle))

    def test_sets_name_and_email_as_node(self):
        plan = _plan(git_user={
            "name": "Eric Bauerfeld",
            "email": "eric@dideric.is",
        })
        bottle = _make_bottle()
        _git._provision_git_user(plan, bottle)
        calls = self._git_config_calls(bottle)
        self.assertEqual(2, len(calls))
        # Both run as node so SmolmachinesBottle.exec sets HOME=/home/node
        # automatically, ensuring --global writes to /home/node/.gitconfig.
        for script, user in calls:
            self.assertEqual("node", user)
            self.assertIn("git config --global", script)
        self.assertIn("user.name", calls[0][0])
        self.assertIn("Eric Bauerfeld", calls[0][0])
        self.assertIn("user.email", calls[1][0])
        self.assertIn("eric@dideric.is", calls[1][0])

    def test_name_only(self):
        plan = _plan(git_user={"name": "Bot"})
        bottle = _make_bottle()
        _git._provision_git_user(plan, bottle)
        calls = self._git_config_calls(bottle)
        self.assertEqual(1, len(calls))
        self.assertIn("user.name", calls[0][0])
        self.assertIn("Bot", calls[0][0])

    def test_email_only(self):
        plan = _plan(git_user={"email": "bot@example.com"})
        bottle = _make_bottle()
        _git._provision_git_user(plan, bottle)
        calls = self._git_config_calls(bottle)
        self.assertEqual(1, len(calls))
        self.assertIn("user.email", calls[0][0])
        self.assertIn("bot@example.com", calls[0][0])


class TestProvisionWorkspace(unittest.TestCase):
    def setUp(self):
        self._tmp = tempfile.TemporaryDirectory(prefix="cb-prov-workspace.")
        self.stage = Path(self._tmp.name)

    def tearDown(self):
        self._tmp.cleanup()

    def test_noop_when_copy_cwd_false(self):
        plan = _plan(copy_cwd=False, stage_dir=self.stage)
        bottle = _make_bottle()
        _workspace.provision_workspace(plan, bottle)
        bottle.cp_in.assert_not_called()
        bottle.exec.assert_not_called()

    def test_copies_workspace_to_plan_path_and_chowns(self):
        cwd = self.stage / "cwd"
        cwd.mkdir()
        plan = _plan(copy_cwd=True, user_cwd=str(cwd), stage_dir=self.stage)
        bottle = _make_bottle()
        _workspace.provision_workspace(plan, bottle)

        bottle.cp_in.assert_called_once_with(
            str(cwd),
            "/home/node/workspace",
        )
        scripts = _exec_scripts(bottle)
        self.assertTrue(
            any("rm -rf /home/node/workspace" in s and "mkdir -p /home/node" in s
                for s in scripts)
        )
        self.assertTrue(
            any("chown -R node:node /home/node/workspace" in s
                and "chmod 755 /home/node/workspace" in s
                for s in scripts)
        )


if __name__ == "__main__":
    unittest.main()