proposa-ai/scripts/tests/test_step4_response_table.py

"""响应表覆盖必须使用Word网格，原要求单元格不可压平或移位。"""
import os
import tempfile
import unittest
import xml.etree.ElementTree as ET

from docx import Document
from docx.enum.text import WD_ALIGN_PARAGRAPH

from models import Chapter, ExtractedTable, TableCell, TenderAnalysis
from step4_writing import (
    _fill_current_response_table, _overlay_table_xml_text,
    _save_chapter_docx, _w_tag,
)


def text(element):
    return ''.join(t.text or '' for t in element.iter(_w_tag('t')))


def canonical_cell(cell):
    return ET.tostring(ET.fromstring(cell._tc.xml))


class ResponseTableTests(unittest.TestCase):
    def fixture(self):
        doc = Document()
        table = doc.add_table(rows=4, cols=5)
        headers = ['类别', '序号', '资格要求', '是否响应', '证明材料']
        for i, value in enumerate(headers):
            table.cell(0, i).text = value
        table.cell(1, 0).merge(table.cell(3, 0)).text = '资格'
        cells = [[TableCell(text=value, col=i) for i, value in enumerate(headers)]]
        for r in range(1, 4):
            table.cell(r, 1).text = str(r)
            p = table.cell(r, 2).paragraphs[0]
            p.alignment = WD_ALIGN_PARAGRAPH.CENTER
            p.add_run(f'{r}、提供').bold = True
            p.add_run('示例证明。').underline = True
            table.cell(r, 2).add_paragraph('原说明必须保留。')
            cells.append([
                TableCell(text='资格' if r == 1 else '', row=r, col=0,
                          rowspan=3 if r == 1 else 1,
                          is_merged_origin=r == 1, is_merged_continuation=r > 1),
                TableCell(text=str(r), row=r, col=1),
                TableCell(text=f'{r}、提供示例证明。\n原说明必须保留。', row=r, col=2),
                TableCell(row=r, col=3), TableCell(row=r, col=4),
            ])
        source = ExtractedTable(rows=4, cols=5, cells=cells,
                                table_type='qualification_response', source_type='tender_pdf')
        return doc, table, source

    def test_vertical_merge_does_not_shift_requirements_or_flatten_runs(self):
        _, table, source = self.fixture()
        xml = ET.fromstring(table._tbl.xml)
        before = [[ET.tostring(tc) for tc in row.findall(_w_tag('tc'))[:3]]
                  for row in xml.findall(_w_tag('tr'))]
        filled = _fill_current_response_table(source, '资格条件响应表')
        _overlay_table_xml_text(xml, filled)
        after = [[ET.tostring(tc) for tc in row.findall(_w_tag('tc'))[:3]]
                 for row in xml.findall(_w_tag('tr'))]
        self.assertEqual(before, after, '原要求、合并续格及多段混合格式必须原样保留')
        for row in xml.findall(_w_tag('tr'))[1:]:
            self.assertEqual(text(row.findall(_w_tag('tc'))[3]), '是')

    def test_horizontal_and_vertical_merge_keep_response_origin_only(self):
        _, table, source = self.fixture()
        table.cell(1, 3).merge(table.cell(2, 3))
        source.cells[1][3].rowspan = 2
        source.cells[1][3].is_merged_origin = True
        source.cells[2][3].is_merged_continuation = True
        table.cell(3, 1).merge(table.cell(3, 2))
        source.cells[3][1].text = '3\n3、提供示例证明。\n原说明必须保留。'
        source.cells[3][1].colspan = 2
        source.cells[3][1].is_merged_origin = True
        source.cells[3][2].is_merged_continuation = True
        filled = _fill_current_response_table(source, '资格条件响应表')
        self.assertEqual(filled.cells[2][3].text, '', '合并续格不得被填充')
        xml = ET.fromstring(table._tbl.xml)
        _overlay_table_xml_text(xml, filled)
        rows = xml.findall(_w_tag('tr'))
        self.assertEqual(text(rows[2].findall(_w_tag('tc'))[3]), '')
        self.assertEqual(text(rows[3].findall(_w_tag('tc'))[2]), '是')

    def test_actual_chapter_write_preserves_step1_requirement_xml(self):
        doc, table, source = self.fixture()
        with tempfile.TemporaryDirectory() as directory:
            source.artifact_path = os.path.join(directory, 'source.docx')
            doc.save(source.artifact_path)
            template = Document()
            template.add_paragraph('第一章 资格', style='Heading 1')
            template.add_paragraph('%%资格条件响应表%%')
            template.add_paragraph('第二章 报价', style='Heading 1')
            template_path = os.path.join(directory, 'template.docx')
            template.save(template_path)
            path = _save_chapter_docx(
                Chapter(id='1', title='资格', template_chapter_id='1'), directory,
                template_path=template_path,
                analysis=TenderAnalysis(project_name='示例项目', tender_tables=[source]),
            )
            result = Document(path).tables[0]
            for r in range(1, 4):
                self.assertEqual(result.cell(r, 2).text, table.cell(r, 2).text)
                self.assertEqual(canonical_cell(result.cell(r, 2)), canonical_cell(table.cell(r, 2)))
                self.assertEqual(result.cell(r, 3).text, '是')
            self.assertEqual(source.cells[1][3].text, '', '填表不得污染Step1元数据')

    def test_grid_before_and_existing_response_are_preserved(self):
        _, table, source = self.fixture()
        table.cell(1, 3).text = '已有响应'
        table.cell(2, 3).paragraphs[0].alignment = WD_ALIGN_PARAGRAPH.RIGHT
        table.cell(2, 3).paragraphs[0].add_run(' ').underline = True
        xml = ET.fromstring(table._tbl.xml)
        rows = xml.findall(_w_tag('tr'))
        first_response = ET.tostring(rows[1].findall(_w_tag('tc'))[3])
        rows[2].remove(rows[2].find(_w_tag('tc')))
        pr = ET.Element(_w_tag('trPr'))
        ET.SubElement(pr, _w_tag('gridBefore')).set(_w_tag('val'), '1')
        rows[2].insert(0, pr)
        filled = _fill_current_response_table(source, '资格条件响应表')
        _overlay_table_xml_text(xml, filled)
        self.assertEqual(first_response, ET.tostring(rows[1].findall(_w_tag('tc'))[3]))
        response = rows[2].findall(_w_tag('tc'))[2]
        self.assertEqual(text(response), '是')
        self.assertEqual(response.find('.//' + _w_tag('jc')).get(_w_tag('val')), 'right')
        self.assertIsNotNone(response.find('.//' + _w_tag('u')))
        self.assertEqual(len(response.findall(_w_tag('p'))), 1, '空单元格沿用原段落')
        once = ET.tostring(xml)
        _overlay_table_xml_text(xml, filled)
        self.assertEqual(once, ET.tostring(xml), '重复写出须幂等')


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