box/tests/test_fleet_heal.py

"""Tests for `box fleet heal` and `box watchdog` in super-cli.py.

Heal drives lock cleanup, watchdog-timer install, tunnel restart, and
egress/relay verification; watchdog status/run expose the systemd
watchdog layer. Shell-outs are faked; no sudo/systemctl/netns touch
the host.
"""
import argparse
import importlib.util
import io
import json
import sys
import unittest
from contextlib import redirect_stdout
from pathlib import Path
from unittest import mock

REPO_ROOT = Path(__file__).resolve().parent.parent
SPEC = importlib.util.spec_from_file_location(
    "super_cli_heal", REPO_ROOT / "bin" / "super-cli.py")
super_cli = importlib.util.module_from_spec(SPEC)
SPEC.loader.exec_module(super_cli)


def _ns(**over):
    kw = dict(node=None, target=None, json=True)
    kw.update(over)
    return argparse.Namespace(**kw)


class FakeSh:
    """Programmable stand-in for super_cli._sh; records calls."""

    def __init__(self):
        self.calls = []
        self.handlers = []

    def on(self, *needles, rc=0, out=""):
        self.handlers.append((needles, rc, out))
        return self

    def __call__(self, cmd, timeout=120, input_text=None):
        self.calls.append((list(cmd), timeout, input_text))
        blob = " ".join(cmd)
        for needles, rc, out in self.handlers:
            if all(n in blob for n in needles):
                return rc, out() if callable(out) else out
        raise AssertionError("unexpected command: %r" % (cmd,))


class HealBase(unittest.TestCase):
    def setUp(self):
        import tempfile
        self.tmp = tempfile.TemporaryDirectory()
        self.addCleanup(self.tmp.cleanup)
        self.root = Path(self.tmp.name)
        self.systemd = self.root / "systemd"
        self.systemd.mkdir()
        self.locks = self.root / "locks"
        self.locks.mkdir()
        self.patchers = [
            mock.patch.object(super_cli, "SYSTEMD_SYSTEM_DIR", self.systemd),
            mock.patch.object(super_cli, "WATCHDOG_LOCK_TMPL",
                              str(self.locks / "chromebox-watchdog-{node}.lock")),
        ]
        for p in self.patchers:
            p.start()
        self.addCleanup(self._unpatch)
        self.sh = FakeSh()
        self.sh_mock = mock.patch.object(super_cli, "_sh", self.sh)
        self.sh_mock.start()
        self.addCleanup(self.sh_mock.stop)

    def _unpatch(self):
        for p in self.patchers:
            p.stop()

    def run_heal(self, node="dev", as_json=True):
        buf = io.StringIO()
        with redirect_stdout(buf):
            with self.assertRaises(SystemExit) as cm:
                super_cli.cmd_fleet_heal(_ns(node=node, json=as_json))
        return cm.exception.code, buf.getvalue()


