""" 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