FabledSteward/tests/integration/test_host_maintenance_job.py

"""Integration (#5767): Steward hands a host a maintenance job and follows it
to the end, through a restart of its own, from the agent's posts alone.

Rows are seeded and read with plain SQL: importing the plugin's models again
after the app loaded the plugin raises "Table already defined" (see
test_host_agent_alerts)."""
from __future__ import annotations

import asyncio
import hashlib
import os
import uuid
from datetime import datetime, timedelta, timezone

import pytest

pytestmark = pytest.mark.integration


def _app():
    if not os.environ.get("STEWARD_DATABASE_URL"):
        pytest.skip("needs Postgres")
    from steward.app import create_app
    return create_app(testing=False)


async def _seed(app):
    from sqlalchemy import text
    from steward.models.hosts import Host
    from steward.models.users import User
    tag = uuid.uuid4().hex[:8]
    token = f"tok-{tag}-{uuid.uuid4().hex}"
    job_id = str(uuid.uuid4())
    async with app.db_sessionmaker() as s:
        async with s.begin():
            user = User(username=f"mj-{tag}", email=f"mj-{tag}@test", password_hash="x")
            host = Host(name=f"maint-{tag}", address="10.9.9.7")
            s.add_all([user, host])
            await s.flush()
            await s.execute(text(
                "INSERT INTO host_agent_registrations "
                "(id, host_id, token_hash, token_created_at, last_seen_at, maint_helper, "
                " created_at, updated_at) "
                "VALUES (:id, :host, :hash, now(), now(), true, now(), now())"),
                {"id": str(uuid.uuid4()), "host": host.id,
                 "hash": hashlib.sha256(token.encode()).hexdigest()})
            await s.execute(text(
                "INSERT INTO host_maintenance_jobs "
                "(id, host_id, state, state_at, reboot_policy, step_seen, created_by, created_at) "
                "VALUES (:id, :host, 'queued', now(), 'if_required', false, :user, now())"),
                {"id": job_id, "host": host.id, "user": user.id})
            host_id = host.id
    return host_id, token, job_id


async def _post(app, token, events=(), sample=True):
    now = datetime.now(timezone.utc).isoformat()
    resp = await app.test_client().post(
        "/plugins/host_agent/ingest",
        json={"samples": [{"ts": now, "cpu_pct": 2.0}] if sample else [],
              "events": list(events)},
        headers={"Authorization": f"Bearer {token}"})
    assert resp.status_code == 200, await resp.get_data(as_text=True)
    return (await resp.get_json()).get("maintenance", [])


async def _job(app, job_id):
    from sqlalchemy import text
    async with app.db_sessionmaker() as s:
        async with s.begin():
            return (await s.execute(text(
                "SELECT j.state, j.packages, j.reboot_needed, j.back_at, w.end_at "
                "FROM host_maintenance_jobs j LEFT JOIN maintenance_windows w ON w.id = j.window_id "
                "WHERE j.id = :id"), {"id": job_id})).one()


def test_a_job_is_followed_through_a_steward_restart_and_the_host_reboot():
    app = _app()

    async def _first_half():
        host_id, token, job_id = await _seed(app)
        # The agent's next post starts the job: update first, under a window.
        assert await _post(app, token) == [{"id": f"{job_id}-update", "action": "update"}]
        state, *_, window_end = await _job(app, job_id)
        assert state == "updating" and window_end > datetime.now(timezone.utc)
        # The helper took it: the step is no longer handed out.
        running = {"kind": "maintenance", "id": f"{job_id}-update", "action": "update",
                   "state": "running"}
        assert await _post(app, token, [running]) == []
        return token, job_id

    token, job_id = asyncio.run(_first_half())

    # Steward restarts: a new process, nothing in memory but the database.
    app = _app()

    async def _second_half():
        done = {"kind": "maintenance", "id": f"{job_id}-update", "action": "update",
                "state": "done", "ok": True, "packages": 12, "reboot_required": True}
        assert await _post(app, token, [done]) == [{"id": f"{job_id}-reboot", "action": "reboot"}]
        assert (await _job(app, job_id))[:3] == ("rebooting", 12, True)
        # The agent restarting on a host that stayed up doesn't count as the reboot.
        old_boot = datetime.now(timezone.utc) - timedelta(days=3)
        await _post(app, token, [{"kind": "started", "at": datetime.now(timezone.utc).isoformat(),
                                  "boot_at": old_boot.isoformat()}])
        assert (await _job(app, job_id))[0] == "rebooting"
        # After the reboot: the helper's result and the agent's boot, in one post.
        boot = datetime.now(timezone.utc)
        rebooted = {"kind": "maintenance", "id": f"{job_id}-reboot", "action": "reboot",
                    "state": "rebooting", "ok": True}
        started = {"kind": "started", "at": (boot + timedelta(seconds=20)).isoformat(),
                   "boot_at": boot.isoformat()}
        assert await _post(app, token, [rebooted, started]) == []
        state, packages, reboot_needed, back_at, window_end = await _job(app, job_id)
        assert state == "done" and back_at is not None
        assert window_end <= datetime.now(timezone.utc)

    asyncio.run(_second_half())


def test_a_swarm_node_s_job_also_holds_the_swarm_wide_alerts():
    """Draining and rebooting a node moves the swarm figures; that is the job,
    not news (#5788). The window over them follows the job's own, and ends
    with it."""
    from types import SimpleNamespace

    from sqlalchemy import text
    from steward.core.plugin_manager import loaded_plugin_submodule
    from steward.models.users import User
    app = _app()
    maintenance = loaded_plugin_submodule("host_agent", "maintenance")
    tag = uuid.uuid4().hex[:8]
    host_name = f"swarm-node-{tag}"

    async def _go():
        now = datetime.now(timezone.utc)
        async with app.db_sessionmaker() as s:
            async with s.begin():
                user = User(username=f"sw-{tag}", email=f"sw-{tag}@test", password_hash="x")
                s.add(user)
                await s.flush()
                job = SimpleNamespace(window_id=None, created_by=user.id, state="draining",
                                      node_id="node-1", created_at=now - timedelta(seconds=1))
                await maintenance._window(s, job, host_name, now)
                await maintenance._window(s, job, host_name, now + timedelta(minutes=1))   # the next tick
                job.state = "done"
                await maintenance._window(s, job, host_name, now + timedelta(minutes=2))   # finished
            return (await s.execute(text(
                "SELECT scope_source, scope_resource, end_at FROM maintenance_windows "
                "WHERE name LIKE :n ORDER BY scope_source"), {"n": f"Maintenance: {host_name}%"})).all(), now

    rows, now = asyncio.run(_go())
    assert [(r[0], r[1]) for r in rows] == [("docker", "swarm"), ("host_agent", host_name)]
    assert all(r[2] == now + timedelta(minutes=2) for r in rows)     # ended with the job, not before