18 lines
626 B
Python
18 lines
626 B
Python
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.")
|