From a9ef01a6f9a71fda5cd0199c3f8600e3c89e3173 Mon Sep 17 00:00:00 2001
From: Andrew Tolopko <atolopko-czi@users.noreply.github.com>
Date: Fri, 29 Jul 2022 11:17:04 -0400
Subject: [PATCH] fix test (#2549)

address issues building mlflow model in GHA test env
---
 test/unit/cli/fixtures/__init__.py            |  5 +++++
 .../{ => fixtures}/mlflow_model_fixture.py    | 11 +----------
 test/unit/cli/test_annotate.py                | 19 +++++++++++++++----
 3 files changed, 21 insertions(+), 14 deletions(-)
 create mode 100644 test/unit/cli/fixtures/__init__.py
 rename test/unit/cli/{ => fixtures}/mlflow_model_fixture.py (55%)

diff --git a/test/unit/cli/fixtures/__init__.py b/test/unit/cli/fixtures/__init__.py
new file mode 100644
index 00000000..c2ee749b
--- /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 5a67f0a9..f481923c 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 17f5127b..c8fe1b4b 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 @@ class TestCliAnnotate(unittest.TestCase):
         """
 
         _, query_dataset_file_path = mkstemp()
-        model = FakeModel()
-        model_file_path = write_model(model)
+        model_file_path = write_model(FakeModel())
 
         result = CliRunner().invoke(
             annotate,