FabledSteward/tests/integration/test_host_maintenance_cadence.py

"""Integration (#5768): a plan updates its group's hosts one at a time; the
second starts only once the first is back, and a failure stops the run and
flags the host for an alert rule.

The plugin's cadence and models are reached through the module the app loaded
(see test_docker_disk_page): importing the plugin's models again raises
"Table already defined"."""
from __future__ import annotations

import asyncio
import hashlib
import importlib
import os
import uuid
from datetime import datetime, 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)


def _ha(sub):
    from steward.core.plugin_manager import get_loaded_plugin_modules
    loaded = get_loaded_plugin_modules().get("host_agent")
    base = loaded.__name__ if loaded is not None else "plugins.host_agent"
    return importlib.import_module(f"{base}.{sub}")


async def _seed(app):
    from sqlalchemy import text
    from steward.models.ansible_inventory import AnsibleGroup, AnsibleTarget
    from steward.models.hosts import Host
    from steward.models.users import User
    tag = uuid.uuid4().hex[:8]
    out = {"tokens": {}, "hosts": {}}
    async with app.db_sessionmaker() as s:
        async with s.begin():
            user = User(username=f"cad-{tag}", email=f"cad-{tag}@test", password_hash="x")
            group = AnsibleGroup(name=f"cad-{tag}")
            s.add_all([user, group])
            await s.flush()
            for n in ("a", "b"):
                host = Host(name=f"cad-{n}-{tag}", address="10.9.9.6")
                s.add(host)
                await s.flush()
                target = AnsibleTarget(name=f"cad-{n}-{tag}", address="10.9.9.6", host_id=host.id)
                target.groups = [group]
                s.add(target)
                token = f"tok-{n}-{tag}-{uuid.uuid4().hex}"
                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, :h, :hash, now(), now(), true, now(), now())"),
                    {"id": str(uuid.uuid4()), "h": host.id,
                     "hash": hashlib.sha256(token.encode()).hexdigest()})
                out["tokens"][n], out["hosts"][n] = token, host
            plan_id = str(uuid.uuid4())
            await s.execute(text(
                "INSERT INTO host_maintenance_plans (id, name, group_id, weekday, at_time, time_zone, "
                " window_minutes, reboot_policy, enabled, created_by, created_at) "
                "VALUES (:id, :name, :g, 6, '03:00', 'UTC', 240, 'if_required', true, :u, now())"),
                {"id": plan_id, "name": f"plan-{tag}", "g": group.id, "u": user.id})
    out["plan_id"] = plan_id
    return out


async def _post(app, token, events=()):
    resp = await app.test_client().post(
        "/plugins/host_agent/ingest",
        json={"samples": [{"ts": datetime.now(timezone.utc).isoformat(), "cpu_pct": 1.0}],
              "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 _tick(app):
    cadence = _ha("cadence")
    async with app.db_sessionmaker() as s:
        async with s.begin():
            await cadence.tick(s, datetime.now(timezone.utc))
            return await cadence.failed_hosts(s, datetime.now(timezone.utc))


async def _run(app, plan_id):
    from sqlalchemy import text
    async with app.db_sessionmaker() as s:
        return (await s.execute(text(
            "SELECT state, position, detail FROM host_maintenance_runs WHERE plan_id = :p"),
            {"p": plan_id})).one()


def test_hosts_go_one_at_a_time_and_a_failure_stops_the_run():
    app = _app()

    async def _go():
        seed = await _seed(app)
        cadence, models = _ha("cadence"), _ha("models")
        async with app.db_sessionmaker() as s:
            async with s.begin():
                plan = await s.get(models.HostMaintenancePlan, seed["plan_id"])
                await cadence.start_run(s, plan, datetime.now(timezone.utc), manual=True)
        await _tick(app)

        # Only the first host (by name) has a job; the second waits.
        ta, tb = seed["tokens"]["a"], seed["tokens"]["b"]
        assert await _post(app, tb) == []
        (step,) = await _post(app, ta)
        assert step["action"] == "update"
        done = {"kind": "maintenance", "id": step["id"], "action": "update", "state": "done",
                "ok": True, "packages": 2, "reboot_required": False}
        await _post(app, ta, [done])

        # Back in service (posting, not a swarm node): the next host starts.
        await _tick(app)
        assert (await _run(app, seed["plan_id"]))[:2] == ("running", 1)
        (step_b,) = await _post(app, tb)
        failed = {"kind": "maintenance", "id": step_b["id"], "action": "update",
                  "state": "failed", "ok": False, "detail": "dpkg was interrupted"}
        await _post(app, tb, [failed])

        flags = await _tick(app)
        state, _, detail = await _run(app, seed["plan_id"])
        assert state == "failed" and seed["hosts"]["b"].name in detail and "dpkg" in detail
        assert flags[seed["hosts"]["b"].name] == 1.0 and flags[seed["hosts"]["a"].name] == 0.0

    asyncio.run(_go())