72 lines
3.4 KiB
Python
72 lines
3.4 KiB
Python
import shutil
|
|
import subprocess
|
|
import sys
|
|
import tempfile
|
|
import unittest
|
|
from pathlib import Path
|
|
from unittest.mock import patch
|
|
|
|
from PIL import Image, ImageDraw, ImageFont, PdfParser
|
|
|
|
sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "scripts"))
|
|
from scan_document import create_document, check_ocr
|
|
|
|
|
|
class ScanOcrTest(unittest.TestCase):
|
|
def test_missing_language_is_reported_before_processing(self):
|
|
with patch("scan_document.run_ocr_command", return_value="List of available languages (1):\neng\n"):
|
|
with self.assertRaisesRegex(RuntimeError, "deu"):
|
|
check_ocr("deu+eng")
|
|
with self.assertRaises(ValueError):
|
|
check_ocr("--bad-language")
|
|
|
|
@unittest.skipUnless(shutil.which("tesseract") and shutil.which("pdftotext"), "Tesseract and Poppler required")
|
|
def test_ocr_keeps_receipt_canvas_blank_back_and_page_order_with_real_extractable_text(self):
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
root = Path(directory)
|
|
pages = []
|
|
for index, text in enumerate(["FEDEO Bon 12345", "", "Zweite Seite 24680"]):
|
|
page = root / f"page-{index}.png"
|
|
image = Image.new("RGB", (1240, 1754), (45, 45, 45) if index == 0 else "white")
|
|
draw = ImageDraw.Draw(image)
|
|
if index == 0:
|
|
draw.rectangle((260, 150, 1000, 1550), fill="white")
|
|
if text:
|
|
draw.text((300, 300), text, fill="black", font=ImageFont.load_default(size=44))
|
|
draw.text((300, 390), "Summe EUR 12,34", fill="black", font=ImageFont.load_default(size=36))
|
|
image.save(page)
|
|
pages.append(page)
|
|
output = root / "searchable.pdf"
|
|
result = create_document(pages, output, 150, ocr=True, background_cleanup=True)
|
|
self.assertEqual(result["pageCount"], 3)
|
|
with PdfParser.PdfParser(filename=str(output)) as pdf:
|
|
self.assertEqual(len(pdf.pages), 3)
|
|
for reference in pdf.pages:
|
|
box = pdf.read_indirect(reference)[b"MediaBox"]
|
|
self.assertAlmostEqual(box[2], 1240 * 72 / 150, delta=.1)
|
|
self.assertAlmostEqual(box[3], 1754 * 72 / 150, delta=.1)
|
|
text = subprocess.check_output(["pdftotext", str(output), "-"], text=True)
|
|
extracted = text.split("\f")
|
|
self.assertIn("12345", extracted[0])
|
|
self.assertEqual(extracted[1].strip(), "")
|
|
self.assertIn("24680", extracted[2])
|
|
self.assertTrue(all(page.exists() for page in pages))
|
|
|
|
def test_ocr_failure_never_publishes_a_partial_document_or_removes_raw_pages(self):
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
root = Path(directory)
|
|
page = root / "raw.png"
|
|
Image.new("RGB", (400, 600), "white").save(page)
|
|
output = root / "failed.pdf"
|
|
with patch("scan_document.check_ocr"), patch("scan_document.run_ocr_command", side_effect=RuntimeError("OCR failed")):
|
|
with self.assertRaisesRegex(RuntimeError, "OCR failed"):
|
|
create_document([page], output, 300, ocr=True)
|
|
self.assertFalse(output.exists())
|
|
self.assertFalse(output.with_name("failed.pdf.partial").exists())
|
|
self.assertTrue(page.exists())
|
|
self.assertEqual(list(root.glob("prepared-*")), [])
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|