import pytest from image_prediction.exceptions import UnknownImageExtractor from image_prediction.image_extractor.extractors.mock import ImageExtractorMock from image_prediction.image_extractor.extractors.parsable import ParsablePDFImageExtractor @pytest.fixture def image_extractor(extractor_type): if extractor_type == "mock": return ImageExtractorMock() elif extractor_type == "parsable_pdf": return ParsablePDFImageExtractor() elif extractor_type == "default": return None else: raise UnknownImageExtractor(f"No image extractor for type {extractor_type} was specified.")