class FleetHealTests(HealBase):
    def _healthy_sh(self):
        (self.sh
         .on("systemctl", "enable", "--now", rc=0, out="")
         .on("netvm-node-up.sh", "dev", rc=0,
             out="tunnel already up (egress=1.2.3.4), skipping handshake wait\n"
                 "node=dev netns=warp-dev egress=1.2.3.4")
         .on("curl", rc=0, out="ip=1.2.3.4\nfoo=bar"))

    def test_heal_recovered(self):
        (self.systemd / "chromebox-watchdog-dev.timer").write_text("x")
        self._healthy_sh()
        with mock.patch.object(super_cli, "probe_cdp_status",
                               return_value={"ok": True, "latency_ms": 12}):
            code, out = self.run_heal()
        self.assertEqual(code, 0)
        data = json.loads(out)
        self.assertEqual(data["verdict"], "RECOVERED")
        self.assertTrue(data["ok"])
        self.assertEqual([s["step"] for s in data["steps"]],
                         ["lock", "timer", "tunnel", "egress", "relay"])
        self.assertTrue(all(s["ok"] for s in data["steps"]))

    def test_heal_down_when_egress_fails(self):
        (self.systemd / "chromebox-watchdog-dev.timer").write_text("x")
        (self.sh
         .on("systemctl", "enable", "--now", rc=0, out="")
         .on("netvm-node-up.sh", "dev", rc=0,
             out="no handshake yet (endpoint=162.159.192.1)\n"
                 "node=dev netns=warp-dev egress=unknown")
         .on("curl", rc=7, out="curl: (7) couldn't connect"))
        with mock.patch.object(super_cli, "probe_cdp_status",
                               return_value={"ok": False, "error": "refused",
                                             "latency_ms": None}):
            code, out = self.run_heal()
        self.assertEqual(code, 1)
        data = json.loads(out)
        self.assertEqual(data["verdict"], "DOWN")
        self.assertFalse(data["ok"])
        by_step = {s["step"]: s for s in data["steps"]}
        self.assertFalse(by_step["tunnel"]["ok"])
        self.assertFalse(by_step["egress"]["ok"])

    def test_heal_down_prints_identity_guidance(self):
        (self.systemd / "chromebox-watchdog-dev.timer").write_text("x")
        (self.sh
         .on("systemctl", rc=0, out="")
         .on("netvm-node-up.sh", rc=0, out="node=dev egress=unknown")
         .on("curl", rc=7, out="fail"))
        with mock.patch.object(super_cli, "probe_cdp_status",
                               return_value={"ok": False, "error": "x",
                                             "latency_ms": None}):
            code, out = self.run_heal(as_json=False)
        self.assertEqual(code, 1)
        self.assertIn("netvm-new-identity.sh dev", out)

    def test_heal_installs_missing_timer(self):
        seen = {}

        def fake(cmd, timeout=120, input_text=None):
            blob = " ".join(cmd)
            self.sh.calls.append((list(cmd), timeout, input_text))
            if "tee" in blob:
                seen["tee_target"] = cmd[-1]
                seen["tee_input"] = input_text
                return 0, ""
            if "daemon-reload" in blob:
                seen["reload"] = True
                return 0, ""
            if "enable" in blob:
                return 0, ""
            if "netvm-node-up.sh" in blob:
                return 0, "node=dev egress=1.2.3.4"
            if "curl" in blob:
                return 0, "ip=1.2.3.4"
            raise AssertionError("unexpected: %r" % (cmd,))

        with mock.patch.object(super_cli, "_sh", fake):
            with mock.patch.object(
                    super_cli, "probe_cdp_status",
                    return_value={"ok": True, "latency_ms": 3}):
                code, _ = self.run_heal()
        self.assertEqual(code, 0)
        self.assertEqual(seen["tee_target"],
                         str(self.systemd / "chromebox-watchdog-dev.timer"))
        self.assertIn("Unit=chromebox-watchdog@dev.service", seen["tee_input"])
        self.assertTrue(seen["reload"])

    def test_heal_clears_unwritable_lock(self):
        (self.systemd / "chromebox-watchdog-dev.timer").write_text("x")
        lock = self.locks / "chromebox-watchdog-dev.lock"
        lock.write_text("")
        lock.chmod(0o444)
        self._healthy_sh()
        self.sh.on("rm", str(lock), rc=0, out="")
        with mock.patch.object(super_cli, "probe_cdp_status",
                               return_value={"ok": True, "latency_ms": 3}):
            code, out = self.run_heal()
        self.assertEqual(code, 0)
        data = json.loads(out)
        by_step = {s["step"]: s for s in data["steps"]}
        self.assertTrue(by_step["lock"]["ok"])
        self.assertIn("removed stale lock", by_step["lock"]["detail"])
        self.assertTrue(any("rm" in " ".join(c[0]) for c in self.sh.calls))

    def test_heal_rejects_unknown_node(self):
        buf = io.StringIO()
        with redirect_stdout(buf):
            with self.assertRaises(SystemExit) as cm:
                super_cli.cmd_fleet_heal(_ns(node="ghost"))
        self.assertEqual(cm.exception.code, 1)


class ShHelperTests(unittest.TestCase):
    def test_timeout_and_oserror(self):
        import subprocess as real_subprocess
        with mock.patch.object(super_cli.subprocess, "run",
                               side_effect=real_subprocess.TimeoutExpired("x", 1)):
            self.assertEqual(super_cli._sh(["x"], timeout=1)[0], 124)
        with mock.patch.object(super_cli.subprocess, "run",
                               side_effect=OSError("nope")):
            self.assertEqual(super_cli._sh(["x"])[0], 127)


