baodan/tests/document_ir_test.py

import importlib.util
import json
import sys
from pathlib import Path

import pytest


sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "api"))


def test_pdf_classifier_uses_page_signals_and_quality_budget():
    from insurance.document.pdf_classifier import PageSignal, classify_pdf, plan_ocr_pages

    signals = [
        PageSignal(1, 800, 0.95, 0.05, 10_000),
        PageSignal(2, 0, 0.0, 0.98, 10_000),
        PageSignal(3, 20, 0.3, 0.2, 10_000),
    ]

    assert classify_pdf(signals) == "mixed"
    plan = plan_ocr_pages(signals, max_pages=1, max_pixels=100_000)
    assert plan["selectedPages"] == [2]
    assert plan["skippedPages"] == [3]
    assert plan["status"] == "partial"


def test_layout_preserves_real_bbox_and_stable_reading_order():
    from insurance.document.layout import order_text_blocks

    blocks = order_text_blocks([
        (200, 10, 280, 30, "右侧", 0, 0),
        (10, 10, 100, 30, "左侧", 0, 0),
        (-5, 40, 120, 60, "下一行", 0, 0),
    ], 1, 300, 400)

    assert [item["text"] for item in blocks] == ["左侧", "右侧", "下一行"]
    assert blocks[2]["bbox"] == [0.0, 40.0, 120.0, 60.0]
    assert all(item["id"].startswith("p1-b") for item in blocks)


def test_table_rows_do_not_merge_only_by_policy_year_and_keep_conflicts():
    from insurance.document.tables import merge_row_candidates

    rows, conflicts = merge_row_candidates([
        {"tableId": "t1", "scenarioType": "base", "policyYear": 10, "rowVariant": "before", "value": 100},
        {"tableId": "t1", "scenarioType": "base", "policyYear": 10, "rowVariant": "after", "value": 80},
        {"tableId": "t1", "scenarioType": "optimistic", "policyYear": 10, "rowVariant": "before", "value": 130},
        {"tableId": "t1", "scenarioType": "base", "policyYear": 20, "rowVariant": "before", "value": 200},
        {"tableId": "t1", "scenarioType": "base", "policyYear": 20, "rowVariant": "before", "value": 210},
    ])

    assert len(rows) == 5
    assert len(conflicts) == 1
    assert sum(1 for row in rows if row.get("status") == "conflict") == 2


def test_cross_page_table_continuation_requires_geometry_and_header_match():
    from insurance.document.tables import link_table_continuations

    linked = link_table_continuations([
        {
            "tableId": "t1", "scenarioType": "base", "sourcePages": [3],
            "headers": [{"text": "年度"}, {"text": "退保价值"}],
            "bbox": [20, 100, 500, 700], "columnX": [20, 200],
        },
        {
            "tableId": "t2", "scenarioType": "base", "sourcePages": [4],
            "headers": [{"text": "年度"}, {"text": "退保价值"}],
            "bbox": [21, 80, 501, 700], "columnX": [21, 201],
        },
    ])

    assert linked[1]["continuationOf"] == "t1"


def test_evidence_ref_matches_frozen_schema_and_supports_reverse_index():
    import jsonschema

    from insurance.document.evidence import EvidenceIndex, build_evidence_record, build_evidence_ref

    evidence = build_evidence_ref(7, 3, [10, 20, 80, 40], "HK$ 100,000", table_id="t1", row_id="r1", column_id="c2")
    schema = json.loads(Path("docs/schemas/field-value.schema.json").read_text(encoding="utf-8"))
    jsonschema.validate({"status": "extracted", "value": 100000, "currency": "HKD", "evidence": [evidence]}, schema)

    record = build_evidence_record("policy.annualPremium", evidence, "HK$ 100,000")
    index = EvidenceIndex([record])
    assert index.for_field("policy.annualPremium") == [record]
    assert index.fields_for_region(7, 3, [10, 20, 80, 40]) == ["policy.annualPremium"]
    assert record["snippet"] == "HK$ 100,000"


