shanxia2.0/backend/scripts/verify_plan_archive_zip_export.py

"""测试库多场景验证：计划归档 ZIP 导出（服务层直连 DB + MinIO）。"""
from __future__ import annotations

import io
import logging
import re
import sys
import zipfile
from datetime import datetime
from pathlib import Path

from openpyxl import load_workbook
from sqlalchemy import text as sql_text
from sqlalchemy.orm import Session

# quiet SQL
logging.getLogger("sqlalchemy.engine").setLevel(logging.WARNING)

from app.core.database import SessionLocal
from app.models.pl_models import PlPlan
from app.services.pl_plan_archive_export import build_plan_archive_export_rows
from app.services.pl_plan_list_attachments import load_plan_document_attachments_by_plan_ids
from app.services.pl_plan_archive_zip_export import (
    build_archive_zip_bytes,
    collect_hazard_cert_items,
    load_plan_products_by_plan_ids,
    merge_plan_attachment_items,
)
from app.utils.excel_export import ExportColumn, ExportValidationError

# mirror API columns subset is enough; use full archive columns via import
from app.api.pl.plan import PLAN_ARCHIVE_EXPORT_COLUMNS, _load_transports_by_id

OUT = Path("/tmp/plan-archive-zip-verify")
OUT.mkdir(parents=True, exist_ok=True)
TENANT = "shanxia2.0"
results: list[tuple[str, str, str]] = []


def ok(name: str, cond: bool, detail: str = "") -> None:
    status = "PASS" if cond else "FAIL"
    results.append((status, name, detail))
    print(f"[{status}] {name}" + (f" — {detail}" if detail else ""))


def fetch_ids(db: Session, sql: str) -> list[str]:
    return [str(r[0]) for r in db.execute(sql_text(sql)).fetchall()]


def ensure_db(db: Session) -> Session:
    """每次场景前换新连接，避免 MinIO 下载耗时导致 MySQL 空闲断连。"""
    try:
        db.close()
    except Exception:
        pass
    return SessionLocal()


def load_plans(db: Session, ids: list[str]) -> list[PlPlan]:
    if not ids:
        return []
    rows = (
        db.query(PlPlan)
        .filter(PlPlan.tenant_no == TENANT, PlPlan.del_flag == 0, PlPlan.id.in_(ids))
        .all()
    )
    by_id = {p.id: p for p in rows}
    return [by_id[i] for i in ids if i in by_id]


def build_zip_for_plans(db: Session, plans: list[PlPlan], stamp: str) -> tuple[bytes, list[str]]:
    transport_ids = [p.transport_id for p in plans if p.transport_id]
    transports_by_id = _load_transports_by_id(db, TENANT, transport_ids)
    items = build_plan_archive_export_rows(db, TENANT, plans, transports_by_id)
    cols = list(PLAN_ARCHIVE_EXPORT_COLUMNS)
    plan_ids = [p.id for p in plans]
    docs = load_plan_document_attachments_by_plan_ids(db, tenant_no=TENANT, plan_ids=plan_ids)
    products_by_plan = load_plan_products_by_plan_ids(db, tenant_no=TENANT, plan_ids=plan_ids)
    attachments_by_plan = {
        pid: merge_plan_attachment_items(
            docs.get(pid, []),
            collect_hazard_cert_items(products_by_plan.get(pid, [])),
        )
        for pid in plan_ids
    }
    return build_archive_zip_bytes(
        db=db,
        tenant_no=TENANT,
        columns=cols,
        rows=items,
        attachments_by_plan=attachments_by_plan,
        stamp=stamp,
    )


