PacketStation/tests/server/feedhelpers.py

"""Helpers for the feed and drawing tests: the generator, fixture dirs, feed settings and one job run.

`run_feed(settings, state)` and `run_drawings(settings)` run the real job functions (visual-feed, drawing-lookup)
once, the way the worker would, over a connection as the worker role, against the fixture files in
`settings.fixture_dir`. Never VISUAL or M-Files.
"""

from __future__ import annotations

import importlib
import json
import sys
from datetime import date, timedelta
from pathlib import Path
from typing import Any

from pshelpers import make_settings
from psycopg.types.json import Jsonb

from packetstation.settings import Settings

ROOT = Path(__file__).resolve().parents[2]
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

gen = importlib.import_module("tests.fixtures.generator")

FEED_JOB = "visual-feed"
DRAWING_JOB = "drawing-lookup"
FEED_COUNT_KEYS = (
    "wos_window",
    "wos_tracked",
    "rows",
    "materials",
    "orders",
    "labor_legs",
    "new_drafts",
    "drafts_updated",
    "cos_opened",
    "cos_reaimed",
    "cos_withdrawn",
    "retired",
    "wo_gone_flagged",
    "closed_observed",
    "ignored_golive",
    "projects_created",
    "units_created",
    "auto_released",
    "unmapped_resources",
    "skipped_removed",
    "skipped_closed",
    "suspect",
)
# The counts that are zero when VISUAL did not change between two runs.
CHANGE_KEYS = (
    "new_drafts",
    "drafts_updated",
    "cos_opened",
    "cos_reaimed",
    "cos_withdrawn",
    "retired",
    "wo_gone_flagged",
    "closed_observed",
    "projects_created",
    "units_created",
    "auto_released",
)


def fixture_dir(tmp: Path, today: date | None = None) -> Path:
    """Generate the fixtures into tmp/fixtures (state 1)."""
    out: Path = gen.generate(tmp / "fixtures", today)
    return out


def set_revisions(directory: Path, dwg: str, revisions: dict[str, Any] | None, state: int = 1) -> None:
    """Replace the `revisions` of one drawing in drawings.<state>.json (None removes them)."""
    path = Path(directory) / f"drawings.{state}.json"
    data = json.loads(path.read_text(encoding="utf-8"))
    if revisions is None:
        data[dwg].pop("revisions", None)
    else:
        data[dwg]["revisions"] = revisions
    path.write_text(json.dumps(data, indent=1), encoding="utf-8")


def switch_state(directory: Path, n: int | str) -> None:
    (Path(directory) / "state").write_text(str(n), encoding="utf-8")


def feed_settings_over(base: Settings, tmp: Path) -> Settings:
    """`base` (a pg_settings) with the fixture sources pointed at freshly generated files under tmp."""
    d = fixture_dir(tmp)
    return base.model_copy(update={"feed_source": "fixture", "drawing_source": "fixture", "fixture_dir": d})


def fixture_settings(tmp: Path, **over: Any) -> Settings:
    """Plain test settings (no database) with fixture sources."""
    return make_settings(tmp, feed_source="fixture", drawing_source="fixture", fixture_dir=tmp / "fixtures", **over)


def set_feed(settings: Settings, **values: Any) -> None:
    """Merge values into the `feed` settings row by SQL (no audit: test setup)."""
    from pshelpers import app_conn

    from packetstation.feed import config

    with app_conn(settings) as conn:
        cur = config.load(conn).model_dump(mode="json")
        cur.update({k: (v.isoformat() if isinstance(v, date) else v) for k, v in values.items()})
        conn.execute(
            "INSERT INTO settings (key, value, version, updated_at) VALUES ('feed', %s, 1, now())"
            " ON CONFLICT (key) DO UPDATE SET value = EXCLUDED.value, version = settings.version + 1",
            (Jsonb(cur),),
        )


def golive_today_offset() -> date:
    return date.today() + timedelta(days=gen.GOLIVE_OFFSET)


# ---------------------------------------------------------------- one job run


def scheduler(settings: Settings) -> Any:
    """An unstarted worker Scheduler: what runs.run_job needs, without APScheduler or the listener."""
    from packetstation.worker.scheduler import Scheduler

    return Scheduler(settings, build="pytest")


