diff --git a/test/conftest.py b/test/conftest.py index 87b8185..ee2b3d5 100644 --- a/test/conftest.py +++ b/test/conftest.py @@ -6,6 +6,7 @@ from image_prediction.utils import get_logger pytest_plugins = [ + "test.fixtures.extractor", "test.fixtures.image", "test.fixtures.image_metadata_pair", "test.fixtures.input", diff --git a/test/fixtures/extractor.py b/test/fixtures/extractor.py new file mode 100644 index 0000000..d51e1bd --- /dev/null +++ b/test/fixtures/extractor.py @@ -0,0 +1,17 @@ +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.") \ No newline at end of file diff --git a/test/fixtures/model.py b/test/fixtures/model.py index 57968a3..729d234 100644 --- a/test/fixtures/model.py +++ b/test/fixtures/model.py @@ -4,9 +4,7 @@ from image_prediction.classifier.classifier import Classifier from image_prediction.classifier.image_classifier import ImageClassifier from image_prediction.estimator.adapter.adapter import EstimatorAdapter from image_prediction.estimator.preprocessor.preprocessors.basic import BasicPreprocessor -from image_prediction.exceptions import UnknownImageExtractor, UnknownEstimatorAdapter -from image_prediction.image_extractor.extractors.mock import ImageExtractorMock -from image_prediction.image_extractor.extractors.parsable import ParsablePDFImageExtractor +from image_prediction.exceptions import UnknownEstimatorAdapter from image_prediction.redai_adapter.model import PredictionModelHandle @@ -27,18 +25,6 @@ def estimator_mock(): return EstimatorMock() -@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.") - - @pytest.fixture def image_classifier(classifier, monkeypatch, batch_of_expected_string_labels): return ImageClassifier(classifier, preprocessor=BasicPreprocessor())