13 lines
498 B
Python

import pytest
from image_prediction.model_loading import load_model_and_classes
@pytest.mark.parametrize("loader_type", ["mock"])
@pytest.mark.parametrize("estimator_type", ["mock"])
@pytest.mark.parametrize("batch_size", [3])
def test_load_model_and_classes(model_loader, estimator_adapter, classes):
model_loaded, classes_loaded = load_model_and_classes("some random identifier", model_loader=model_loader)
assert model_loaded == estimator_adapter
assert classes_loaded == classes