53 lines
1.4 KiB
Python
53 lines
1.4 KiB
Python
import logging
|
|
|
|
import numpy as np
|
|
import pytest
|
|
|
|
from image_prediction.estimator.mock import EstimatorMock
|
|
from image_prediction.service_estimator.mock import ServiceEstimatorMock
|
|
from image_prediction.utils import get_logger
|
|
|
|
logger = get_logger()
|
|
logger.setLevel(logging.DEBUG)
|
|
|
|
|
|
@pytest.fixture(scope="session")
|
|
def estimator():
|
|
return EstimatorMock()
|
|
|
|
|
|
@pytest.fixture(scope="session")
|
|
def batches(batch_size, classes):
|
|
input_batch = np.random.normal(size=(batch_size, 10, 15))
|
|
output_batch = np.random.randint(low=0, high=len(classes), size=batch_size)
|
|
return input_batch, output_batch
|
|
|
|
|
|
@pytest.fixture(scope="session")
|
|
def classes():
|
|
return ["A", "B", "C"]
|
|
|
|
|
|
def map_labels(numeric_labels, classes):
|
|
return [classes[nl] for nl in numeric_labels]
|
|
|
|
|
|
@pytest.fixture(scope="session")
|
|
def service_estimator(model_type, estimator, classes):
|
|
if model_type == "mock":
|
|
return ServiceEstimatorMock(estimator, classes)
|
|
|
|
|
|
@pytest.mark.parametrize("model_type", ["mock"], scope="session")
|
|
@pytest.mark.parametrize("batch_size", [0, 1, 2, 16, 32, 64], scope="session")
|
|
def test_predict(service_estimator, batches, classes):
|
|
|
|
input_batch, output_batch = batches
|
|
expected_predictions = map_labels(output_batch, classes)
|
|
|
|
service_estimator.estimator.output_batch = output_batch
|
|
|
|
predictions = service_estimator.predict(input_batch)
|
|
|
|
assert predictions == expected_predictions
|