split input tests

Author: ianardee

tests/extraction/test_image_extractor.py

@@ -6,7 +6,7 @@
 from mindee.extraction.common.image_extractor import extract_multiple_images_from_source
 from mindee.input.sources.path_input import PathInput
 from mindee.product.barcode_reader.barcode_reader_v1 import BarcodeReaderV1
-from tests.input.test_inputs import PRODUCT_DATA_DIR
+from tests.utils import PRODUCT_DATA_DIR
 
 
 @pytest.fixture

tests/extraction/test_invoice_splitter_auto_extraction.py

@@ -8,9 +8,8 @@
 from mindee.parsing.common.document import Document
 from mindee.product.invoice.invoice_v4 import InvoiceV4
 from mindee.product.invoice_splitter.invoice_splitter_v1 import InvoiceSplitterV1
-from tests.input.test_inputs import PRODUCT_DATA_DIR
 from tests.product import get_id, get_version
-from tests.utils import levenshtein_ratio
+from tests.utils import PRODUCT_DATA_DIR, levenshtein_ratio
 
 
 @pytest.fixture

tests/extraction/test_multi_receipts_extractor.py

@@ -10,7 +10,7 @@
 from mindee.product.multi_receipts_detector.multi_receipts_detector_v1 import (
     MultiReceiptsDetectorV1,
 )
-from tests.input.test_inputs import PRODUCT_DATA_DIR
+from tests.utils import PRODUCT_DATA_DIR
 
 
 @pytest.fixture

tests/extraction/test_pdf_extractor.py

@@ -8,7 +8,7 @@
 from mindee.product.invoice_splitter.invoice_splitter_v1_document import (
     InvoiceSplitterV1Document,
 )
-from tests.input.test_inputs import PRODUCT_DATA_DIR
+from tests.utils import PRODUCT_DATA_DIR
 
 
 @pytest.fixture

tests/input/test_apply_page_options.py