def test_extraction_evidence_only_uses_value_found_in_real_page_blocks():
    from insurance.document.evidence import map_extraction_evidence

    document_ir = {
        "pages": [{
            "pageNumber": 2,
            "textBlocks": [
                {"id": "p2-b1", "text": "Annual Premium HK$ 100,000", "bbox": [20, 40, 240, 60], "source": "native", "confidence": None},
                {"id": "p2-b2", "text": "Policy Year 10 Total Surrender 350,000", "bbox": [20, 100, 300, 120], "source": "native", "confidence": None},
            ],
            "ocrBlocks": [],
        }],
    }
    data = {
        "policy": {"annual_premium": 100000, "sum_insured": 999999},
        "benefit_illustration": [{
            "policy_year": 10,
            "total_surrender_value": 350000,
            "source_page": 2,
            "table_id": "t1",
            "scenario_type": "base",
        }],
    }

    evidence = map_extraction_evidence(3, data, document_ir)
    paths = {item["fieldPath"] for item in evidence}
    assert "policy.annual_premium" in paths
    assert "benefit.base.10.default.total_surrender_value" in paths
    assert "policy.sum_insured" not in paths
    assert all(item["evidence"]["bbox"] for item in evidence)


def test_extraction_evidence_maps_iul_account_and_death_benefit_columns():
    from insurance.document.evidence import map_extraction_evidence

    document_ir = {
        "pages": [{
            "pageNumber": 3,
            "textBlocks": [{
                "id": "p3-b1",
                "text": "1 21,100 0 2,000,000",
                "bbox": [20, 100, 520, 120],
                "source": "native",
                "confidence": None,
            }],
            "ocrBlocks": [],
        }],
    }
    data = {
        "benefit_illustration": [{
            "policy_year": 1,
            "non_guaranteed_account_value": 21100,
            "non_guaranteed_death_benefit": 2000000,
            "source_page": 3,
            "scenario_type": "base",
        }],
    }

    evidence = map_extraction_evidence(9, data, document_ir)
    paths = {item["fieldPath"] for item in evidence}
    assert "benefit.base.1.default.non_guaranteed_account_value" in paths
    assert "benefit.base.1.default.non_guaranteed_death_benefit" in paths
    assert all(item["evidence"]["pageNumber"] == 3 for item in evidence)


def test_document_ir_text_is_not_globally_truncated():
    from insurance.ppt.extraction import _extract_document_text

    long_text = "保单年度 退保价值 100000\n" * 8000
    document_ir = {
        "document": {"status": "parsed"},
        "pages": [{
            "pageNumber": 1,
            "pageClass": "native",
            "nativeTextQuality": 0.99,
            "ocrQuality": None,
            "textBlocks": [{"text": long_text}],
            "ocrBlocks": [],
        }],
    }

    text, _qualities, used_ir, status = _extract_document_text("unused.pdf", document_ir)
    assert used_ir is True
    assert status == "parsed"
    assert len(text) > 120_000
    assert text.endswith(long_text.rstrip())


def test_invalid_document_returns_structured_error(tmp_path):
    from insurance.document.errors import DocumentIngestError
    from insurance.document.ingest import ingest_pdf

    invalid = tmp_path / "not-a-pdf.pdf"
    invalid.write_bytes(b"plain text")
    with pytest.raises(DocumentIngestError) as captured:
        ingest_pdf(str(invalid))

    payload = captured.value.to_dict()
    assert payload["category"] == "quality"
    assert payload["userMessage"] == payload["message"]
    assert {key: payload[key] for key in ("code", "message", "retryable", "details")} == {
        "code": "DOCUMENT_TYPE_INVALID",
        "message": "文件内容不是有效 PDF",
        "retryable": False,
        "details": {"detectedMimeType": "application/pdf"},
    }


