fegger
2e81fa57bd
Adds the uv lockfile generated from pyproject.toml and updates tool.mypy.python_version to 3.12 so third-party stubs parse cleanly while keeping runtime compatibility at Python 3.11.
87 lines
2.6 KiB
Python
87 lines
2.6 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
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_empty_pages_raises() -> None:
|
|
"""The VLM fallback must fail cleanly for documents with no renderable pages."""
|
|
from odoo_ocr.config import Settings
|
|
from tests.conftest import FakeVLMClient
|
|
|
|
client = FakeVLMClient()
|
|
with (
|
|
patch("odoo_ocr.pipeline.classifier._heuristic_classify", return_value=None),
|
|
pytest.raises(ValueError, match="no renderable pages"),
|
|
):
|
|
await classify_document(Path("empty.pdf"), [], client, Settings())
|