diff --git a/test/unit/cli/fixtures/__init__.py b/test/unit/cli/fixtures/__init__.py new file mode 100644 index 000000000..c2ee749b2 --- /dev/null +++ b/test/unit/cli/fixtures/__init__.py @@ -0,0 +1,5 @@ +from .mlflow_model_fixture import FakeModel + + +def _load_pyfunc(data_path): + return FakeModel() diff --git a/test/unit/cli/mlflow_model_fixture.py b/test/unit/cli/fixtures/mlflow_model_fixture.py similarity index 55% rename from test/unit/cli/mlflow_model_fixture.py rename to test/unit/cli/fixtures/mlflow_model_fixture.py index 5a67f0a9f..f481923c2 100644 --- a/test/unit/cli/mlflow_model_fixture.py +++ b/test/unit/cli/fixtures/mlflow_model_fixture.py @@ -1,20 +1,11 @@ -import shutil -from tempfile import TemporaryDirectory, mkstemp - import mlflow -def write_model(model) -> str: - with TemporaryDirectory() as mlflow_model_dir: - mlflow.pyfunc.save_model(mlflow_model_dir, python_model=model) - return shutil.make_archive(mkstemp()[1], "zip", mlflow_model_dir) - - class FakeModel(mlflow.pyfunc.PythonModel): def __init__(self, input_to_output: dict = {}): self.input_to_output = input_to_output - def predict(self, context, model_input) -> None: + def predict(self, model_input) -> None: # this stdout output is useful for validating the input in a test, noting that this model will be invoked in a # subprocess, so stdout is one means of communicating information back to the test code print(f"__MODEL_INPUT__={model_input.iloc[0][0]}") diff --git a/test/unit/cli/test_annotate.py b/test/unit/cli/test_annotate.py index 17f5127bb..c8fe1b4bc 100644 --- a/test/unit/cli/test_annotate.py +++ b/test/unit/cli/test_annotate.py @@ -1,10 +1,22 @@ +import os +import shutil import unittest -from tempfile import mkstemp +from tempfile import mkstemp, TemporaryDirectory +import mlflow from click.testing import CliRunner from server.cli.annotate import annotate -from test.unit.cli.mlflow_model_fixture import FakeModel, write_model +from test.unit.cli.fixtures.mlflow_model_fixture import FakeModel + + +def write_model(model) -> str: + with TemporaryDirectory() as mlflow_model_dir: + fixtures_path = os.path.join(os.path.dirname(__file__), 'fixtures') + mlflow.pyfunc.save_model(mlflow_model_dir, + loader_module='fixtures', + code_path=[fixtures_path]) + return shutil.make_archive(mkstemp()[1], "zip", mlflow_model_dir) class TestCliAnnotate(unittest.TestCase): @@ -29,8 +41,7 @@ def test__annotate__loads_and_runs(self): """ _, query_dataset_file_path = mkstemp() - model = FakeModel() - model_file_path = write_model(model) + model_file_path = write_model(FakeModel()) result = CliRunner().invoke( annotate,