"""
test_model.py
The MLOps framework needs `ServiceTest` to test services locally and make
test calls after deploying services.
It's highly recommended to define tests for training, predictions and data
validation, but it's not required by the framework.
"""
import os
from typing import Any, List

import pytest

from model_code.model import ModelIO, initialization, predict


class ServiceTest:
    """
    MLOps framework uses this object to test services
    Training-related methods and attributes are not required by the framework
    if there's no training service.
    """

    def __init__(self):
        """
        Define known inputs: `training_input` and `prediction_input`.
        Note that the values need to be customized to your model.
        """
        self.prediction_input = {
            "input": [
                {
                    "well": "dummy",
                    "input_logs": {
                        "ACS": [-1],
                        "RDEP": [1],
                        "DEN": [1],
                    },
                }
            ],
            "return_file": False,
        }

    def prediction_check(self, y: List):
        """
        Evaluate the prediction service's output when the input is
        `self.prediction_input`
        Input:
        - y: output generated by `predict`
        Output:
        - `True` if `y` is correct, `False` otherwise.
        """
        assert y == [1, 1, 1]
        return True


################################################################################
# The rest of the file are tests for the example model. It's recommended that
# you include tests to check your code. They aren't required by the framework,
# but if they're present they will be run by pytest during deployment
################################################################################

# We will reuse the inputs and checks in the unit tests. Not required though.
service_test = ServiceTest()  # type: ignore


def _check_model(model: Any) -> None:
    x = service_test.prediction_input
    y = predict(x, model)
    assert service_test.prediction_check(y)


def test_model_loading_error():
    with pytest.raises(ModelIO):
        initialization(folder_path="path_doesnt_exist")


def test_predict():
    model = initialization("model_artifact")  # model_storage[0]
    _check_model(model)


def test_predict_local_model():
    """
    Test prediction service using a model artifact to load the model.
    This will fail in the training service, since model artifacts are not
    deployed in this case. It's important to skip this test as shown below:
    """
    if os.getenv("SERVICE_NAME") == "training":
        pytest.skip("This test is specific to prediction services")
    model = initialization("model_artifact")  # Same as in settings file!
    _check_model(model)