def test_document_cache_key_changes_with_every_runtime_version():
    from insurance.document.ingest import build_cache_key

    base = build_cache_key(
        "a" * 64,
        parser_version="parser-1",
        ocr_version="ocr-1",
        profile_code="profile-a",
        profile_version="1",
        rules_version="rules-1",
        model_version="model-1",
        prompt_version="prompt-1",
    )
    dimensions = {
        "parser_version": "parser-2",
        "ocr_version": "ocr-2",
        "profile_code": "profile-b",
        "profile_version": "2",
        "rules_version": "rules-2",
        "model_version": "model-2",
        "prompt_version": "prompt-2",
    }
    defaults = {
        "parser_version": "parser-1",
        "ocr_version": "ocr-1",
        "profile_code": "profile-a",
        "profile_version": "1",
        "rules_version": "rules-1",
        "model_version": "model-1",
        "prompt_version": "prompt-1",
    }
    for key, changed in dimensions.items():
        options = dict(defaults)
        options[key] = changed
        assert build_cache_key("a" * 64, **options) != base


def test_public_payload_removes_server_paths_recursively():
    from insurance.utils.public_payload import strip_server_paths

    cleaned = strip_server_paths({
        "files": [{"name": "plan.pdf", "path": "D:/private/plan.pdf"}],
        "output": {"filePath": "/app/private/output.pptx", "downloadUrl": "/download/1"},
    })

    assert cleaned == {
        "files": [{"name": "plan.pdf"}],
        "output": {"downloadUrl": "/download/1"},
    }


def test_model_public_payloads_do_not_expose_local_paths():
    from insurance.models.generation_task import GenerationTask
    from insurance.models.poster_case_upload import PosterCaseUpload
    from insurance.models.ppt_session import PptSession

    session = PptSession(
        id="session-1",
        user_id="user-1",
        files_json=json.dumps([{"name": "plan.pdf", "path": "D:/private/plan.pdf"}]),
        extractions_json=json.dumps([{"pdfName": "plan.pdf", "pdfPath": "/app/private/plan.pdf"}]),
        ppt_path="/app/private/output.pptx",
    )
    session_payload = session.to_dict()
    assert "path" not in session_payload["files"][0]
    assert "pdfPath" not in session_payload["extractions"][0]
    assert "ppt_path" not in session_payload
    assert session_payload["has_ppt"] is True

    task = GenerationTask(
        id="task-1",
        user_id="user-1",
        artifact_type="ppt",
        operation="generate",
        workspace_id="session-1",
        output_json=json.dumps({"filePath": "/app/private/output.pptx", "downloadUrl": "/download/1"}),
    )
    assert task.to_dict()["output"] == {"downloadUrl": "/download/1"}

    case = PosterCaseUpload(
        id=1,
        user_id="user-1",
        product_id="product-1",
        source_file_url="D:/private/plan.pdf",
    )
    case_payload = case.to_dict()
    assert "sourceFileUrl" not in case_payload
    assert case_payload["sourceFileAvailable"] is True


def test_migration_036_creates_document_ir_tables():
    from flask import Flask
    from sqlalchemy import inspect

    from insurance.db.database import db
    from insurance.db.migrate_036 import migrate

    app = Flask(__name__)
    app.config.update(
        SQLALCHEMY_DATABASE_URI="sqlite:///:memory:",
        SQLALCHEMY_TRACK_MODIFICATIONS=False,
    )
    db.init_app(app)
    with app.app_context():
        migrate()
        tables = set(inspect(db.engine).get_table_names())

    assert {
        "insurance_documents",
        "insurance_document_pages",
        "insurance_document_tables",
    } <= tables


