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)

    # 4) CSV export (same columns as xlsx, single header row, utf-8-sig).
    resp3 = client.get(f"/api/plans/{plan_pk}/export-csv")
    check("csv: export 200", resp3.status_code == 200)
    check("csv: content-type", "text/csv" in resp3.headers.get("content-type", ""))
    disp3 = resp3.headers.get("content-disposition", "")
    check("csv: attachment disposition", "attachment" in disp3 and "filename*=" in disp3)
    text = resp3.content.decode("utf-8-sig")
    import csv as _csv
    import io as _io
    rows3 = list(_csv.reader(_io.StringIO(text)))
    check("csv: header + 2 data rows", len(rows3) == 3)
    check("csv: header has seq label", rows3[0][0] == "\u5e8f\u53f7")
    check("csv: header has winding temp",
          "\u7ed5\u7ec4\u6e29\u5ea6[C]" in rows3[0])
    check("csv: row1 winding temp value",
          abs(float(rows3[1][rows3[0].index("\u7ed5\u7ec4\u6e29\u5ea6[C]")]) - 52.59) < 1e-6)
    check("csv: 404 for unknown plan",
          client.get("/api/plans/99999/export-csv").status_code == 404)
    check("csv: 400 when no results",
          client.get(f"/api/plans/{empty_pk}/export-csv").status_code == 400)

    # 5) Raw full-fidelity columns: attach a raw archive to row 1 of the
    # thermal plan and verify all fields (with units) appear in the export.
    with SessionLocal() as db:
        r1 = (db.query(SimulationResult)
              .filter(SimulationResult.plan_id == plan_pk,
                      SimulationResult.run_index == 1).first())
        r1.set_raw([
            {"section": "\u9a71\u52a8", "name": "\u76f4\u6d41\u6bcd\u7ebf\u7535\u538b",
             "value": 48.0, "unit": "Volts"},
            {"section": "\u7535\u78c1", "name": "\u8f6c\u77e9\u5bc6\u5ea6",
             "value": 23.333, "unit": "kNm/m3"},
            {"section": "Thermal-\u6e29\u5ea6", "name": "T[\u7ed5\u7ec4\u5e73\u5747]",
             "value": 52.591, "unit": "\u00b0C"},
            {"section": "Thermal-\u70ed\u963b", "name": "Rt[\u6c14\u9699]",
             "value": 5.076, "unit": "\u00b0C/W"},
        ])
        db.commit()

    resp4 = client.get(f"/api/plans/{plan_pk}/export-xlsx")
    check("raw: export 200", resp4.status_code == 200)
    out4 = Path(tempfile.mkdtemp()) / "raw_export.xlsx"
    out4.write_bytes(resp4.content)
    import openpyxl as _oxl
    ws4 = _oxl.load_workbook(out4).active
    h4 = [ws4.cell(2, c).value for c in range(1, ws4.max_column + 1)]
    c4 = []
    prev = None
    for c in range(1, ws4.max_column + 1):
        v = ws4.cell(1, c).value
        if v and v != prev:
            c4.append(v); prev = v
    check("raw: EM raw header with unit",
          "\u76f4\u6d41\u6bcd\u7ebf\u7535\u538b[Volts]" in h4)
    check("raw: thermal raw header with unit",
          "T[\u7ed5\u7ec4\u5e73\u5747][\u00b0C]" in h4)
    check("raw: thermal category translated",
          "Thermal-\u6e29\u5ea6".replace("Thermal-", "\u70ed\u4eff\u771f-") in c4
          or "\u70ed\u4eff\u771f-\u6e29\u5ea6" in c4)
    r3v = [ws4.cell(3, c).value for c in range(1, ws4.max_column + 1)]
    check("raw: raw value in data row",
          52.591 in r3v and 48.0 in r3v)

    # parse_export_full unit test on the real coupled-run export files.
    from afmcore.metrics import parse_export_full
    em_csv = REPO_ROOT / "output" / "thermal_validation_20260904_124142" / "raw" / "emagnetic.csv"
    th_csv = REPO_ROOT / "output" / "thermal_validation_20260904_124142" / "raw" / "thermal_steadystate.csv"
    if em_csv.exists() and th_csv.exists():
        em_rows = parse_export_full(em_csv)
        th_rows = parse_export_full(th_csv)
        check("full: EM export parsed", len(em_rows) > 200,
              f"got {len(em_rows)}")
        check("full: thermal export parsed", len(th_rows) > 500,
              f"got {len(th_rows)}")
        check("full: units captured",
              any(r["unit"] == "\u00b0C" for r in th_rows))
        check("full: sections captured",
              any(r["section"] == "\u6e29\u5ea6" for r in th_rows))
    else:
        check("full: real export files present", False,
              "thermal_validation_20260904_124142 raw files missing")

    # INF / NaN safety: non-finite floats must not enter raw rows as floats
    # (they break JSON reporting: "Out of range float values are not JSON
    # compliant" - observed live on 2026-09-04 task e53ff277).
    import math as _math
    tmp_csv = Path(tempfile.mkdtemp()) / "inf_test.csv"
    tmp_csv.write_text(
        "\u7535\u78c1\n"
        "\u6052\u8f6c\u77e9\u8f6c\u901f\u9650\u5236;INF;rpm\n"
        "\u5e73\u5747\u8f6c\u77e9;0.5;Nm\n"
        "Node to Node Thermal Resistances\n"
        "0;Ambient;2\n"
        "\u6e29\u5ea6\n"
        "T[\u7ed5\u7ec4];52.5;\u00b0C\n",
        encoding="utf-8",
    )
    full_rows = parse_export_full(tmp_csv)
    inf_row = [r for r in full_rows if r["name"] == "\u6052\u8f6c\u77e9\u8f6c\u901f\u9650\u5236"]
    check("full: INF kept as string",
          len(inf_row) == 1 and inf_row[0]["value"] == "INF")
    check("full: no non-finite floats",
          all(not (isinstance(r["value"], float) and not _math.isfinite(r["value"]))
              for r in full_rows))
    check("full: node-to-node matrix skipped",
          all(r["section"] != "Node to Node Thermal Resistances" for r in full_rows))
    import json as _json
    _json.dumps(full_rows)  # must not raise
    check("full: raw rows JSON-serializable", True)

    # Normalized parser must also skip INF (not extract it as a metric).
    from afmcore.metrics import parse_export as _pe
    parsed = _pe(tmp_csv)
    flat_vals = [v for sec in parsed.values() for v in sec.values()]
    check("full: parse_export skips INF",
          all(_math.isfinite(v) for v in flat_vals))

    # Node Temperatures 4-column layout: "node_id;name;value;unit" must be
    # remapped to (name, value, unit) - regression for 2026-09-04 user report
    # where headers became "0[25]" and data cells held part names.
    nt_csv = Path(tempfile.mkdtemp()) / "node_temps.csv"
    nt_csv.write_text(
        "Node Temperatures\n"
        "0;\"Ambient\";25;\u00b0C\n"
        "1;\"Housing\";39.99;\u00b0C\n",
        encoding="utf-8",
    )
    nt_rows = parse_export_full(nt_csv)
    check("full: node temps remapped",
          nt_rows[0]["name"] == "Ambient"
          and nt_rows[0]["value"] == 25.0
          and nt_rows[0]["unit"] == "\u00b0C")
    check("full: node temps no id leak",
          all(r["name"] != "0" for r in nt_rows))


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