image-classification-service/test/fixtures/extractor.py

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.")