PacketStation/tests/server/pkthelpers.py

"""Helpers for the packet API tests: projects and units by SQL, hand drafts, PDFs, transitions, the happy path."""

from __future__ import annotations

import io
import secrets
from collections.abc import Iterator
from contextlib import contextmanager
from typing import Any

from fakedir import FAKE_PASSWORD, FakeDirectory
from fastapi.testclient import TestClient
from pshelpers import ORIGIN_HEADERS, app_conn, sign_in
from psycopg.types.json import Jsonb

from packetstation.domain import checklists
from packetstation.settings import Settings

SIGN = {"factor": "password", "secret": FAKE_PASSWORD}
USERS = ("admin.ann", "clerk.clara", "wh.wade", "qc.quinn", "pm.paula", "weld.andy", "lead.lou", "sup.wes", "dup.dave")


@contextmanager
def phase2_client(settings: Settings, fake_dir: FakeDirectory) -> Iterator[TestClient]:
    from packetstation.main import create_app

    app = create_app(settings, directory=fake_dir)  # main.py registers every phase 2 router
    with TestClient(app, headers=ORIGIN_HEADERS) as client:
        yield client


def seed_users(client: TestClient, users: tuple[str, ...] = USERS) -> None:
    """Sign each test user in once so their users rows exist."""
    for sam in users:
        assert sign_in(client, sam).status_code == 200, sam


def as_user(client: TestClient, sam: str) -> None:
    r = sign_in(client, sam)
    assert r.status_code == 200, r.text


def user_id(conn: Any, sam: str) -> str:
    row = conn.execute("SELECT id FROM users WHERE sam_account = %s", (sam,)).fetchone()
    assert row is not None, sam
    return str(row[0])


def station_id(conn: Any, dept: str = "03", name: str = "Stall 1") -> int:
    row = conn.execute("SELECT id FROM stations WHERE dept_code = %s AND name = %s", (dept, name)).fetchone()
    assert row is not None
    return int(row[0])


def make_unit(conn: Any, *, number: str | None = None, lot: str = "001", status: str = "open") -> tuple[int, int]:
    """A project and one unit by SQL. Returns (project_id, unit_id)."""
    number = number or "L9" + secrets.token_hex(3).upper()
    pid = conn.execute(
        "INSERT INTO projects (number, name, customer, po) VALUES (%s, %s, %s, %s) RETURNING id",
        (number, "Test project", "Acme Refining", "PO-77"),
    ).fetchone()[0]
    uid = conn.execute(
        "INSERT INTO units (project_id, lot, name, status) VALUES (%s, %s, %s, %s::project_status) RETURNING id",
        (pid, lot, "Unit one", status),
    ).fetchone()[0]
    return int(pid), int(uid)


def add_unit(conn: Any, project_id: int, lot: str, status: str = "open") -> int:
    return int(
        conn.execute(
            "INSERT INTO units (project_id, lot, status) VALUES (%s, %s, %s::project_status) RETURNING id",
            (project_id, lot, status),
        ).fetchone()[0]
    )


def draft_body(unit_id: int, **over: Any) -> dict[str, Any]:
    body: dict[str, Any] = {
        "unit_id": unit_id,
        "dept_code": "03",
        "op": "WELD-PART",
        "title": "Spool 4 inlet",
        "customer": "Acme Refining",
        "drawing": "D-1001",
        "rev": "B",
        "wps": "WPS-11",
        "work_orders": [
            {"wo_base": "L670001", "lot": "001", "split": "0", "sub": "0", "qty": 2, "target_hrs": 1.5},
        ],
        "bom": [],
    }
    body.update(over)
    return body


def make_draft(client: TestClient, unit_id: int, **over: Any) -> dict[str, Any]:
    """A hand draft through the API as whoever is signed in (Clerical). Returns the read model."""
    r = client.post("/api/packets", json=draft_body(unit_id, **over))
    assert r.status_code == 201, r.text
    out: dict[str, Any] = r.json()
    return out


def tiny_pdf(pages: int = 1) -> bytes:
    from pypdf import PdfWriter

    w = PdfWriter()
    for _ in range(pages):
        w.add_blank_page(width=612, height=792)
    buf = io.BytesIO()
    w.write(buf)
    return buf.getvalue()


def upload(client: TestClient, data: bytes, name: str = "drawing.pdf", mime: str = "application/pdf") -> Any:
    return client.post("/api/assets", files={"file": (name, data, mime)})


def attach_drawing(client: TestClient, packet_id: int, pages: int = 1) -> dict[str, Any]:
    a = upload(client, tiny_pdf(pages))
    assert a.status_code == 201, a.text
    r = client.post(
        f"/api/packets/{packet_id}/pages", json={"asset_id": a.json()["id"], "type": "drawing", "sign": SIGN}
    )
    assert r.status_code == 200, r.text
    out: dict[str, Any] = r.json()
    return out


def set_header(client: TestClient, packet_id: int, **fields: Any) -> Any:
    return client.patch(f"/api/packets/{packet_id}", json={"set": fields, "base": {}})


def ready_draft(client: TestClient, unit_id: int, **over: Any) -> dict[str, Any]:
    """A hand draft with the heat map answered and a drawing attached: it can leave Clerical."""
    pkt = make_draft(client, unit_id, **over)
    assert set_header(client, pkt["id"], heat_map=True).status_code == 200
    return attach_drawing(client, pkt["id"])


def step(client: TestClient, pkt: dict[str, Any], action: str, **fields: Any) -> Any:
    """POST a transition with the expect taken from the packet answer."""
    body = {"action": action, "expect": {"status": pkt["status"], "dept": pkt["dept"]["code"]}, **fields}
    return client.post(f"/api/packets/{pkt['id']}/transitions", json=body)


