diff --git a/test/unit_tests/conftest.py b/test/unit_tests/conftest.py index d9a702c..611cc3f 100644 --- a/test/unit_tests/conftest.py +++ b/test/unit_tests/conftest.py @@ -1,5 +1,6 @@ import random import tempfile +from itertools import starmap from operator import itemgetter import fpdf @@ -12,6 +13,7 @@ from image_prediction.classifier.image_classifier import ImageClassifier from image_prediction.estimator.adapter.adapters.keras import KerasEstimatorAdapter from image_prediction.estimator.adapter.adapters.mock import EstimatorMock, EstimatorAdapterMock from image_prediction.exceptions import UnknownEstimatorAdapter, UnknownImageExtractor +from image_prediction.image_extractor.extractor import ImageMetadataPair from image_prediction.image_extractor.extractors.mock import ImageExtractorMock from image_prediction.image_extractor.extractors.parsable import ParsablePDFImageExtractor @@ -161,32 +163,41 @@ def metadata(images): @pytest.fixture -def pdf(images, metadata): +def image_metadata_pairs(images, metadata): + return list(starmap(ImageMetadataPair, zip(images, metadata))) - def add_image(image, metadata): - def fewer_pages_then_required(): - return metadata["page_idx"] > pdf.page - 1 - - def add_image_to_last_page(): - x, y, w, h = itemgetter("x1", "y1", "width", "height")(metadata) - - with tempfile.NamedTemporaryFile(suffix=".png") as temp_image: - image.save(temp_image.name) - pdf.image(temp_image.name, x=x, y=y, w=w, h=h) - - while fewer_pages_then_required(): - pdf.add_page() - - add_image_to_last_page() - - def pdf_object_to_actual_pdf(): - return pdf.output(dest="S").encode("latin1") +@pytest.fixture +def pdf(image_metadata_pairs): pdf = fpdf.FPDF(unit="pt") - pdf.add_page() - for image, metadata in zip(images, metadata): - add_image(image, metadata) + for pair in image_metadata_pairs: + add_image(pdf, pair) - return pdf_object_to_actual_pdf() + return pdf_object_to_actual_pdf(pdf) + + +def add_image(pdf, image_metadata_pair): + + while fewer_pages_then_required(image_metadata_pair.metadata["page_idx"], pdf): + pdf.add_page() + + add_image_to_last_page(pdf, image_metadata_pair) + + +def fewer_pages_then_required(page_idx, pdf): + return page_idx > pdf.page - 1 + + +def pdf_object_to_actual_pdf(pdf): + return pdf.output(dest="S").encode("latin1") + + +def add_image_to_last_page(pdf, image_metadata_pair): + image, metadata = image_metadata_pair + x, y, w, h = itemgetter("x1", "y1", "width", "height")(metadata) + + with tempfile.NamedTemporaryFile(suffix=".png") as temp_image: + image.save(temp_image.name) + pdf.image(temp_image.name, x=x, y=y, w=w, h=h)