pcb-afm-simulation-system/scripts/test_xlsx_export.py

"""Tests for the plan-results Excel export (TEST-066).

Covers:
1. afmcore.xlsx_report unit level: column assembly (thermal columns appear
   only when thermal metrics are present) and the two-row-header writer.
2. API level: GET /api/plans/{id}/export-xlsx on a temp SQLite DB via
   FastAPI TestClient - verifies xlsx payload, header structure, category
   groups, error paths (404 / 400).

Run: python scripts/test_xlsx_export.py   (from repo root)
All source is ASCII only.
"""
from __future__ import annotations

import os
import sys
import tempfile
from pathlib import Path

REPO_ROOT = Path(__file__).resolve().parents[1]
BACKEND_DIR = REPO_ROOT / "web" / "backend"
SRC_DIR = REPO_ROOT / "src"
for p in (str(BACKEND_DIR), str(SRC_DIR)):
    if p not in sys.path:
        sys.path.insert(0, p)

# Isolate the DB before any app module is imported.
_tmp_db = tempfile.NamedTemporaryFile(prefix="afm_xlsx_test_", suffix=".db", delete=False)
_tmp_db.close()
os.environ["AFM_DB_PATH"] = _tmp_db.name

FAILURES = []


def check(name: str, cond: bool, detail: str = "") -> None:
    tag = "PASS" if cond else "FAIL"
    print(f"[{tag}] {name}" + (f" -- {detail}" if detail and not cond else ""))
    if not cond:
        FAILURES.append(name)


# ---------------------------------------------------------------------------
# Part 1: unit tests for afmcore.xlsx_report
# ---------------------------------------------------------------------------

from afmcore.xlsx_report import (  # noqa: E402
    CAT_EM,
    CAT_SCAN_INFO,
    CAT_THERMAL,
    build_export_columns,
    write_summary_xlsx,
)


def unit_tests() -> None:
    params = ["Airgap", "Magnet_Thickness"]

    # Case A: EM-only data (thermal_mode off) - no thermal columns.
    keys, cats, headers = build_export_columns(
        params, ["tavg_nm", "ripple_pct", "efficiency_pct", "total_losses_w"]
    )
    check("unit: fixed cols first", keys[:4] == ["run_index", "status", "solve_time_s", "error"])
    check("unit: params in scan-info group",
          keys[4:6] == params and all(c == CAT_SCAN_INFO for c in cats[4:6]))
    check("unit: no thermal columns for EM-only data", CAT_THERMAL not in cats)
    check("unit: EM group present", CAT_EM in cats)
    check("unit: tavg header", headers[keys.index("tavg_nm")] == "\u5e73\u5747\u8f6c\u77e9[Nm]")

    # Case B: thermal data present - thermal group appears after EM.
    keys, cats, headers = build_export_columns(
        params,
        ["tavg_nm", "efficiency_pct", "winding_temp_c", "magnet_temp_c",
         "temp_rise_c", "thermal_resistance_k_w"],
        {"Airgap": "\u6c14\u9699[mm]"},
    )
    check("unit: thermal group present", CAT_THERMAL in cats)
    check("unit: thermal after EM", cats.index(CAT_THERMAL) > cats.index(CAT_EM))
    check("unit: param label used", headers[keys.index("Airgap")] == "\u6c14\u9699[mm]")
    check("unit: winding temp header",
          headers[keys.index("winding_temp_c")] == "\u7ed5\u7ec4\u6e29\u5ea6[C]")

    # Case C: unknown metric key still exported (no data loss).
    keys, cats, headers = build_export_columns([], ["tavg_nm", "mystery_metric"])
    check("unit: unknown metric kept", "mystery_metric" in keys)

    # Writer round-trip: write 2 data rows, read back structure.
    out = Path(tempfile.mkdtemp()) / "roundtrip.xlsx"
    rows = [
        {"run_index": 1, "status": "OK", "solve_time_s": 170.3, "error": "",
         "Airgap": 1.2, "tavg_nm": 0.545, "winding_temp_c": 54.86},
        {"run_index": 2, "status": "FAILED", "solve_time_s": 0, "error": "boom",
         "Airgap": 1.4, "tavg_nm": "", "winding_temp_c": ""},
    ]
    write_summary_xlsx(out, keys, cats, headers, rows)

    import openpyxl
    wb = openpyxl.load_workbook(out)
    ws = wb.active
    check("unit: sheet name", ws.title == "scan_results")
    check("unit: row1 is category groups", ws.cell(1, 1).value == CAT_SCAN_INFO)
    check("unit: row2 is headers", ws.cell(2, 1).value == "\u5e8f\u53f7")
    check("unit: data starts row 3", ws.cell(3, 1).value == 1)
    check("unit: failed row error cell", ws.cell(4, 4).value == "boom")
    check("unit: freeze panes", ws.freeze_panes == "B3")
    check("unit: autofilter", ws.auto_filter.ref == f"A2:{openpyxl.utils.get_column_letter(len(keys))}4")
    check("unit: merged category groups", len(ws.merged_cells.ranges) > 0)


# ---------------------------------------------------------------------------
# Part 2: API tests against a temp DB via TestClient
# ---------------------------------------------------------------------------