def run_job(settings: Settings, job: str, *, requested: bool = False, allow_fail: bool = False) -> dict[str, Any]:
    """Run one registered job once through runs.run_job, as the worker does. Returns the job_runs row.

    `requested` inserts a job_requests row first, so the run is a "run now". A failed run raises AssertionError
    with its errors unless allow_fail.
    """
    from packetstation import db
    from packetstation.worker import runs

    if requested:
        with db.connect_as("worker", settings, autocommit=True) as conn:
            conn.execute("INSERT INTO job_requests (job, note) VALUES (%s, 'pytest')", (job,))
    row = runs.run_job(scheduler(settings), job, "schedule")
    if not allow_fail:
        assert row["status"] != "failed", row.get("errors")
    return dict(row)


def _counts(row: dict[str, Any]) -> dict[str, Any]:
    counts = dict(row.get("counts") or {})
    counts.update(
        {"_status": row["status"], "_errors": row.get("errors") or [], "_run_id": row["id"], "_trigger": row["trigger"]}
    )
    return counts


def run_feed(
    settings: Settings, state: int | None = None, *, requested: bool = False, allow_fail: bool = False
) -> dict[str, Any]:
    """Point the fixtures at `state` (when given) and run the visual-feed job once. Returns the run's counts with
    `_status`, `_errors`, `_run_id` and `_trigger` added."""
    if state is not None:
        switch_state(settings.fixture_dir, state)
    return _counts(run_job(settings, FEED_JOB, requested=requested, allow_fail=allow_fail))


def run_drawings(settings: Settings, state: int | None = None, *, requested: bool = False) -> dict[str, Any]:
    if state is not None:
        switch_state(settings.fixture_dir, state)
    return _counts(run_job(settings, DRAWING_JOB, requested=requested))


# ---------------------------------------------------------------- reading back


def packet_by_key(conn: Any, key: str, *, live: bool = True) -> dict[str, Any] | None:
    from psycopg.rows import dict_row

    with conn.cursor(row_factory=dict_row) as cur:
        live_sql = "SELECT * FROM packets WHERE visual_key = %s AND removed_at IS NULL ORDER BY id DESC LIMIT 1"
        any_sql = "SELECT * FROM packets WHERE visual_key = %s ORDER BY id DESC LIMIT 1"
        row = cur.execute(live_sql if live else any_sql, (key,)).fetchone()
    return dict(row) if row else None


def packet_id(conn: Any, key: str) -> int:
    p = packet_by_key(conn, key)
    assert p is not None, key
    return int(p["id"])


def legs_of(conn: Any, pid: int) -> list[tuple[Any, ...]]:
    return [
        tuple(r)
        for r in conn.execute(
            "SELECT wo_base, lot, split, sub, seq, qty, target_hrs, visual_status, missing_runs, gone_at IS NOT NULL"
            " FROM packet_work_orders WHERE packet_id = %s ORDER BY wo_base, lot, split, sub, seq",
            (pid,),
        ).fetchall()
    ]


def cos_of(conn: Any, pid: int) -> list[dict[str, Any]]:
    from psycopg.rows import dict_row

    with conn.cursor(row_factory=dict_row) as cur:
        return [
            dict(r)
            for r in cur.execute("SELECT * FROM change_orders WHERE packet_id = %s ORDER BY co_no", (pid,)).fetchall()
        ]


def move_status(conn: Any, pid: int, status: str) -> None:
    """Move a packet past draft by SQL (test setup: the transitions are tested elsewhere)."""
    conn.execute(
        "UPDATE packets SET status = %s::packet_status, status_at = now(), heat_map = coalesce(heat_map, true),"
        " version = version + 1 WHERE id = %s",
        (status, pid),
    )


def system_audit(conn: Any, action: str) -> list[tuple[Any, ...]]:
    return [
        tuple(r)
        for r in conn.execute(
            "SELECT role_snapshot, user_id, subject_type, subject_id, packet_id FROM audit_events WHERE action = %s"
            " ORDER BY id",
            (action,),
        ).fetchall()
    ]


__all__ = [
    "CHANGE_KEYS",
    "FEED_COUNT_KEYS",
    "cos_of",
    "feed_settings_over",
    "fixture_dir",
    "fixture_settings",
    "gen",
    "golive_today_offset",
    "legs_of",
    "move_status",
    "packet_by_key",
    "packet_id",
    "run_drawings",
    "run_feed",
    "run_job",
    "set_feed",
    "switch_state",
    "system_audit",
]