#!/usr/bin/env python3 """Exercise a real JSON file through multipart + streaming validation. Optional --import-falkor writes only a UUID-named local test graph and removes it in finally. System project metadata/auth are test doubles; this is NOT a production importer, nor evidence of a browser/proxy/production deployment test. """ from __future__ import annotations import argparse import asyncio import json from pathlib import Path import resource import shutil import sys import tempfile import time import uuid from unittest.mock import patch import httpx from fastapi import FastAPI ROOT = Path(__file__).resolve().parents[1] if str(ROOT) not in sys.path: sys.path.insert(0, str(ROOT)) from app.api import graph_file_import as file_api from app.api.plaza import map_pois from app.auth import get_current_user from app.config import settings from app.project_context import ProjectContext from app.project_lifecycle import _falkor_client def graph_names(db) -> set[str]: return {name.decode() if isinstance(name, bytes) else str(name) for name in db.list_graphs()} async def verify(args): if settings.falkordb_host not in {"localhost", "127.0.0.1", "::1"}: raise RuntimeError("验证仅允许本机 FalkorDB，不连接生产服务器") if args.report.exists(): raise FileExistsError("报告已存在，请选择新路径") if args.schema.stat().st_size > file_api.SCHEMA_BYTES: raise ValueError("Schema 超过大小上限") schema = json.loads(args.schema.read_text(encoding="utf-8")) scratch = "codex_upload_verify_" + uuid.uuid4().hex metadata = {"project_id": scratch, "display_name": "城市知识图谱-贵州*贵阳（本地导入验证）", "schema": schema, "spatial_map": {"enabled": True, "scope": "guizhou", "region_name": "贵阳市", "region_adcode": "520100", "region_level": "city"}} app = FastAPI() app.include_router(file_api.router, prefix="/v1/admin") # Isolated in-process test app only; never bypass the running app's login. app.dependency_overrides[get_current_user] = lambda: {"username": "test_only", "roles": ["admin"]} @app.get("/health") async def health(): return {"ok": True} report = {"file_bytes": args.data.stat().st_size, "metadata_and_auth": "isolated_test_doubles", "transport": "HTTPX ASGI multipart", "falkordb_import": bool(args.import_falkor), "temp_free_before_bytes": shutil.disk_usage(tempfile.gettempdir()).free} created = False db = _falkor_client() if args.import_falkor else None existing_graphs = graph_names(db) if db else set() if scratch in existing_graphs: raise FileExistsError("临时图名冲突，未开始测试") if db: report["existing_graphs_before"] = len(existing_graphs) real_validator = file_api.validate_graph_file last_print = time.monotonic() def validate(source, staging, body, progress): def observed(phase, nodes, relations): nonlocal last_print size = staging.stat().st_size if staging.exists() else 0 report["peak_staging_bytes"] = max(report.get("peak_staging_bytes", 0), size) report["last_progress"] = {"phase": phase, "nodes": nodes, "relations": relations, "staging_bytes": size, "temp_free_bytes": shutil.disk_usage(tempfile.gettempdir()).free} progress(phase, nodes, relations) if time.monotonic() - last_print > 8: print(json.dumps({"phase": phase, "nodes": nodes, "relations": relations}, ensure_ascii=False), flush=True) last_print = time.monotonic() return real_validator(source, staging, body, observed) async def publish(payload, user, *, graph_importer): nonlocal created report["validated_counts"] = payload["counts"] report["file_sha256"] = payload["file_sha256"] if db: created = True counts = await asyncio.to_thread(graph_importer, scratch, payload["graph_data"]) report["actual_graph_counts"] = counts assert counts["nodes"] == payload["counts"]["nodes"] assert counts["relations"] == payload["counts"]["relations"] return {"project": {"project_id": scratch}, "counts": payload["counts"]} started = time.monotonic() latencies = [] failure = None try: async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://local-test", timeout=None) as client: with args.data.open("rb") as source, patch.object(file_api, "validate_graph_file", side_effect=validate), patch( "app.api.projects.provision_validated_project", side_effect=publish, ): upload = asyncio.create_task(client.post( "/v1/admin/projects/provision-file", data={"metadata": json.dumps(metadata, ensure_ascii=False)}, files={"file": (args.data.name, source, "application/json")}, )) while not upload.done(): t = time.monotonic() health = await client.get("/health") assert health.status_code == 200 latencies.append(round((time.monotonic() - t) * 1000, 2)) await asyncio.sleep(0.5) response = await upload report["status_code"] = response.status_code if response.status_code != 200: raise RuntimeError(response.text[:1500]) messages = [json.loads(line) for line in response.text.splitlines()] report["response_events"] = len(messages) if not messages or messages[-1]["event"] != "result": raise RuntimeError(str(messages[-1] if messages else "没有返回结果")) if db: pois = await map_pois(context=ProjectContext("local_test", scratch, scratch), _user=None) report["map_pois"] = pois["total"] assert pois["source"] == "falkordb" assert pois["total"] == report["validated_counts"]["map_pois"] del pois report["passed"] = True except Exception as exc: failure = exc report.update(passed=False, error=str(exc)[:2000]) finally: if db: try: if created and scratch in graph_names(db): db.select_graph(scratch).delete() final_graphs = graph_names(db) if scratch in final_graphs: raise RuntimeError("UUID 临时图清理失败") if not existing_graphs.issubset(final_graphs): raise RuntimeError("验证期间发现原有图谱缺失") report.update(test_graph_removed=True, existing_graphs_preserved=len(existing_graphs)) finally: db.close() report.update(seconds=round(time.monotonic() - started, 2), peak_rss_bytes=resource.getrusage(resource.RUSAGE_SELF).ru_maxrss * (1 if sys.platform == "darwin" else 1024), health_checks=len(latencies), max_health_latency_ms=max(latencies, default=0), temp_free_after_bytes=shutil.disk_usage(tempfile.gettempdir()).free) args.report.parent.mkdir(parents=True, exist_ok=True) with args.report.open("x", encoding="utf-8") as stream: json.dump(report, stream, ensure_ascii=False, indent=2) print(json.dumps(report, ensure_ascii=False, indent=2), flush=True) if failure: raise RuntimeError("验证未通过，原因及资源占用已写入报告") from failure if __name__ == "__main__": parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--schema", type=Path, required=True) parser.add_argument("--data", type=Path, required=True) parser.add_argument("--report", type=Path, required=True) parser.add_argument("--import-falkor", action="store_true") asyncio.run(verify(parser.parse_args()))