"""Unit tests for the extraction layer. All OpenAI calls are mocked — these tests run offline in <1s and never spend money. Real end-to-end tests belong in tests/integration/ and require OPENAI_API_KEY. """ from __future__ import annotations from datetime import date from unittest.mock import MagicMock, patch import pytest from src.extractors import DocumentExtractor, compute_overall_confidence, make_envelope from src.extractors.prompts import PROMPTS, get_prompt from src.schemas import ( ExtractionResult, ExtractionWarning, FieldConfidence, Invoice, Party, Receipt, ) # --- Prompt registry ------------------------------------------------------- def test_prompt_registry_has_expected_types(): assert "invoice" in PROMPTS assert "receipt" in PROMPTS def test_get_prompt_invoice_mentions_key_rules(): p = get_prompt("invoice") assert "CRITICAL EXTRACTION RULES" in p assert "ISO 8601" in p assert "ISO 4217" in p assert "field_confidences" in p def test_get_prompt_unknown_raises(): with pytest.raises(KeyError): get_prompt("bogus") # --- Envelope factory ------------------------------------------------------ def test_make_envelope_invoice_shape(): env_cls = make_envelope(Invoice) fields = env_cls.model_fields assert "data" in fields assert "field_confidences" in fields assert "warnings" in fields # Can instantiate env = env_cls( data=Invoice( invoice_number="INV-1", vendor=Party(name="Acme"), total=10.00, currency="USD", ), field_confidences=[FieldConfidence(field="total", score=0.95)], ) assert env.data.invoice_number == "INV-1" def test_make_envelope_is_cached(): """Envelope classes are cached per domain to avoid re-generating on every call.""" a = make_envelope(Invoice) b = make_envelope(Invoice) assert a is b def test_compute_overall_confidence_mean(): confs = [ FieldConfidence(field="a", score=1.0), FieldConfidence(field="b", score=0.8), FieldConfidence(field="c", score=0.6), ] assert compute_overall_confidence(confs) == 0.8 def test_compute_overall_confidence_empty_is_zero(): assert compute_overall_confidence([]) == 0.0 # --- Extractor end-to-end (mocked) ----------------------------------------- def _mock_openai_call_returning(envelope_cls, envelope_instance): """Build a MagicMock imitating the OpenAI beta.parse response shape.""" usage = MagicMock(prompt_tokens=800, completion_tokens=200) message = MagicMock(parsed=envelope_instance, refusal=None) choice = MagicMock(message=message) return MagicMock(choices=[choice], usage=usage) @patch("src.extractors.openai_client.OpenAI") def test_extractor_invoice_end_to_end_mocked(mock_openai_cls): """DocumentExtractor.extract wires loader -> LLM -> ExtractionResult correctly.""" # Build a fake envelope response the mocked LLM will "return". env_cls = make_envelope(Invoice) envelope_instance = env_cls( data=Invoice( invoice_number="INV-42", vendor=Party(name="Widgets Ltd"), invoice_date=date(2026, 6, 15), total=125.50, subtotal=115.00, tax=10.50, currency="USD", ), field_confidences=[ FieldConfidence(field="invoice_number", score=0.98), FieldConfidence(field="total", score=0.95), FieldConfidence(field="vendor.name", score=0.99), ], warnings=[ ExtractionWarning(field="customer", message="Customer not present on doc", severity="info") ], ) # Wire the OpenAI client mock end-to-end. fake_response = _mock_openai_call_returning(env_cls, envelope_instance) fake_client = MagicMock() fake_client.beta.chat.completions.parse.return_value = fake_response mock_openai_cls.return_value = fake_client # Feed a tiny fake "PDF" — loader will fail gracefully (source_type=empty). # So we patch the loader too, to inject a fake LoadedDocument. with patch("src.extractors.extractor.load_document") as mock_load: from src.extractors.document_loader import LoadedDocument mock_load.return_value = LoadedDocument( text="INVOICE INV-42 from Widgets Ltd. Total: $125.50", images_b64=[], source_type="text_pdf", page_count=1, filename="invoice.pdf", ) extractor = DocumentExtractor() result, metrics = extractor.extract( b"%PDF-fake", filename="invoice.pdf", doc_type="invoice" ) # Assertions on the result assert isinstance(result, ExtractionResult) assert result.document_type == "invoice" assert result.data.invoice_number == "INV-42" assert result.data.total == 125.50 assert 0.9 <= result.overall_confidence <= 1.0 assert len(result.warnings) == 1 assert result.raw_text_snippet.startswith("INVOICE INV-42") # Assertions on metrics assert metrics.input_tokens == 800 assert metrics.output_tokens == 200 assert metrics.cost_usd > 0 @patch("src.extractors.openai_client.OpenAI") def test_extractor_receipt_end_to_end_mocked(mock_openai_cls): """Same shape check for the receipt schema.""" env_cls = make_envelope(Receipt) envelope_instance = env_cls( data=Receipt( merchant="Corner Coffee", transaction_date=date(2026, 6, 20), total=4.75, currency="USD", ), field_confidences=[ FieldConfidence(field="merchant", score=0.99), FieldConfidence(field="total", score=0.97), ], warnings=[], ) fake_response = _mock_openai_call_returning(env_cls, envelope_instance) fake_client = MagicMock() fake_client.beta.chat.completions.parse.return_value = fake_response mock_openai_cls.return_value = fake_client with patch("src.extractors.extractor.load_document") as mock_load: from src.extractors.document_loader import LoadedDocument mock_load.return_value = LoadedDocument( text="Corner Coffee\nTotal: $4.75", source_type="text_pdf", page_count=1, filename="receipt.pdf", ) extractor = DocumentExtractor() result, _ = extractor.extract( b"%PDF-fake", filename="receipt.pdf", doc_type="receipt" ) assert result.document_type == "receipt" assert result.data.merchant == "Corner Coffee" assert result.data.total == 4.75 @patch("src.extractors.openai_client.OpenAI") def test_extractor_rejects_unknown_doc_type(_mock): extractor = DocumentExtractor() with pytest.raises(KeyError): extractor.extract(b"", filename="x.pdf", doc_type="bogus") def test_document_loader_handles_txt(): """Regression guard for the .txt handler added on 2026-07-04. The eval CLI packages inline text records as `.txt`, so the loader must decode them into source_type='text' with no images rendered. Previously fell through the 'unknown format' branch and killed live eval. """ from src.extractors.document_loader import load_document body = "TAN WOON YANN\nDATE: 2018-06-25\nTOTAL SGD 72.00" doc = load_document(body.encode("utf-8"), filename="receipt.txt") assert doc.source_type == "text" assert doc.text == body assert doc.images_b64 == [] assert doc.page_count == 1 # md/log/text also route through this branch for ext in (".md", ".log", ".text"): d = load_document(b"hello", filename=f"x{ext}") assert d.source_type == "text" assert d.text == "hello" # empty text still valid, but reports as empty empty = load_document(b"", filename="blank.txt") assert empty.source_type == "empty"