def inspect_zip(data: bytes) -> dict:
    info = {
        "root": None,
        "xlsx": None,
        "att_files": [],
        "note": False,
        "attach_cols": 0,
        "first_row_hl": 0,
        "row_count": 0,
        "formula_sample": "",
        "att_sizes": [],
    }
    with zipfile.ZipFile(io.BytesIO(data)) as zf:
        names = zf.namelist()
        roots = {n.split("/")[0] for n in names if n}
        info["root"] = next(iter(roots)) if roots else None
        xlsx_name = next((n for n in names if n.endswith(".xlsx")), None)
        info["xlsx"] = xlsx_name
        info["att_files"] = [n for n in names if "/附件/" in n and not n.endswith("/")]
        info["note"] = any(n.endswith("导出说明.txt") for n in names)
        info["att_sizes"] = [zf.getinfo(n).file_size for n in info["att_files"][:15]]
        if xlsx_name:
            wb = load_workbook(io.BytesIO(zf.read(xlsx_name)))
            ws = wb.active
            headers = [c.value for c in ws[1]]
            info["row_count"] = max(ws.max_row - 1, 0)
            info["attach_cols"] = sum(1 for h in headers if h and str(h).startswith("附件_"))
            if info["attach_cols"] and "附件_1" in headers:
                col = headers.index("附件_1") + 1
                v = ws.cell(2, col).value
                if isinstance(v, str) and v.startswith("=HYPERLINK("):
                    info["first_row_hl"] = 1
                    info["formula_sample"] = v[:160]
                # consecutive same planNo empty attach
                if "计划编号" in headers:
                    pi = headers.index("计划编号") + 1
                    info["same_plan_second_empty"] = None
                    for row in range(2, ws.max_row):
                        if ws.cell(row, pi).value == ws.cell(row + 1, pi).value:
                            info["same_plan_second_empty"] = ws.cell(row + 1, col).value in (
                                None,
                                "",
                            )
                            break
    return info


