proposa-ai/scripts/tests/test_step1_postamble.py

import tempfile
import unittest
from pathlib import Path
from unittest.mock import patch

from docx import Document

from models import ExtractedTable, TableCell, TenderAnalysis
from pdf_table_to_docx.table_parser import CellInfo, TableInfo
from step6_exporting.docx_builder import _generate_table_data_from_tender
from step1_parsing.document_parser import (
    _extract_items_from_pdf,
    _write_items_to_docx,
)
from step1_parsing.table_extractor import (
    _build_reference_doc_stream,
    _build_reference_fragment_impl,
    _clean_reference_table_artifacts,
    _write_fragment_docx,
    extract_reference_tables_with_llm,
    extract_tables_from_template_docx,
    write_reference_tables_to_docx,
)


class Step1PostambleTests(unittest.TestCase):
    def _make_bid_detail_docx(self, path: str) -> None:
        """构造“三、投标报价分项明细表 + 明细1/明细2 两张表”的参考文档。"""
        doc = Document()
        doc.add_paragraph("三、投标报价分项明细表")
        doc.add_paragraph("明细1：人员费用")
        table1 = doc.add_table(rows=2, cols=3)
        table1.cell(0, 0).text = "编号"
        table1.cell(0, 1).text = "内容"
        table1.cell(0, 2).text = "金额"
        table1.cell(1, 0).text = "1"
        table1.cell(1, 1).text = "人员费用"
        table1.cell(1, 2).text = "4930728"
        doc.add_paragraph("明细2：办公费用")
        table2 = doc.add_table(rows=2, cols=3)
        table2.cell(0, 0).text = "编号"
        table2.cell(0, 1).text = "内容"
        table2.cell(0, 2).text = "金额"
        table2.cell(1, 0).text = "1"
        table2.cell(1, 1).text = "办公杂费"
        table2.cell(1, 2).text = "3600"
        doc.save(path)

    def test_reference_table_entry_marks_source_provenance(self):
        table = ExtractedTable(table_id="REF-T01", source_type="template_docx")
        with (
            patch(
                "step1_parsing.table_extractor.extract_tables_from_template_docx",
                return_value=[table],
            ),
            patch(
                "step1_parsing.table_extractor.enrich_reference_tables_with_llm",
                side_effect=lambda tables, docx_path: tables,
            ),
        ):
            result = extract_reference_tables_with_llm("reference.docx")

        self.assertEqual(result[0].source_type, "reference_docx")

    def test_pdf_table_postamble_survives_model_and_docx_output(self):
        postamble = (
            "说明：（1）所有价格均系用人民币表示。\n"
            "（2）投标人应按照招标需求报价。"
        )
        llm_items = [{
            "name": "3.投标报价汇总表格式",
            "type": "表",
            "pages": [1],
            "preamble": "项目名称：",
            "postamble": postamble,
            "headers": ["序号", "分项名称"],
            "rows": [["1", "人员费用"]],
        }]
        table = TableInfo(
            page_num=0,
            rows=2,
            cols=2,
            cells=[
                [CellInfo(text="序号"), CellInfo(text="分项名称")],
                [CellInfo(text="1"), CellInfo(text="人员费用")],
            ],
            title_hint="3.投标报价汇总表格式",
            original_pages=[1],
        )

        with patch(
            "step1_parsing.document_parser._read_pdf_pages",
            return_value=[(1, "PDF 正文")],
        ), patch(
            "step1_parsing.document_parser._llm_extract_all_items",
            return_value=llm_items,
        ), patch(
            "pdf_table_to_docx.extractor.extract_tables_from_pdf",
            return_value=[table],
        ):
            items = _extract_items_from_pdf("tender.pdf", "招标文件", object())

        self.assertEqual(len(items), 1)
        self.assertEqual(items[0].postamble, postamble)
        self.assertEqual(items[0].source_table.postamble, postamble)

        with tempfile.TemporaryDirectory() as temp_dir:
            _write_items_to_docx(items, temp_dir)
            output = next(Path(temp_dir).glob("*.docx"))
            doc = Document(output)

            self.assertEqual(items[0].artifact_path, str(output.resolve()))
            self.assertEqual(
                items[0].source_table.artifact_path, str(output.resolve())
            )

        texts = [p.text for p in doc.paragraphs if p.text.strip()]
        self.assertEqual(texts[-2:], postamble.splitlines())

    def test_step6_table_payload_keeps_tender_postamble(self):
        table = ExtractedTable(
            table_id="T-1",
            rows=2,
            cols=2,
            cells=[
                [TableCell(text="序号"), TableCell(text="分项名称")],
                [TableCell(text="1"), TableCell(text="人员费用")],
            ],
            title_hint="投标报价汇总表",
            table_type="bid_summary_table",
            postamble="说明：合计应与开标一览表报价相等。",
        )
        payload = _generate_table_data_from_tender(
            "投标报价汇总表",
            {"bid_summary_table": [table]},
            TenderAnalysis(project_name="测试项目"),
            None,
        )

        self.assertEqual(payload[3], table.postamble)

    def test_reference_table_artifact_clones_native_merge_xml(self):
        with tempfile.TemporaryDirectory() as temp_dir:
            source_path = Path(temp_dir) / "reference.docx"
            output_dir = Path(temp_dir) / "tables"

            source = Document()
            source.add_paragraph("设备维修记录")
            table = source.add_table(rows=4, cols=4)
            table.cell(0, 0).merge(table.cell(2, 1)).text = "跨行跨列"
            table.cell(0, 2).merge(table.cell(0, 3)).text = "横向合并"
            table.cell(1, 2).text = "内容"
            source.save(source_path)

            extracted = extract_tables_from_template_docx(str(source_path))
            written = write_reference_tables_to_docx(
                extracted, str(output_dir)
            )

            self.assertEqual(written, 1)
            self.assertTrue(extracted[0].artifact_path)
            artifact = Document(extracted[0].artifact_path)
            self.assertEqual(len(artifact.tables), 1)
            xml = artifact.tables[0]._tbl.xml
            self.assertIn("w:gridSpan", xml)
            self.assertIn("w:vMerge", xml)
            self.assertIn("跨行跨列", xml)

    def test_reference_doc_stream_runs_in_isolated_worker(self):
        with tempfile.TemporaryDirectory() as temp_dir:
            source_path = Path(temp_dir) / "reference.docx"
            source = Document()
            source.add_paragraph("表格前说明")
            table = source.add_table(rows=1, cols=2)
            table.cell(0, 0).text = "字段一"
            table.cell(0, 1).text = "字段二"
            source.add_paragraph("表格后说明")
            source.save(source_path)

            stream = _build_reference_doc_stream(str(source_path))

            self.assertEqual([item["kind"] for item in stream], ["p", "tbl", "p"])
            self.assertEqual(stream[0]["text"], "表格前说明")
            self.assertIn("字段一", stream[1]["text"])
            self.assertEqual(stream[2]["text"], "表格后说明")

    def test_reference_doc_stream_native_crash_degrades_to_empty_context(self):
        with patch(
            "step1_parsing.table_extractor._run_reference_xml_worker",
            return_value=None,
        ):
            self.assertEqual(_build_reference_doc_stream("large.docx"), [])

    def test_bid_detail_fragment_preserves_paragraphs_and_multiple_tables(self):
        with tempfile.TemporaryDirectory() as temp_dir:
            source_path = Path(temp_dir) / "reference.docx"
            self._make_bid_detail_docx(str(source_path))

            fragment_xml = _build_reference_fragment_impl(
                str(source_path), 0, 4
            )
            self.assertEqual(len(fragment_xml), 5)

            output_path = Path(temp_dir) / "投标报价分项明细表.docx"
            _write_fragment_docx(fragment_xml, str(output_path))

            doc = Document(output_path)
            texts = [paragraph.text for paragraph in doc.paragraphs]
            self.assertIn("三、投标报价分项明细表", texts)
            self.assertIn("明细1：人员费用", texts)
            self.assertIn("明细2：办公费用", texts)
            self.assertEqual(len(doc.tables), 2)
            self.assertIn("人员费用", doc.tables[0].cell(1, 1).text)
            self.assertIn("办公杂费", doc.tables[1].cell(1, 1).text)

    def test_bid_detail_fragment_written_as_single_named_docx(self):
        with tempfile.TemporaryDirectory() as temp_dir:
            source_path = Path(temp_dir) / "reference.docx"
            self._make_bid_detail_docx(str(source_path))
            fragment_xml = _build_reference_fragment_impl(
                str(source_path), 0, 4
            )

            output_dir = Path(temp_dir) / "tables"
            output_dir.mkdir()
            stale_numbered = output_dir / "08_旧表.docx"
            stale_numbered.write_bytes(b"old")
            stale_named = output_dir / "投标报价分项明细表.docx"
            stale_named.write_bytes(b"old")

            fragment = ExtractedTable(
                table_id="FRAGMENT",
                source_type="reference_docx",
                title_hint="投标报价分项明细表",
                caption="投标报价分项明细表",
                table_type="bid_detail_table",
            )
            fragment.is_fragment = True
            fragment.fragment_xml = fragment_xml

            normal = ExtractedTable(
                table_id="NORMAL",
                source_type="reference_docx",
                title_hint="其他表",
                table_type="other",
                rows=1,
                cols=1,
                cells=[[TableCell(text="内容")]],
            )

            written = write_reference_tables_to_docx(
                [fragment, normal], str(output_dir)
            )

            self.assertEqual(written, 2)
            self.assertFalse(stale_numbered.exists())
            self.assertTrue((output_dir / "投标报价分项明细表.docx").exists())
            self.assertTrue(
                fragment.artifact_path.endswith("投标报价分项明细表.docx")
            )
            self.assertTrue(
                fragment.fragment_docx_path.endswith("投标报价分项明细表.docx")
            )
            artifact = Document(
                output_dir / "投标报价分项明细表.docx"
            )
            self.assertEqual(len(artifact.tables), 2)
            artifacts = sorted(path.name for path in output_dir.glob("*.docx"))
            self.assertEqual(artifacts.count("投标报价分项明细表.docx"), 1)
            self.assertEqual(len(artifacts), 2)
            self.assertTrue(any(name.startswith("02_") for name in artifacts))

    def test_clean_reference_table_artifacts_removes_numbered_and_named_stale_files(self):
        with tempfile.TemporaryDirectory() as temp_dir:
            output_dir = Path(temp_dir)
            (output_dir / "01_旧表.docx").write_bytes(b"old")
            (output_dir / "投标报价分项明细表.docx").write_bytes(b"old")
            (output_dir / "keep.txt").write_text("keep")

            removed = _clean_reference_table_artifacts(str(output_dir))

            self.assertEqual(removed, 2)
            self.assertFalse((output_dir / "01_旧表.docx").exists())
            self.assertFalse((output_dir / "投标报价分项明细表.docx").exists())
            self.assertTrue((output_dir / "keep.txt").exists())

    def test_extract_reference_tables_with_llm_builds_bid_detail_fragment(self):
        with tempfile.TemporaryDirectory() as temp_dir:
            source_path = Path(temp_dir) / "reference.docx"
            self._make_bid_detail_docx(str(source_path))
            output_dir = Path(temp_dir) / "tables"

            with patch(
                "step1_parsing.table_extractor.enrich_reference_tables_with_llm",
                side_effect=lambda tables, docx_path: tables,
            ):
                result = extract_reference_tables_with_llm(
                    str(source_path), output_dir=str(output_dir)
                )

            self.assertEqual(len(result), 1)
            fragment = result[0]
            self.assertTrue(getattr(fragment, "is_fragment", False))
            self.assertEqual(fragment.title_hint, "投标报价分项明细表")
            self.assertEqual(fragment.fragment_source_indices, [0, 1])
            self.assertTrue(
                (output_dir / "投标报价分项明细表.docx").exists()
            )
            artifact = Document(
                output_dir / "投标报价分项明细表.docx"
            )
            self.assertEqual(len(artifact.tables), 2)
            texts = [paragraph.text for paragraph in artifact.paragraphs]
            self.assertIn("三、投标报价分项明细表", texts)
            self.assertIn("明细2：办公费用", texts)


if __name__ == "__main__":
    unittest.main()