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()