115 lines
4.1 KiB
Python
115 lines
4.1 KiB
Python
import tempfile
|
|
import unittest
|
|
from pathlib import Path
|
|
|
|
from lernotafiscal import db
|
|
from lernotafiscal.extraction import detect_documents, extract_pdf_text_naive, segment_blocks
|
|
from lernotafiscal.extraction import PageSource
|
|
|
|
|
|
SINGLE_DOC = """
|
|
MERCADO CENTRAL LTDA
|
|
CNPJ 12.345.678/0001-90
|
|
DATA 13/07/2026
|
|
ITEM ARROZ 10,00
|
|
VALOR TOTAL R$ 42,50
|
|
"""
|
|
|
|
|
|
MULTI_DOC = """
|
|
MERCADO CENTRAL LTDA
|
|
CNPJ 12.345.678/0001-90
|
|
DATA 13/07/2026
|
|
VALOR TOTAL R$ 42,50
|
|
|
|
FARMACIA SAUDE
|
|
CNPJ 98.765.432/0001-10
|
|
DATA 14/07/2026
|
|
TOTAL A PAGAR R$ 18,90
|
|
"""
|
|
|
|
|
|
class ExtractionTests(unittest.TestCase):
|
|
def test_single_image_text_creates_one_candidate(self) -> None:
|
|
pages = [PageSource(1, SINGLE_DOC, None, "image")]
|
|
docs = detect_documents(pages, "cupom.png")
|
|
|
|
self.assertEqual(len(docs), 1)
|
|
self.assertEqual(docs[0].purchase_date, "2026-07-13")
|
|
self.assertEqual(docs[0].supplier_name, "MERCADO CENTRAL LTDA")
|
|
self.assertEqual(docs[0].total_paid, 42.50)
|
|
|
|
def test_pdf_page_with_multiple_blocks_creates_multiple_candidates(self) -> None:
|
|
blocks = segment_blocks(MULTI_DOC)
|
|
self.assertEqual(len(blocks), 2)
|
|
|
|
docs = detect_documents([PageSource(1, MULTI_DOC, None, "pdf")], "lote.pdf")
|
|
self.assertEqual(len(docs), 2)
|
|
self.assertEqual(docs[0].supplier_name, "MERCADO CENTRAL LTDA")
|
|
self.assertEqual(docs[1].supplier_name, "FARMACIA SAUDE")
|
|
|
|
def test_pdf_with_one_document_per_page_creates_review_item_per_page(self) -> None:
|
|
pages = [
|
|
PageSource(1, SINGLE_DOC, None, "pdf"),
|
|
PageSource(2, SINGLE_DOC.replace("MERCADO CENTRAL LTDA", "POSTO AVENIDA"), None, "pdf"),
|
|
]
|
|
docs = detect_documents(pages, "duas-paginas.pdf")
|
|
|
|
self.assertEqual(len(docs), 2)
|
|
self.assertEqual(docs[0].source_location, "pagina 1, bloco 1")
|
|
self.assertEqual(docs[1].source_location, "pagina 2, bloco 1")
|
|
self.assertEqual(docs[1].supplier_name, "POSTO AVENIDA")
|
|
|
|
def test_naive_pdf_extraction_rejects_binary_streams_without_fake_pages(self) -> None:
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
pdf_path = Path(tmp) / "scan.pdf"
|
|
pdf_path.write_bytes(
|
|
b"%PDF-1.4\n"
|
|
b"1 0 obj << /Type /Page >> endobj\n"
|
|
b"2 0 obj << /Length 28 >> stream\n"
|
|
b"\x01\x02binary\x0cnoise\x0cnot text\xff\n"
|
|
b"endstream\n%%EOF"
|
|
)
|
|
|
|
pages = extract_pdf_text_naive(pdf_path)
|
|
|
|
self.assertEqual(pages, [""])
|
|
|
|
def test_low_confidence_candidate_can_be_corrected_and_confirmed(self) -> None:
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
conn = db.connect(Path(tmp) / "test.sqlite3")
|
|
try:
|
|
db.init_db(conn)
|
|
upload_id = db.insert_upload(conn, "scan.png", Path(tmp) / "scan.png", "image/png", 12)
|
|
candidate = detect_documents([PageSource(1, "", None, "image")], "scan.png")[0]
|
|
detected_id = db.insert_detected_document(conn, upload_id, candidate)
|
|
|
|
db.confirm_document(conn, detected_id, "2026-07-15", "PADARIA BOA", 9.75)
|
|
docs = db.dashboard_documents(conn)
|
|
|
|
self.assertEqual(len(docs), 1)
|
|
self.assertEqual(docs[0]["supplier_name"], "PADARIA BOA")
|
|
self.assertEqual(docs[0]["total_paid"], 9.75)
|
|
finally:
|
|
conn.close()
|
|
|
|
def test_ignored_documents_do_not_reach_dashboard(self) -> None:
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
conn = db.connect(Path(tmp) / "test.sqlite3")
|
|
try:
|
|
db.init_db(conn)
|
|
upload_id = db.insert_upload(conn, "cupom.png", Path(tmp) / "cupom.png", "image/png", 12)
|
|
candidate = detect_documents([PageSource(1, SINGLE_DOC, None, "image")], "cupom.png")[0]
|
|
detected_id = db.insert_detected_document(conn, upload_id, candidate)
|
|
|
|
db.ignore_document(conn, detected_id)
|
|
docs = db.dashboard_documents(conn)
|
|
|
|
self.assertEqual(docs, [])
|
|
finally:
|
|
conn.close()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|