image-classification-service/test/unit_tests/image_classifier_test.py

import pytest


@pytest.mark.parametrize("estimator_type", ["mock", "keras"])
def test_predict(image_classifier, images, batch_of_expected_string_labels):
    predictions = list(image_classifier.predict(images))
    assert predictions == batch_of_expected_string_labels