Files
shanxia2.0/backend/scripts/verify_plan_archive_zip_export.py
T
2026-09-26 11:59:39 +08:00

392 lines
16 KiB
Python

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