"""测试库多场景验证:计划归档 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())