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