fegger
7e706793fa
- Add AGENTS.md and project-specific Zed skills (odoo-ocr-pipeline, odoo-xml-import, local-vlm-client). - Implement Pydantic schemas for documents, invoices, review results, and VLM responses. - Add unified BaseVLMClient with Ollama implementation and llama.cpp stub. - Build pipeline stages: loader, classifier, digital_pdf/scanned_print/handwritten/mixed_unknown branches, extractor, reviewer, xml_builder. - Add CLI entry point with sidecar JSON and confidence-gated XML output. - Include prompts for classifier, OCR, extraction, and review models. - Add tests with FakeVLMClient; pytest, ruff, and mypy all pass.
73 lines
2.1 KiB
Python
73 lines
2.1 KiB
Python
"""Tests for the document classifier."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from pathlib import Path
|
|
from unittest.mock import patch
|
|
|
|
import pytest
|
|
from PIL import Image
|
|
|
|
from odoo_ocr.pipeline.classifier import classify_document
|
|
from odoo_ocr.schemas.document import DocumentClass
|
|
|
|
|
|
@pytest.fixture
|
|
def blank_image() -> Image.Image:
|
|
return Image.new("RGB", (100, 100), color="white")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_heuristic_digital_pdf(tmp_path: Path) -> None:
|
|
pdf_path = tmp_path / "invoice.pdf"
|
|
# Minimal valid PDF with text via pymupdf
|
|
import fitz
|
|
|
|
doc = fitz.open()
|
|
page = doc.new_page()
|
|
page.insert_text((50, 50), "Invoice number 12345\nVendor: Acme")
|
|
doc.save(str(pdf_path))
|
|
doc.close()
|
|
|
|
from odoo_ocr.config import Settings
|
|
from tests.conftest import FakeVLMClient
|
|
|
|
client = FakeVLMClient()
|
|
result = await classify_document(pdf_path, [], client, Settings())
|
|
assert result.category == DocumentClass.digital_pdf
|
|
assert result.confidence >= 0.90
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_heuristic_scanned_print(blank_image: Image.Image) -> None:
|
|
# Blank image has low handwriting score
|
|
from odoo_ocr.config import Settings
|
|
from tests.conftest import FakeVLMClient
|
|
|
|
client = FakeVLMClient()
|
|
result = await classify_document(Path("scan.jpg"), [blank_image], client, Settings())
|
|
assert result.category == DocumentClass.scanned_print
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_vlm_classifier_fallback(blank_image: Image.Image) -> None:
|
|
|
|
from odoo_ocr.config import Settings
|
|
from tests.conftest import FakeVLMClient
|
|
|
|
client = FakeVLMClient(
|
|
responses={
|
|
"document classifier": {
|
|
"category": "handwritten",
|
|
"confidence": 0.91,
|
|
"reasoning": "handwritten text visible",
|
|
}
|
|
}
|
|
)
|
|
with patch(
|
|
"odoo_ocr.pipeline.classifier._heuristic_classify", return_value=None
|
|
):
|
|
result = await classify_document(Path("unknown.png"), [blank_image], client, Settings())
|
|
assert result.category == DocumentClass.handwritten
|
|
assert result.confidence == 0.91
|