def test_persisted_document_ir_can_be_reloaded_against_frozen_schema():
    from flask import Flask
    import jsonschema

    from insurance.db.database import db
    from insurance.db.migrate_036 import migrate
    from insurance.document.service import load_document_ir
    from insurance.models.insurance_document import InsuranceDocument, InsuranceDocumentPage

    app = Flask(__name__)
    app.config.update(
        SQLALCHEMY_DATABASE_URI="sqlite:///:memory:",
        SQLALCHEMY_TRACK_MODIFICATIONS=False,
    )
    db.init_app(app)
    with app.app_context():
        migrate()
        document = InsuranceDocument(
            user_id="user-1", sha256="d" * 64, storage_key="/private/plan.pdf",
            original_name="plan.pdf", mime_type="application/pdf", file_size=1,
            page_count=1, pdf_kind="native", status="parsed", parser_version="test-v1",
            quality_json="{}",
        )
        db.session.add(document)
        db.session.flush()
        db.session.add(InsuranceDocumentPage(
            document_id=document.id, page_number=1, width=595, height=842,
            page_class="native", native_text_quality=1.0,
            text_blocks_json=json.dumps([{
                "id": "p1-b1", "text": "Premium 100000", "bbox": [10, 10, 100, 30],
                "source": "native", "confidence": None,
            }]),
            ocr_blocks_json="[]", images_json="[]",
        ))
        db.session.commit()
        payload = load_document_ir(document)
        schema = json.loads(Path("docs/schemas/document-ir.schema.json").read_text(encoding="utf-8"))
        jsonschema.validate(payload, schema)

    assert "documentId" not in payload["pages"][0]


@pytest.mark.skipif(
    not (importlib.util.find_spec("reportlab") and importlib.util.find_spec("pdfplumber")),
    reason="需要 PDF 生成和版面解析组件验证持久化",
)
def test_document_ingest_persists_pages_without_exposing_storage_key(tmp_path):
    from flask import Flask
    from reportlab.pdfgen import canvas

    from insurance.db.database import db
    from insurance.db.migrate_036 import migrate
    from insurance.document.service import ingest_for_processing

    pdf_path = tmp_path / "persist.pdf"
    pdf = canvas.Canvas(str(pdf_path), pagesize=(595, 842))
    pdf.drawString(72, 770, "Policy Premium 100000")
    pdf.save()

    app = Flask(__name__)
    app.config.update(
        SQLALCHEMY_DATABASE_URI="sqlite:///:memory:",
        SQLALCHEMY_TRACK_MODIFICATIONS=False,
    )
    db.init_app(app)
    with app.app_context():
        migrate()
        document, _result = ingest_for_processing(str(pdf_path), "user-1", original_name="persist.pdf")
        payload = document.to_public_dict()

    assert document.id
    assert payload["pageCount"] == 1
    assert "storageKey" not in payload


@pytest.mark.skipif(
    not (
        importlib.util.find_spec("reportlab")
        and (
            importlib.util.find_spec("fitz")
            or importlib.util.find_spec("pymupdf")
            or importlib.util.find_spec("pdfplumber")
        )
    ),
    reason="需要 PDF 生成和版面解析组件验证真实坐标摄取",
)
def test_ingest_real_pdf_matches_document_ir_schema(tmp_path):
    import jsonschema
    from reportlab.pdfgen import canvas

    pdf_path = tmp_path / "native.pdf"
    document = canvas.Canvas(str(pdf_path), pagesize=(595, 842))
    document.drawString(72, 770, "Policy Year Premium Cash Value 100000")
    document.save()

    from insurance.document.ingest import ingest_pdf

    result = ingest_pdf(str(pdf_path), ocr_page_budget=0)
    schema = json.loads(Path("docs/schemas/document-ir.schema.json").read_text(encoding="utf-8"))
    jsonschema.validate(result.document_ir, schema)
    assert result.document_ir["document"]["pageCount"] == 1
    assert result.document_ir["pages"][0]["textBlocks"][0]["bbox"]
    page_text = " ".join(block["text"] for block in result.document_ir["pages"][0]["textBlocks"])
    assert "Policy Year" in page_text