UmbraTrans/umbratrans_web/tests/test_node_agent.py

"""The node agent, driven as the process it actually is.

Every test here starts a real ``node_agent.py`` subprocess and points it at a stub
controller over real HTTP. That is heavier than calling the agent's methods, and
the weight is the point: the two behaviours worth testing here - stopping without
a signal, and dying without releasing anything - are about *processes*. ``os._exit``
cannot be observed in-process at all, and a stop file has no meaning without
something polling for it in a loop. A test that called ``NodeAgent.run`` in-process
would have had to mock both away and would have proved only that a mock was called.
"""

from __future__ import annotations

import json
import os
import subprocess
import sys
import threading
import time
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path
from typing import Any

import pytest

AGENT = Path(__file__).resolve().parents[1] / "node_agent.py"
NODE_KEY = "test-key-never-sent-anywhere-real"


class _StubController:
    """Just enough of ``/api/node`` to let an agent run.

    Counts what it was asked for, because that is how the tests tell a stopped
    agent from a running one and a dead one from a working one. A test that
    asserted on a log line would pass on a log line written by an agent that had
    already stopped.
    """

    def __init__(self) -> None:
        self.lock = threading.Lock()
        self.heartbeats = 0
        self.claims = 0
        self.renews = 0
        self.completes = 0
        self.keys: set[str] = set()
        self.offers_job = True
        self.server: ThreadingHTTPServer | None = None

    @property
    def url(self) -> str:
        assert self.server is not None
        return f"http://127.0.0.1:{self.server.server_address[1]}"

    def __enter__(self) -> "_StubController":
        controller = self

        class Handler(BaseHTTPRequestHandler):
            def log_message(self, *args: Any) -> None:  # noqa: A003 - stdlib name
                pass

            def _read(self) -> dict[str, Any]:
                length = int(self.headers.get("Content-Length") or 0)
                if not length:
                    return {}
                return json.loads(self.rfile.read(length) or b"{}")

            def _send(self, payload: Any) -> None:
                body = b"" if payload is None else json.dumps(payload).encode()
                self.send_response(200)
                self.send_header("Content-Type", "application/json")
                self.send_header("Content-Length", str(len(body)))
                self.end_headers()
                self.wfile.write(body)

            def do_POST(self) -> None:  # noqa: N802 - stdlib name
                self._read()
                authorization = self.headers.get("Authorization") or ""
                with controller.lock:
                    controller.keys.add(authorization)
                    if self.path.startswith("/api/node/heartbeat"):
                        controller.heartbeats += 1
                        self._send({"status": "online", "schedulable": True})
                    elif self.path.startswith("/api/node/jobs/claim"):
                        controller.claims += 1
                        if not controller.offers_job:
                            self._send(None)
                            return
                        self._send(
                            {
                                "job": {
                                    "id": "job-1",
                                    "kind": "sleep",
                                    "payload": {"seconds": 0},
                                },
                                "placement": {
                                    "node_name": "stub",
                                    "device_kind": "cpu",
                                    "reason": "test",
                                },
                                "lease_expires_at": "2026-01-01T00:00:00+00:00",
                            }
                        )
                    elif "/renew" in self.path:
                        controller.renews += 1
                        self._send({"ok": True})
                    elif self.path.endswith("/complete"):
                        controller.completes += 1
                        self._send({"ok": True})
                    else:
                        self._send({})

        self.server = ThreadingHTTPServer(("127.0.0.1", 0), Handler)
        threading.Thread(target=self.server.serve_forever, daemon=True).start()
        return self

    def __exit__(self, *_: Any) -> None:
        assert self.server is not None
        self.server.shutdown()
        self.server.server_close()

    def snapshot(self) -> tuple[int, int, int]:
        with self.lock:
            return self.heartbeats, self.claims, self.completes


def _write_config(path: Path, controller: _StubController, **extra: Any) -> Path:
    path.write_text(
        json.dumps(
            {
                "controller_url": controller.url,
                "node_name": "stub-node",
                "kinds": [],
                "lease_seconds": 30,
                "poll_interval_seconds": 0.05,
                "node_key": "",
                **extra,
            }
        )
    )
    return path


def _agent_environment(*, with_key: bool = True) -> dict[str, str]:
    """A database-free environment, because the agent never had one.

    The key is here unless a test is specifically about how it arrives, since that
    is the arrangement a node actually runs under.
    """

    environment = {k: v for k, v in os.environ.items() if not k.startswith("UMBRATRANS")}
    if with_key:
        environment["UMBRA_NODE_KEY"] = NODE_KEY
    else:
        environment.pop("UMBRA_NODE_KEY", None)
    return environment


def _run_agent(config: Path, *args: str, timeout: float = 20.0, with_key: bool = True):
    environment = _agent_environment(with_key=with_key)
    return subprocess.run(
        [sys.executable, str(AGENT), "--config", str(config), "--log-level", "INFO", *args],
        capture_output=True,
        text=True,
        env=environment,
        timeout=timeout,
    )


