tests/test_generation.py

"""
tests/test_generation.py — Tests for generation layer (no API required)
"""
import sys
from pathlib import Path
import pytest

sys.path.insert(0, str(Path(__file__).parent.parent))

from models import EvidenceChunk
from generation.citation_builder import extract_citations, format_evidence_context
from generation.prompts import DRAFT_SYSTEM_V2, DRAFT_USER_V2, EXPECTED_SECTIONS


@pytest.fixture
def sample_evidence():
    return [
        EvidenceChunk(chunk_id="aabb0011ccdd2233", text="Contract signed Jan 15, 2024.", source_file="contract.pdf", score=0.95),
        EvidenceChunk(chunk_id="eeff4455aabb6677", text="Monthly fee of $25,000 due.", source_file="contract.pdf", score=0.88),
    ]


class TestCitationBuilder:
    def test_extract_valid_citations(self, sample_evidence):
        draft = "The contract was signed. [CHUNK-aabb0011] Payment is $25k. [CHUNK-eeff4455]"
        citations = extract_citations(draft, sample_evidence)
        assert len(citations) == 2
        ids = {c.short_id for c in citations}
        assert "aabb0011" in ids
        assert "eeff4455" in ids

    def test_extract_no_citations(self, sample_evidence):
        draft = "This text has no citations."
        citations = extract_citations(draft, sample_evidence)
        assert len(citations) == 0

    def test_unknown_citation(self, sample_evidence):
        draft = "Claim here. [CHUNK-deadbeef]"
        citations = extract_citations(draft, sample_evidence)
        assert len(citations) == 1
        assert citations[0].chunk is None

    def test_format_evidence_context(self, sample_evidence):
        ctx = format_evidence_context(sample_evidence)
        assert "[CHUNK-aabb0011]" in ctx
        assert "[CHUNK-eeff4455]" in ctx
        assert "contract.pdf" in ctx


class TestPrompts:
    def test_system_prompt_has_placeholders(self):
        rendered = DRAFT_SYSTEM_V2.format(
            draft_type="Case Fact Summary",
            injected_preferences="- Always list dates chronologically",
        )
        assert "Case Fact Summary" in rendered
        assert "Always list dates" in rendered

    def test_user_prompt_has_placeholders(self):
        rendered = DRAFT_USER_V2.format(
            draft_type="Case Fact Summary",
            doc_names="contract.pdf",
            evidence_context="[evidence here]",
            user_query="Summarize the contract",
        )
        assert "contract.pdf" in rendered
        assert "Summarize" in rendered

    def test_expected_sections_defined(self):
        assert len(EXPECTED_SECTIONS) == 6
        assert "Document Overview" in EXPECTED_SECTIONS