fix test (#2549)

address issues building mlflow model in GHA test env

test/unit/cli/fixtures/__init__.py

@@ -0,0 +1,5 @@
+from .mlflow_model_fixture import FakeModel
+
+
+def _load_pyfunc(data_path):
+    return FakeModel()

test/unit/cli/mlflow_model_fixture.py → 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]}")

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,