def test_a_stop_file_stops_an_agent_that_cannot_be_signalled(tmp_path: Path) -> None:
    """The agent stops when a file appears, with no signal involved.

    A node agent normally runs as a service account, so its supervisor is not
    allowed to signal it - kill(2) across accounts fails. Stopping it therefore has
    to work without a signal, or every stop leaks a running agent. That is not
    hypothetical: the proof of concept leaked one on every run, because the only
    thing it could signal was the ``sudo`` wrapper it had itself spawned.
    """

    with _StubController() as controller:
        stop_file = tmp_path / "node.stop"
        config = _write_config(tmp_path / "node.json", controller, stop_file=str(stop_file))

        started = subprocess.Popen(
            [sys.executable, str(AGENT), "--config", str(config), "--log-level", "INFO"],
            stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True,
            env=_agent_environment(),
        )
        try:
            deadline = time.time() + 10
            while time.time() < deadline and controller.snapshot()[0] == 0:
                time.sleep(0.05)
            before = controller.snapshot()
            assert before[0] > 0, "the agent never reached the controller, so the stop " \
                                   "would not prove anything"

            stop_file.write_text("stop\n")
            assert started.wait(timeout=15) == 0, "the agent did not stop cleanly"

            # And it really did stop: no further traffic at all after the file.
            settled = controller.snapshot()
            time.sleep(0.6)
            assert controller.snapshot() == settled, "traffic continued after the stop file"
        finally:
            if started.poll() is None:
                started.kill()
            started.wait(timeout=5)


def test_fault_injection_leaves_a_job_unclaimed_by_anything_that_finished_it(
    tmp_path: Path,
) -> None:
    """The agent vanishes holding a lease, having done nothing.

    Exit code 137 and no completion call. The absence of the completion is the
    assertion that matters: an agent that ran the job and then died would also
    exit abruptly, and the lease would already be released by the completion, so
    the recovery path would never see a stranded job at all.
    """

    with _StubController() as controller:
        config = _write_config(tmp_path / "node.json", controller)
        finished = _run_agent(
            config, "--fault-exit-after-claim", "1", "--max-iterations", "400"
        )

        assert finished.returncode == 137, (
            f"expected the fault exit (137), got {finished.returncode}: "
            f"{finished.stderr[-400:]}"
        )
        assert "without running it" in finished.stderr, finished.stderr[-400:]
        assert controller.claims >= 1, "it never claimed anything, so nothing was stranded"
        assert controller.completes == 0, "the job was completed, so no lease was stranded"


def test_the_node_key_is_never_in_the_config_file_it_is_written_beside(tmp_path: Path) -> None:
    """A config file holding a secret is a secret that ends up in a backup.

    The key arrives by environment or by stdin, never from the file, even when the
    file is the only thing that could plausibly supply it.
    """

    with _StubController() as controller:
        config = _write_config(tmp_path / "node.json", controller)
        assert NODE_KEY not in config.read_text()

        by_env = _run_agent(config, "--max-iterations", "1")
        assert by_env.returncode == 0, by_env.stderr[-400:]

        by_stdin = _run_agent_stdin(config)
        assert by_stdin.returncode == 0, by_stdin.stderr[-400:]

        # The stub recorded the header the agent actually sent.
        assert f"Bearer {NODE_KEY}" in controller.keys


def _run_agent_stdin(config: Path):
    return subprocess.run(
        [sys.executable, str(AGENT), "--config", str(config), "--log-level", "INFO",
         "--node-key-stdin", "--max-iterations", "1"],
        input=NODE_KEY + "\n",
        capture_output=True,
        text=True,
        env=_agent_environment(),
        timeout=20,
    )


def test_the_key_must_arrive_by_one_of_the_two_ways(tmp_path: Path) -> None:
    """Neither route supplied means a clear message naming the sudo cause.

    The bare message this replaces said only "no node_key". Reaching it in
    practice means sudo reset the environment, which is a fact about the host's
    sudoers policy and not about the agent, so the agent says so.
    """

    with _StubController() as controller:
        config = _write_config(tmp_path / "node.json", controller)
        finished = _run_agent(config, "--max-iterations", "1", with_key=False)
        assert finished.returncode != 0
        assert "UMBRA_NODE_KEY is not set" in finished.stderr, finished.stderr[-300:]
        assert "sudo" in finished.stderr, (
            "the failure under sudo is the environment being reset; the message should "
            f"say so, got: {finished.stderr[-300:]}"
        )


def test_an_empty_pipe_is_reported_as_an_empty_pipe(tmp_path: Path) -> None:
    """A closed stdin with nothing written is not a key that happens to be blank."""

    with _StubController() as controller:
        config = _write_config(tmp_path / "node.json", controller)
        finished = subprocess.run(
            [sys.executable, str(AGENT), "--config", str(config), "--node-key-stdin",
             "--max-iterations", "1"],
            input="", capture_output=True, text=True,
            env=_agent_environment(with_key=False), timeout=20,
        )
        assert finished.returncode != 0
        assert "stdin was empty" in finished.stderr, finished.stderr[-300:]


if __name__ == "__main__":  # pragma: no cover
    sys.exit(pytest.main([__file__, "-v"]))