File size: 7,905 Bytes
44c2f50 4e6fd2a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 | """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 `<id>.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"
|