def api_tests() -> None:
    from fastapi.testclient import TestClient
    from app.main import app
    from app.database import SessionLocal, init_db
    from app.models.simulation_plan import SimulationPlan
    from app.models.simulation_result import SimulationResult

    init_db()
    client = TestClient(app)

    with SessionLocal() as db:
        from app.models.project import Project

        project = Project(name="xlsx-export-test-project", topology="SSSR")
        db.add(project)
        db.commit()
        db.refresh(project)
        project_pk = project.id

        plan = SimulationPlan(
            plan_id="SP-XLSX-TEST", name="\u70ed\u4eff\u771f\u5bfc\u51fa\u9a8c\u8bc1",
            status="completed", project_id=project_pk,
        )
        plan.set_plan_dict({
            "variables": [
                {"name": "Airgap", "name_cn": "\u6c14\u9699", "unit": "mm",
                 "values": [1.0, 1.2]},
            ],
        })
        db.add(plan)
        db.commit()
        db.refresh(plan)
        plan_pk = plan.id

        rows = [
            (1, "OK", 178.0, "", {"Airgap": 1.0},
             {"tavg_nm": 0.522, "ripple_pct": 2.82, "efficiency_pct": 86.06,
              "total_losses_w": 41.9, "winding_temp_c": 52.59,
              "magnet_temp_c": 67.61, "temp_rise_c": 27.57,
              "thermal_resistance_k_w": 5.657}),
            (2, "OK", 181.2, "", {"Airgap": 1.2},
             {"tavg_nm": 0.545, "ripple_pct": 2.99, "efficiency_pct": 84.91,
              "total_losses_w": 47.1, "winding_temp_c": 54.86,
              "magnet_temp_c": 62.74, "temp_rise_c": 29.91,
              "thermal_resistance_k_w": 2.058}),
        ]
        for idx, status, secs, err, params, metrics in rows:
            r = SimulationResult(
                plan_id=plan_pk, run_index=idx, status=status,
                solve_time_s=secs, error_message=err,
            )
            r.set_params(params)
            r.set_metrics(metrics)
            db.add(r)
        # A second plan with EM-only results (thermal_mode off).
        plan2 = SimulationPlan(plan_id="SP-XLSX-EM", name="\u7eaf\u7535\u626b\u63cf", status="completed", project_id=project_pk)
        plan2.set_plan_dict({"variables": []})
        db.add(plan2)
        db.commit()
        db.refresh(plan2)
        r2 = SimulationResult(plan_id=plan2.id, run_index=1, status="OK",
                              solve_time_s=170.0)
        r2.set_params({})
        r2.set_metrics({"tavg_nm": 0.5, "efficiency_pct": 86.0})
        db.add(r2)
        db.commit()
        plan2_pk = plan2.id

    # 1) Happy path: thermal plan export.
    resp = client.get(f"/api/plans/{plan_pk}/export-xlsx")
    check("api: export 200", resp.status_code == 200, f"got {resp.status_code}")
    check("api: xlsx content-type",
          "spreadsheetml" in resp.headers.get("content-type", ""))
    disp = resp.headers.get("content-disposition", "")
    check("api: attachment disposition", "attachment" in disp and "filename*=" in disp)

    out = Path(tempfile.mkdtemp()) / "api_export.xlsx"
    out.write_bytes(resp.content)
    import openpyxl
    wb = openpyxl.load_workbook(out)
    ws = wb.active
    headers = [ws.cell(2, c).value for c in range(1, ws.max_column + 1)]
    cats = [ws.cell(1, c).value for c in range(1, ws.max_column + 1)]
    check("api: thermal headers present",
          "\u7ed5\u7ec4\u6e29\u5ea6[C]" in headers and "\u70ed\u963b[K/W]" in headers)
    check("api: thermal category present", CAT_THERMAL in cats)
    check("api: param header uses Chinese label", "\u6c14\u9699[mm]" in headers)
    check("api: two data rows", ws.max_row == 4)
    check("api: row3 winding temp", abs(ws.cell(3, headers.index("\u7ed5\u7ec4\u6e29\u5ea6[C]") + 1).value - 52.59) < 1e-6)

    # 2) EM-only plan: no thermal columns.
    resp2 = client.get(f"/api/plans/{plan2_pk}/export-xlsx")
    check("api: EM-only export 200", resp2.status_code == 200)
    out2 = Path(tempfile.mkdtemp()) / "api_export_em.xlsx"
    out2.write_bytes(resp2.content)
    ws2 = openpyxl.load_workbook(out2).active
    headers2 = [ws2.cell(2, c).value for c in range(1, ws2.max_column + 1)]
    check("api: EM-only has no thermal columns",
          "\u7ed5\u7ec4\u6e29\u5ea6[C]" not in headers2 and CAT_THERMAL not in
          [ws2.cell(1, c).value for c in range(1, ws2.max_column + 1)])

    # 3) Error paths.
    check("api: 404 for unknown plan", client.get("/api/plans/99999/export-xlsx").status_code == 404)
    with SessionLocal() as db:
        empty = SimulationPlan(plan_id="SP-XLSX-EMPTY", name="\u7a7a\u65b9\u6848", status="draft", project_id=project_pk)
        empty.set_plan_dict({"variables": []})
        db.add(empty)
        db.commit()
        db.refresh(empty)
        empty_pk = empty.id
    check("api: 400 when no results",
          client.get(f"/api/plans/{empty_pk}/export-xlsx").status_code == 400)


if __name__ == "__main__":
    unit_tests()
    api_tests()
    print()
    if FAILURES:
        print(f"FAILED: {len(FAILURES)} -> {FAILURES}")
        sys.exit(1)
    print("ALL PASS")