class WatchdogTests(unittest.TestCase):
    def setUp(self):
        self.sh = FakeSh()
        self.sh_mock = mock.patch.object(super_cli, "_sh", self.sh)
        self.sh_mock.start()
        self.addCleanup(self.sh_mock.stop)

    def _unit_states(self, missing=()):
        def fake(cmd, timeout=120, input_text=None):
            self.sh.calls.append((list(cmd), timeout, input_text))
            unit = cmd[-1]
            if "is-enabled" in cmd or "is-active" in cmd:
                return (1, "") if unit in missing else (0, "")
            if cmd[:3] == ["sudo", "systemctl", "start"]:
                return 0, ""
            raise AssertionError("unexpected: %r" % (cmd,))
        return fake

    def test_status_json(self):
        stub = mock.Mock()
        stub.collect.return_value = {
            n: {"browser": "healthy", "browser_detail": "silent",
                "cdp": "healthy", "cdp_detail": "silent"}
            for n in super_cli.VALID_NODES}
        with mock.patch.object(super_cli, "_sh", self._unit_states(
                missing=("chromebox-watchdog-def.timer",))):
            with mock.patch.dict(sys.modules, {"host_evidence": stub}):
                buf = io.StringIO()
                with redirect_stdout(buf):
                    super_cli.cmd_watchdog_status(_ns(json=True))
        data = json.loads(buf.getvalue())
        self.assertTrue(data["ok"])
        self.assertEqual(len(data["nodes"]), 6)
        by_node = {n["node"]: n for n in data["nodes"]}
        self.assertFalse(by_node["def"]["enabled"])
        self.assertFalse(by_node["def"]["active"])
        self.assertTrue(by_node["dev"]["active"])
        self.assertEqual(by_node["dev"]["browser"], "healthy")
        self.assertIn("timer", data["relay"])

    def test_status_survives_missing_evidence(self):
        with mock.patch.object(super_cli, "_sh", self._unit_states()):
            with mock.patch.dict(sys.modules, {"host_evidence": None}):
                buf = io.StringIO()
                with redirect_stdout(buf):
                    super_cli.cmd_watchdog_status(_ns(json=True))
        data = json.loads(buf.getvalue())
        self.assertTrue(data["ok"])
        self.assertEqual(data["nodes"][0]["browser"], "unknown")

    def test_run_node_and_relay(self):
        for target, unit in (("dev", "chromebox-watchdog@dev.service"),
                             ("relay", "cdp-relay-watchdog.service")):
            with self.subTest(target=target):
                with mock.patch.object(super_cli, "_sh",
                                       self._unit_states()) as _:
                    buf = io.StringIO()
                    with redirect_stdout(buf):
                        with self.assertRaises(SystemExit) as cm:
                            super_cli.cmd_watchdog_run(
                                _ns(target=target, json=True))
                    self.assertEqual(cm.exception.code, 0)
                    data = json.loads(buf.getvalue())
                    self.assertTrue(data["ok"])
                    self.assertEqual(data["unit"], unit)

    def test_run_rejects_bad_target(self):
        with self.assertRaises(SystemExit) as cm:
            super_cli.cmd_watchdog_run(_ns(target="ghost"))
        self.assertEqual(cm.exception.code, 1)

    def test_run_reports_start_failure(self):
        def fake(cmd, timeout=120, input_text=None):
            return 1, "Failed to start"

        with mock.patch.object(super_cli, "_sh", fake):
            buf = io.StringIO()
            with redirect_stdout(buf):
                with self.assertRaises(SystemExit) as cm:
                    super_cli.cmd_watchdog_run(_ns(target="dev", json=True))
            self.assertEqual(cm.exception.code, 1)
            self.assertFalse(json.loads(buf.getvalue())["ok"])


class ParserTests(unittest.TestCase):
    def test_fleet_heal_parses(self):
        args = super_cli.build_parser().parse_args(["fleet", "heal", "dev"])
        self.assertEqual((args.domain, args.action, args.node),
                         ("fleet", "heal", "dev"))

    def test_watchdog_parses(self):
        args = super_cli.build_parser().parse_args(["watchdog", "run", "relay"])
        self.assertEqual((args.domain, args.action, args.target),
                         ("watchdog", "run", "relay"))
        args = super_cli.build_parser().parse_args(["watchdog"])
        self.assertEqual((args.domain, args.action), ("watchdog", "status"))


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