6krrt/tests/test_admin_triggers.py

"""Tests for the /admin/api operational trigger endpoints.

refresh-catalog, seed-energy, and apply-feedback run repo maintenance scripts
via ``asyncio.create_subprocess_exec``; restart-service schedules systemctl
through FastAPI BackgroundTasks. These tests monkeypatch
``asyncio.create_subprocess_exec`` (and ``subprocess.run``) so no real process
or network is ever touched — they assert the endpoint maps query params and
CLI args onto the spawned command and returns the documented job shape.
"""

from __future__ import annotations

import sqlite3
import sys
from pathlib import Path

import pytest
from starlette.testclient import TestClient

import admin
import dispatcher
from config import load_config

ROOT = Path(__file__).resolve().parent.parent
CFG = load_config(str(ROOT / "config" / "config.yaml"))


class FakeProcess:
    """A stand-in for ``asyncio.subprocess.Process`` with a known outcome."""

    def __init__(self, stdout: bytes = b"ok\n", returncode: int = 0):
        self._stdout = stdout
        self.returncode = returncode
        self.killed = False

    async def communicate(self):
        return self._stdout, None

    async def wait(self):
        return self.returncode

    def kill(self):
        self.killed = True


class _Recorder:
    """Captures every ``create_subprocess_exec`` call as ``(args, kwargs)``."""

    def __init__(self):
        self.calls: list[tuple[tuple, dict]] = []


@pytest.fixture
def fake_spawn(monkeypatch):
    """Replace ``asyncio.create_subprocess_exec`` with a success-bound fake."""
    recorder = _Recorder()

    async def _fake(*args, **kwargs):
        recorder.calls.append((list(args), kwargs))
        return FakeProcess(stdout=b"ok\n", returncode=0)

    monkeypatch.setattr(admin.asyncio, "create_subprocess_exec", _fake)
    return recorder


@pytest.fixture
def failing_spawn(monkeypatch):
    """A fake that makes any spawned command fail with returncode 3."""
    recorder = _Recorder()

    async def _fake(*args, **kwargs):
        recorder.calls.append((list(args), kwargs))
        return FakeProcess(stdout=b"boom\n", returncode=3)

    monkeypatch.setattr(admin.asyncio, "create_subprocess_exec", _fake)
    return recorder


@pytest.fixture
def seeded_client(tmp_path, monkeypatch):
    """A TestClient wired to dispatcher.app with a temp DB (mounts /admin)."""
    conn = sqlite3.connect(str(tmp_path / "test.db"))
    conn.executescript((ROOT / "config" / "schema.sql").read_text())
    conn.close()
    monkeypatch.setattr(dispatcher.cfg.database, "path", str(tmp_path / "test.db"))
    monkeypatch.setenv("NEURALWATT_API_KEY", "test-key")
    with TestClient(dispatcher.app) as client:
        yield client


# --- refresh-catalog --------------------------------------------------------


def test_refresh_catalog_runs_poller_then_tier(seeded_client, fake_spawn):
    """POST /admin/api/refresh-catalog spawns poller then tier, &&-semantics."""
    resp = seeded_client.post("/admin/api/refresh-catalog")
    assert resp.status_code == 200
    job = resp.json()
    assert job["status"] == "success"
    assert job["returncode"] == 0
    assert "poller" in job["command"]
    assert "tier" in job["command"]

    # One spawn per step, in order, with the venv python as the interpreter.
    argv = [args for args, _ in fake_spawn.calls]
    assert argv == [
        [sys.executable, "-m", "poller"],
        [sys.executable, "-m", "tier"],
    ]
    for _, kwargs in fake_spawn.calls:
        assert kwargs["cwd"] == str(ROOT)


def test_refresh_catalog_job_shape(seeded_client, fake_spawn):
    """The job object carries id/command/status/returncode/output_tail."""
    job = seeded_client.post("/admin/api/refresh-catalog").json()
    assert isinstance(job["id"], str) and job["id"]
    assert "command" in job
    assert job["status"] == "success"
    assert job["returncode"] == 0
    # Two steps each emit "ok\n", so the merged tail carries both.
    assert job["output_tail"] == "ok\nok\n"


def test_failed_command_reports_failure(seeded_client, failing_spawn):
    """A non-zero returncode is reported as failed with that returncode."""
    job = seeded_client.post("/admin/api/apply-feedback").json()
    assert job["status"] == "failed"
    assert job["returncode"] == 3
    assert job["output_tail"] == "boom\n"


# --- seed-energy ------------------------------------------------------------


def test_seed_energy_default_samples(seeded_client, fake_spawn):
    """Without ?samples=, seed-energy defaults to ``--samples 5``."""
    seeded_client.post("/admin/api/seed-energy")
    argv = fake_spawn.calls[0][0]
    assert argv == [sys.executable, "-m", "seed_energy", "--samples", "5"]


def test_seed_energy_accepts_samples_query(seeded_client, fake_spawn):
    """POST /admin/api/seed-energy?samples=3 forwards ``--samples 3``."""
    resp = seeded_client.post("/admin/api/seed-energy?samples=3")
    assert resp.status_code == 200
    assert resp.json()["status"] == "success"
    argv = fake_spawn.calls[0][0]
    assert argv == [sys.executable, "-m", "seed_energy", "--samples", "3"]


# --- apply-feedback ---------------------------------------------------------


def test_apply_feedback_apply_mode(seeded_client, fake_spawn):
    """Without dry_run, feedback.py runs with no extra flag."""
    resp = seeded_client.post("/admin/api/apply-feedback")
    assert resp.status_code == 200
    assert resp.json()["status"] == "success"
    argv = fake_spawn.calls[0][0]
    assert argv == [sys.executable, "-m", "feedback"]


def test_apply_feedback_dry_run(seeded_client, fake_spawn):
    """?dry_run=true appends ``--dry-run`` to the feedback command."""
    seeded_client.post("/admin/api/apply-feedback?dry_run=true")
    argv = fake_spawn.calls[0][0]
    assert argv == [sys.executable, "-m", "feedback", "--dry-run"]


# --- restart-service --------------------------------------------------------


def test_restart_service_returns_immediately(seeded_client, monkeypatch):
    """POST /admin/api/restart-service returns {"status": "restarting"} 200.

    The systemctl call is scheduled via BackgroundTasks, never awaited inline,
    so the response body is the immediate "restarting" status and the actual
    restart fires as a post-response background task.
    """
    calls: list[tuple] = []
    monkeypatch.setattr(
        admin.subprocess, "run", lambda *a, **k: calls.append((a, k))
    )

    resp = seeded_client.post("/admin/api/restart-service")
    assert resp.status_code == 200
    assert resp.json() == {"status": "restarting"}

    # The BackgroundTask ran after the response was produced; it must have
    # scheduled exactly the systemctl restart command.
    assert calls, "BackgroundTask never fired systemctl"
    spawned = calls[0][0][0]
    assert spawned == ["systemctl", "--user", "restart", "llm-router.service"]