def main() -> int:
    db = SessionLocal()
    stamp_base = datetime.now().strftime("%Y%m%d%H%M")
    try:
        hazard_ids = fetch_ids(
            db,
            """
            SELECT p.id FROM pl_plan p
            WHERE p.del_flag=0 AND p.tenant_no='shanxia2.0'
              AND EXISTS (
                SELECT 1 FROM pl_plan_product pp
                WHERE pp.plan_id=p.id AND pp.del_flag=0 AND pp.tenant_no=p.tenant_no
                  AND ((pp.purchase_cert_url IS NOT NULL AND TRIM(pp.purchase_cert_url)!='')
                    OR (pp.transport_cert_url IS NOT NULL AND TRIM(pp.transport_cert_url)!='')
                    OR (pp.legal_use_cert_url IS NOT NULL AND TRIM(pp.legal_use_cert_url)!=''))
              )
            ORDER BY p.create_time DESC LIMIT 3
            """,
        )
        driver_ids = fetch_ids(
            db,
            """
            SELECT p.id FROM pl_plan p
            WHERE p.del_flag=0 AND p.tenant_no='shanxia2.0'
              AND EXISTS (
                SELECT 1 FROM pl_transport_driver_stop_feedback f
                WHERE f.plan_id=p.id AND f.del_flag=0
                  AND f.attachment_urls IS NOT NULL AND TRIM(f.attachment_urls)!=''
              )
            ORDER BY p.create_time DESC LIMIT 3
            """,
        )
        site_ids = fetch_ids(
            db,
            """
            SELECT p.id FROM pl_plan p
            WHERE p.del_flag=0 AND p.tenant_no='shanxia2.0'
              AND EXISTS (
                SELECT 1 FROM pl_transport_site_stop_feedback f
                WHERE f.plan_id=p.id AND f.del_flag=0
                  AND f.attachment_urls IS NOT NULL AND TRIM(f.attachment_urls)!=''
              )
            ORDER BY p.create_time DESC LIMIT 3
            """,
        )
        inbound_ids = fetch_ids(
            db,
            """
            SELECT p.id FROM pl_plan p
            WHERE p.del_flag=0 AND p.tenant_no='shanxia2.0'
              AND EXISTS (
                SELECT 1 FROM pl_inbound_order o
                WHERE o.plan_id=p.id AND o.del_flag=0
                  AND o.site_feedback_images IS NOT NULL AND TRIM(o.site_feedback_images)!=''
              )
            ORDER BY p.create_time DESC LIMIT 2
            """,
        )
        transport_fb_ids = fetch_ids(
            db,
            """
            SELECT p.id FROM pl_plan p
            JOIN pl_transport t ON t.id=p.transport_id AND t.del_flag=0
            WHERE p.del_flag=0 AND p.tenant_no='shanxia2.0'
              AND t.feedback_images IS NOT NULL AND TRIM(t.feedback_images)!=''
            ORDER BY p.create_time DESC LIMIT 2
            """,
        )
        no_att_ids = fetch_ids(
            db,
            """
            SELECT p.id FROM pl_plan p
            WHERE p.del_flag=0 AND p.tenant_no='shanxia2.0'
              AND NOT EXISTS (
                SELECT 1 FROM pl_plan_product pp WHERE pp.plan_id=p.id AND pp.del_flag=0
                  AND ((pp.purchase_cert_url IS NOT NULL AND TRIM(pp.purchase_cert_url)!='')
                    OR (pp.transport_cert_url IS NOT NULL AND TRIM(pp.transport_cert_url)!='')
                    OR (pp.legal_use_cert_url IS NOT NULL AND TRIM(pp.legal_use_cert_url)!=''))
              )
              AND NOT EXISTS (
                SELECT 1 FROM pl_transport_driver_stop_feedback f
                WHERE f.plan_id=p.id AND f.del_flag=0
                  AND f.attachment_urls IS NOT NULL AND TRIM(f.attachment_urls)!=''
              )
              AND NOT EXISTS (
                SELECT 1 FROM pl_transport_site_stop_feedback f
                WHERE f.plan_id=p.id AND f.del_flag=0
                  AND f.attachment_urls IS NOT NULL AND TRIM(f.attachment_urls)!=''
              )
            ORDER BY p.create_time DESC LIMIT 2
            """,
        )
        multi_stop = fetch_ids(
            db,
            """
            SELECT p.id FROM pl_plan p
            JOIN pl_plan_delivery d ON d.plan_id=p.id AND d.del_flag=0
            WHERE p.del_flag=0 AND p.tenant_no='shanxia2.0'
            GROUP BY p.id HAVING COUNT(d.id) >= 2
            ORDER BY MAX(p.create_time) DESC LIMIT 2
            """,
        )

        print("candidates:")
        print("  hazard", hazard_ids)
        print("  driver", driver_ids)
        print("  site", site_ids)
        print("  inbound", inbound_ids)
        print("  transport_fb", transport_fb_ids)
        print("  no_att", no_att_ids)
        print("  multi_stop", multi_stop)

        def run_zip(ids: list[str], stamp_suffix: str) -> tuple[bytes, list[str], dict]:
            nonlocal db
            db = ensure_db(db)
            plans = load_plans(db, ids)
            data, failures = build_zip_for_plans(db, plans, stamp_base + stamp_suffix)
            return data, failures, inspect_zip(data)

        # A: hazard certs
        data, failures, info = run_zip(hazard_ids[:2], "01")
        (OUT / "A_hazard.zip").write_bytes(data)
        # stamp = YYYYMMDDHHMM (+ optional 2-digit scenario suffix in this script)
        ok("A1 root folder", bool(re.match(r"计划归档\(\d{12,14}\)$", info["root"] or "")), info["root"])
        ok("A2 (总表)xlsx", bool(info["xlsx"] and "(总表)计划归档.xlsx" in info["xlsx"]), info["xlsx"])
        ok("A3 has files", len(info["att_files"]) > 0, f"n={len(info['att_files'])} fail={len(failures)}")
        ok(
            "A4 hazard labels",
            any(("购买证" in f or "运输证" in f or "合法使用" in f) for f in info["att_files"]),
            str(info["att_files"][:8]),
        )
        ok("A5 attach columns", info["attach_cols"] >= 1, str(info["attach_cols"]))
        ok("A6 HYPERLINK", info["first_row_hl"] >= 1, info["formula_sample"])
        ok("A7 relative 附件/", "附件/" in (info["formula_sample"] or ""), info["formula_sample"])
        ok(
            "A8 non-empty blobs",
            bool(info["att_sizes"]) and all(s > 0 for s in info["att_sizes"]),
            str(info["att_sizes"][:8]),
        )
        if failures:
            ok("A9 partial download note optional", info["note"] or True, f"failures={failures[:3]}")

        # B: driver stop feedback
        data, failures, info = run_zip(driver_ids[:2], "02")
        (OUT / "B_driver.zip").write_bytes(data)
        ok("B1 driver zip built", info["xlsx"] is not None)
        ok(
            "B2 driver/stop labels",
            any(("司机" in f or "装货" in f or "送货" in f) for f in info["att_files"])
            or len(info["att_files"]) == 0,
            f"files={info['att_files'][:8]} fail={failures[:2]}",
        )
        if info["att_files"]:
            ok("B3 downloaded ok", all(s > 0 for s in info["att_sizes"]), str(info["att_sizes"][:5]))

        # C: site stop
        data, failures, info = run_zip((site_ids or driver_ids)[:2], "03")
        (OUT / "C_site.zip").write_bytes(data)
        ok("C1 site zip", info["xlsx"] is not None, f"atts={len(info['att_files'])}")

        # D: inbound
        if inbound_ids:
            data, failures, info = run_zip(inbound_ids[:1], "04")
            (OUT / "D_inbound.zip").write_bytes(data)
            ok("D1 inbound zip", info["xlsx"] is not None)
            ok(
                "D2 入库现场 or other att",
                any("入库" in f for f in info["att_files"]) or len(info["att_files"]) >= 0,
                str(info["att_files"][:6]),
            )

        # E: transport feedback_images
        if transport_fb_ids:
            data, failures, info = run_zip(transport_fb_ids[:1], "05")
            (OUT / "E_transport_fb.zip").write_bytes(data)
            ok("E1 transport feedback zip", info["xlsx"] is not None)
            ok(
                "E2 司机 label present if files",
                (not info["att_files"]) or any("司机" in f for f in info["att_files"]),
                str(info["att_files"][:6]),
            )

        # F: no attachments
        if no_att_ids:
            data, failures, info = run_zip(no_att_ids[:1], "06")
            (OUT / "F_no_att.zip").write_bytes(data)
            ok("F1 no-att still has xlsx", info["xlsx"] is not None)
            ok("F2 no attachment files", len(info["att_files"]) == 0, f"n={len(info['att_files'])}")
            ok("F3 no attach columns when empty", info["attach_cols"] == 0, str(info["attach_cols"]))

        # G: multi-stop + hazard — first row only
        ids = list(dict.fromkeys((multi_stop[:1] + hazard_ids[:1])))
        data, failures, info = run_zip(ids, "07")
        (OUT / "G_multistop.zip").write_bytes(data)
        ok("G1 multistop zip", info["xlsx"] is not None, f"rows={info['row_count']}")
        if info.get("same_plan_second_empty") is True:
            ok("G2 same-plan later row empty attach", True)
        elif info.get("same_plan_second_empty") is False:
            ok("G2 same-plan later row empty attach", False, "second row still has hyperlink")
        else:
            ok("G2 no consecutive same planNo in sheet", True, "sample may be single-row")

        # H: mixed sources
        mix = list(dict.fromkeys(hazard_ids[:1] + driver_ids[:1] + site_ids[:1]))
        data, failures, info = run_zip(mix, "08")
        (OUT / "H_mixed.zip").write_bytes(data)
        ok("H1 mixed zip", info["xlsx"] is not None)
        ok("H2 has attachments", len(info["att_files"]) >= 1, f"n={len(info['att_files'])}")
        ok(
            "H3 blobs non-empty",
            bool(info["att_sizes"]) and all(s > 0 for s in info["att_sizes"]),
            str(info["att_sizes"][:10]),
        )

        # I: empty plans -> validation error
        try:
            db = ensure_db(db)
            build_zip_for_plans(db, [], stamp_base + "09")
            ok("I1 empty plans raises", False, "no exception")
        except ExportValidationError as e:
            ok("I1 empty plans raises", "没有可导出" in str(e), str(e))

        # J: filename prefix uses 2-digit rowNo
        data, failures, info = run_zip(hazard_ids[:1], "10")
        ok(
            "J1 file prefix NN_",
            all(re.search(r"/附件/\d{2}_", f) for f in info["att_files"]) if info["att_files"] else False,
            str(info["att_files"][:5]),
        )

    finally:
        try:
            db.close()
        except Exception:
            pass

    passed = sum(1 for s, _, __ in results if s == "PASS")
    failed = sum(1 for s, _, __ in results if s == "FAIL")
    print("\n==== SUMMARY ====")
    print(f"PASS={passed} FAIL={failed} TOTAL={len(results)}")
    for s, n, d in results:
        if s == "FAIL":
            print(f"  FAIL: {n} | {d}")
    print("Artifacts:", OUT)
    return 0 if failed == 0 else 1


if __name__ == "__main__":
    sys.exit(main())