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, expected_batch_string_labels):
    extractor_classifier = ExtractorClassifier(image_extractor, image_classifier)
    results = extractor_classifier(images)
    labels = list(map(itemgetter("prediction"), results))
    assert labels == expected_batch_string_labels