def ok_step(client: TestClient, pkt: dict[str, Any], action: str, **fields: Any) -> dict[str, Any]:
    r = step(client, pkt, action, **fields)
    assert r.status_code == 200, r.text
    out: dict[str, Any] = r.json()
    return out


def get(client: TestClient, packet_id: int) -> dict[str, Any]:
    r = client.get(f"/api/packets/{packet_id}")
    assert r.status_code == 200, r.text
    out: dict[str, Any] = r.json()
    return out


def finish_stage_by_sql(settings: Settings, packet_id: int) -> None:
    """The checklist, filler and one weld, which the child endpoints (built later) would write."""
    keys = checklists.checklist_keys(True)
    with app_conn(settings) as conn:
        conn.execute(
            "UPDATE packet_stages SET checklist = %s, filler = 'ER70S-6' WHERE packet_id = %s AND closed_at IS NULL",
            (Jsonb({k: {"done": True, "stamp": "A7"} for k in keys}), packet_id),
        )
        conn.execute(
            "INSERT INTO welds (packet_id, weld_no, items, type, size) VALUES (%s, 1, '{1}', 'Fillet', '1/4')",
            (packet_id,),
        )


def finish_qc_by_sql(settings: Settings, packet_id: int) -> None:
    with app_conn(settings) as conn:
        conn.execute(
            "UPDATE qc_inspections SET checks = %s, weld_exam = %s, inspector = 'Quinn QC'"
            " WHERE packet_id = %s AND closed_at IS NULL",
            (Jsonb({k: {"done": True} for k in checklists.qc_keys()}), Jsonb({"1": {"result": "accept"}}), packet_id),
        )


def happy_path(client: TestClient, settings: Settings, unit_id: int) -> dict[str, Any]:
    """Draft -> warehouse -> ready -> assigned -> in progress -> review -> Q.C. -> complete (data book) -> filed."""
    with app_conn(settings) as conn:
        stall = station_id(conn)
        andy = user_id(conn, "weld.andy")
    as_user(client, "clerk.clara")
    pkt = ready_draft(client, unit_id)
    pkt = ok_step(client, pkt, "send_to_warehouse")
    assert pkt["status"] == "warehouse" and pkt["wh"]["stage"] == "incoming"

    as_user(client, "wh.wade")
    pkt = ok_step(client, pkt, "wh_release")
    assert pkt["status"] == "ready" and pkt["stage"]["dept_code"] == "03" and pkt["wh"] is None

    as_user(client, "lead.lou")
    r = step(client, pkt, "return_to_clerical", reason="Wrong department for this one")
    assert r.status_code == 403 and r.json()["code"] == "wrong_status", r.text

    as_user(client, "sup.wes")
    pkt = ok_step(client, pkt, "assign", station_id=stall, welder_id=andy)
    assert pkt["status"] == "assigned" and pkt["station"]["name"] == "Stall 1"
    assert pkt["welder"]["id"] == andy and pkt["stage"]["welder"]["id"] == andy

    as_user(client, "dup.dave")
    r = step(client, pkt, "start")
    assert r.status_code == 403 and r.json()["code"] == "forbidden", r.text

    as_user(client, "weld.andy")
    pkt = ok_step(client, pkt, "start")
    assert pkt["status"] == "in_progress" and pkt["stage"]["clock"]["running_since"] is not None
    r = step(client, pkt, "submit", sign=SIGN)
    assert r.status_code == 403 and r.json()["code"] == "precondition", r.text
    finish_stage_by_sql(settings, pkt["id"])
    r = step(client, pkt, "submit")
    assert r.status_code == 400 and r.json()["code"] == "signature_required", r.text
    pkt = ok_step(client, pkt, "submit", sign=SIGN)
    assert pkt["status"] == "review" and pkt["stage"]["clock"]["running_since"] is None

    as_user(client, "sup.wes")
    r = step(client, pkt, "send_to_qc", book=True, spool=True, sign=SIGN)
    assert r.status_code == 403 and r.json()["code"] == "precondition", r.text
    pkt = ok_step(client, pkt, "send_to_qc", book=True, spool=True, location="Bay 3 rack B", sign=SIGN)
    assert pkt["status"] == "qc" and pkt["qc"]["location"] == "Bay 3 rack B" and pkt["stage"] is not None

    as_user(client, "qc.quinn")
    r = step(client, pkt, "qc_accept", routed_to="databook", sign=SIGN)
    assert r.status_code == 403 and r.json()["code"] == "precondition", r.text
    assert "Q.C. still needs" in r.json()["detail"]
    finish_qc_by_sql(settings, pkt["id"])
    pkt = ok_step(client, pkt, "qc_accept", routed_to="databook", sign=SIGN)
    assert pkt["status"] == "complete" and pkt["routed_to"] == "databook" and pkt["qc"] is None
    assert pkt["stage"] is None and len(pkt["stages"]) == 1 and pkt["qc_history"][0]["result"] == "accepted"

    as_user(client, "clerk.clara")
    r = step(client, pkt, "file", sign=SIGN)
    assert r.status_code == 403 and r.json()["code"] == "precondition"
    assert r.json()["detail"] == "Close the job in VISUAL first."
    r = client.post(f"/api/packets/{pkt['id']}/vclosed", json={"closed": True})
    assert r.status_code == 200 and r.json()["vclosed"]["by"]["display_name"] == "Clara Clerk", r.text
    pkt = ok_step(client, r.json(), "file", sign=SIGN)
    assert pkt["status"] == "filed" and pkt["seal"]["broken_at"] is None and len(pkt["seal"]["key_id"]) == 8
    return pkt