fix for test?

Author: ianardee

mindee/image/extracted_images.py

@@ -1,5 +1,12 @@
+from pathlib import Path
+
 from mindee.image.extracted_image import ExtractedImage
 
 
 class ExtractedImages(list[ExtractedImage]):
     """List of extracted images."""
+
+    def save_all_to_disk(self, output_dir: Path | str) -> None:
+        """Save all extracted images to disk."""
+        for image in self:
+            image.save_to_file(output_dir)

tests/v2/file_operations/test_split_operation_integration.py

@@ -13,19 +13,18 @@
 from tests.utils import OUTPUT_DIR, V2_PRODUCT_DATA_DIR, cleanup_output_files
 
 
-@pytest.fixture
-def invoice_splitter_5p_path():
-    return V2_PRODUCT_DATA_DIR / "split" / "invoice_5p.pdf"
-
-
 def check_findoc_return(findoc_response: ExtractionResponse):
     assert len(findoc_response.inference.model.id) > 0
     assert findoc_response.inference.result.fields.get("total_amount").value > 0
 
 
+output_files = ["default_sample_pages-001-001.pdf", "default_sample_pages-002-002.pdf"]
+
+
 @pytest.mark.pypdfium2
 @pytest.mark.integration
 def test_pdf_should_extract_splits():
+
     client = Client()
     split_input = PathInput(V2_PRODUCT_DATA_DIR / "split" / "default_sample.pdf")
     response = client.enqueue_and_get_result(
@@ -41,8 +40,8 @@ def test_pdf_should_extract_splits():
     extracted_pdfs = response.inference.result.extract_from_input_source(split_input)
 
     assert len(extracted_pdfs) == 2
-    assert extracted_pdfs[0].filename == "default_sample_001-001.pdf"
-    assert extracted_pdfs[1].filename == "default_sample_002-002.pdf"
+    assert extracted_pdfs[0].filename == output_files[0]
+    assert extracted_pdfs[1].filename == output_files[1]
 
     invoice_0 = client.enqueue_and_get_result(
         ExtractionResponse,
@@ -54,7 +53,7 @@ def test_pdf_should_extract_splits():
     check_findoc_return(invoice_0)
     extracted_pdfs.save_all_to_disk(OUTPUT_DIR)
     for i in range(len(extracted_pdfs)):
-        local_input = PathInput(OUTPUT_DIR / f"split_{i + 1:03d}.pdf")
+        local_input = PathInput(OUTPUT_DIR / output_files[i])
         try:
             assert local_input.page_count == extracted_pdfs[i].get_page_count()
         finally:
@@ -65,4 +64,4 @@ def test_pdf_should_extract_splits():
 @pytest.fixture(scope="module", autouse=True)
 def cleanup():
     yield
-    cleanup_output_files(["split_001.pdf", "split_002.pdf"])
+    cleanup_output_files(output_files)