15 lines
599 B
Python
15 lines
599 B
Python
from operator import itemgetter
|
|
|
|
import pytest
|
|
|
|
from image_prediction.extractor_classifier.extractor_classifier import ExtractorClassifier
|
|
|
|
|
|
@pytest.mark.parametrize("extractor_type", ["mock"])
|
|
@pytest.mark.parametrize("estimator_type", ["mock", "keras"])
|
|
def test_extractor_classifier(image_extractor, image_classifier, images, batch_of_expected_string_labels):
|
|
extractor_classifier = ExtractorClassifier(image_extractor, image_classifier)
|
|
results = extractor_classifier(images)
|
|
labels = list(map(itemgetter("prediction"), results))
|
|
assert labels == batch_of_expected_string_labels
|