backend/tests/test_docx_provenance_export.py

"""DOCX export preserves reference provenance through assembly."""

from __future__ import annotations

import io
import zipfile

from backend.core.compat_adapter import (
    payload_to_generated_section,
    reference_sources_from_payload,
    section_to_payload,
)
from backend.core.report_assembler import from_sections_payload, to_docx
from backend.models.report import GeneratedSection, ReferenceSource
from backend.models.schema import SectionDefinition, TemplateSchema


def _minimal_schema() -> TemplateSchema:
    return TemplateSchema(
        version=1,
        sections=[
            SectionDefinition(id="D2", title="Roof coverings", order=0, level=2),
        ],
    )


def test_section_payload_round_trips_reference_sources():
    section = GeneratedSection(
        section_id="D2",
        title="Roof coverings",
        text="The roof comprises slate tiles.",
        reference_sources=[
            ReferenceSource(
                report_filename="past.docx",
                section_id="D2",
                section_title="Roof coverings",
                paragraph_index=2,
            )
        ],
        rag_sources=['Past report "past.docx", section E2 (Roof coverings), paragraph 2'],
    )
    payload = section_to_payload(section, interference_level="minimum", mode="generate")
    assert payload["reference_sources"][0]["report_filename"] == "past.docx"
    assert payload["provenance"]

    rebuilt = payload_to_generated_section("E2", "Roof coverings", payload, _minimal_schema())
    assert rebuilt.reference_sources[0].report_filename == "past.docx"
    assert rebuilt.reference_sources[0].paragraph_index == 2


def test_provenance_only_payload_rebuilds_reference_sources():
    payload = {
        "text": "Slate tiles to pitched rafters.",
        "provenance": [
            {
                "doc_id": "ref.docx",
                "chunk_id": "D2:p2",
                "filename": "ref.docx",
                "section_hint": "D2",
            }
        ],
    }
    sources = reference_sources_from_payload(payload, _minimal_schema())
    assert len(sources) == 1
    assert sources[0].report_filename == "ref.docx"
    assert sources[0].section_id == "D2"
    assert sources[0].paragraph_index == 2


def test_from_sections_payload_docx_omits_internal_source_attribution():
    schema = _minimal_schema()
    payload = {
        "D2": {
            "text": "The roof comprises slate tiles to pitched rafters.",
            "provenance": [
                {
                    "filename": "ref.docx",
                    "chunk_id": "D2:p2",
                    "section_hint": "D2",
                }
            ],
        }
    }
    result = from_sections_payload("tenant-a", schema, payload)
    assert result.sections[0].reference_sources
    docx_bytes = to_docx(result, schema)
    with zipfile.ZipFile(io.BytesIO(docx_bytes)) as zf:
        xml = zf.read("word/document.xml").decode("utf-8")
    assert "slate tiles" in xml
    assert "Source:" not in xml
    assert "UNMATCHED_OBSERVATION" not in xml
    assert "Grounding review required" not in xml