@@ -0,0 +1,161 @@
+import io
+
+import pypdfium2 as pdfium
+import pytest
+
+from mindee.error import MindeeError
+from mindee.input.page_options import KEEP_ONLY, REMOVE, PageOptions
+from mindee.input.sources import (
+    Base64Input,
+    BytesInput,
+    FileInput,
+    LocalInputSource,
+    PathInput,
+)
+from tests.utils import FILE_TYPES_DIR, PRODUCT_DATA_DIR
+
+
+def _assert_page_options(input_source: LocalInputSource, numb_pages: int):
+    assert input_source.is_pdf() is True
+    # Currently the least verbose way of comparing pages with pypdfium2
+    # I.e., each page is read and rendered as a rasterized image.
+    # These images are then compared as raw byte sequences.
+    cut_pdf = pdfium.PdfDocument(input_source.file_object)
+    pdf = pdfium.PdfDocument(FILE_TYPES_DIR / "pdf" / f"multipage_cut-{numb_pages}.pdf")
+    for idx in range(len(pdf)):
+        pdf_page = pdf.get_page(idx)
+        pdf_page_render = pdfium.PdfPage.render(pdf_page)
+        cut_pdf_page = cut_pdf.get_page(idx)
+        cut_pdf_page_render = pdfium.PdfPage.render(cut_pdf_page)
+
+        assert bytes(pdf_page_render.buffer) == bytes(cut_pdf_page_render.buffer)
+    cut_pdf.close()
+    pdf.close()
+
+
+def test_pdf_reconstruct_ok():
+    input_source = PathInput(FILE_TYPES_DIR / "pdf" / "multipage.pdf")
+    input_source.process_pdf(behavior=KEEP_ONLY, on_min_pages=2, page_indexes=range(5))
+    assert isinstance(input_source.file_object, io.BytesIO)
+
+
+@pytest.mark.parametrize("numb_pages", [1, 2, 3])
+def test_process_pdf_cut_n_pages(numb_pages: int):
+    input_source = PathInput(FILE_TYPES_DIR / "pdf" / "multipage.pdf")
+    input_source.process_pdf(
+        behavior=KEEP_ONLY, on_min_pages=2, page_indexes=[0, -2, -1][:numb_pages]
+    )
+    assert input_source.page_count == numb_pages
+    _assert_page_options(input_source, numb_pages)
+
+
+@pytest.mark.parametrize("numb_pages", [1, 2, 3])
+def test_apply_pages_pdf_cut_n_pages(numb_pages: int):
+    input_source = PathInput(FILE_TYPES_DIR / "pdf" / "multipage.pdf")
+    input_source.apply_page_options(
+        PageOptions(on_min_pages=2, page_indexes=[0, -2, -1][:numb_pages])
+    )
+    assert input_source.count_doc_pages() == numb_pages
+    _assert_page_options(input_source, numb_pages)
+
+
+def test_pdf_keep_5_first_pages():
+    input_source = PathInput(FILE_TYPES_DIR / "pdf" / "multipage.pdf")
+    assert input_source.is_pdf() is True
+    input_source.process_pdf(
+        behavior=KEEP_ONLY, on_min_pages=2, page_indexes=[0, 1, 2, 3, 4]
+    )
+    assert input_source.count_doc_pages() == 5
+
+
+def test_pdf_keep_invalid_pages():
+    input_source = PathInput(FILE_TYPES_DIR / "pdf" / "multipage.pdf")
+    assert input_source.is_pdf() is True
+    input_source.process_pdf(
+        behavior=KEEP_ONLY, on_min_pages=2, page_indexes=[0, 1, 17]
+    )
+    assert input_source.count_doc_pages() == 2
+
+
+def test_pdf_remove_5_last_pages():
+    input_source = PathInput(FILE_TYPES_DIR / "pdf" / "multipage.pdf")
+    assert input_source.is_pdf() is True
+    input_source.process_pdf(
+        behavior=REMOVE, on_min_pages=2, page_indexes=[-5, -4, -3, -2, -1]
+    )
+    assert input_source.count_doc_pages() == 7
+
+
+def test_pdf_remove_5_first_pages():
+    input_source = PathInput(FILE_TYPES_DIR / "pdf" / "multipage.pdf")
+    assert input_source.is_pdf() is True
+    input_source.process_pdf(
+        behavior=REMOVE, on_min_pages=2, page_indexes=list(range(5))
+    )
+    assert input_source.count_doc_pages() == 7
+
+
+def test_pdf_remove_invalid_pages():
+    input_source = PathInput(FILE_TYPES_DIR / "pdf" / "multipage.pdf")
+    assert input_source.is_pdf() is True
+    input_source.process_pdf(behavior=REMOVE, on_min_pages=2, page_indexes=[16])
+    assert input_source.count_doc_pages() == 12
+
+
+def test_pdf_keep_no_pages():
+    input_source = PathInput(FILE_TYPES_DIR / "pdf" / "multipage.pdf")
+    assert input_source.is_pdf() is True
+    # empty page indexes
+    with pytest.raises(RuntimeError):
+        input_source.process_pdf(behavior=KEEP_ONLY, on_min_pages=2, page_indexes=[])
+    # all invalid pages
+    with pytest.raises(RuntimeError):
+        input_source.process_pdf(
+            behavior=KEEP_ONLY, on_min_pages=2, page_indexes=[16, 17]
+        )
+
+
+def test_pdf_remove_all_pages():
+    input_source = PathInput(FILE_TYPES_DIR / "pdf" / "multipage.pdf")
+    assert input_source.is_pdf() is True
+    with pytest.raises(RuntimeError):
+        input_source.process_pdf(
+            behavior=REMOVE, on_min_pages=2, page_indexes=list(range(15))
+        )
+
+
+def test_pdf_input_from_file():
+    with open(FILE_TYPES_DIR / "pdf" / "multipage.pdf", "rb") as fp:
+        input_source = FileInput(fp)
+        assert input_source.is_pdf() is True
+        input_source.process_pdf(behavior=KEEP_ONLY, on_min_pages=2, page_indexes=[0])
+    assert input_source.count_doc_pages() == 1
+
+
+def test_pdf_input_from_base64():
+    with open(PRODUCT_DATA_DIR / "invoices" / "invoice_10p.txt", "rt") as fp:
+        input_source = Base64Input(fp.read(), filename="invoice_10p.pdf")
+    assert input_source.is_pdf() is True
+    input_source.process_pdf(behavior=KEEP_ONLY, on_min_pages=2, page_indexes=[0])
+    assert input_source.count_doc_pages() == 1
+
+
+def test_pdf_input_from_bytes():
+    with open(PRODUCT_DATA_DIR / "invoices" / "invoice_10p.pdf", "rb") as fp:
+        input_source = BytesInput(fp.read(), filename="invoice_10p.pdf")
+    assert input_source.is_pdf() is True
+    input_source.process_pdf(behavior=KEEP_ONLY, on_min_pages=2, page_indexes=[0])
+    assert input_source.count_doc_pages() == 1
+
+
+def test_pdf_blank_check():
+    with pytest.raises(MindeeError):
+        input_source = PathInput(FILE_TYPES_DIR / "pdf" / "blank.pdf")
+        input_source.process_pdf(behavior=KEEP_ONLY, on_min_pages=2, page_indexes=[0])
+
+    with pytest.raises(MindeeError):
+        input_source = PathInput(FILE_TYPES_DIR / "pdf" / "blank_1.pdf")
+        input_source.process_pdf(behavior=KEEP_ONLY, on_min_pages=2, page_indexes=[0])
+
+    input_not_blank = PathInput(FILE_TYPES_DIR / "pdf" / "not_blank_image_only.pdf")
+    assert input_not_blank.count_doc_pages() == 1

tests/input/test_fix_pdf.py

@@ -0,0 +1,22 @@
+import pytest
+
+from mindee import PathInput
+from mindee.error import MimeTypeError
+from tests.utils import FILE_TYPES_DIR
+
+
+def test_broken_unfixable_pdf():
+    with pytest.raises(MimeTypeError):
+        input_source = PathInput(FILE_TYPES_DIR / "pdf" / "broken_unfixable.pdf")
+        input_source.fix_pdf()
+
+
+def test_broken_fixable_pdf():
+    input_source = PathInput(FILE_TYPES_DIR / "pdf" / "broken_fixable.pdf")
+    input_source.fix_pdf()
+    assert input_source.page_count == 1
+
+
+def test_broken_fixable_invoice_pdf():
+    input_source = PathInput(FILE_TYPES_DIR / "pdf" / "broken_invoice.pdf")
+    input_source.fix_pdf()