fix: test cases

Signed-off-by: Leila Wang <[email protected]>
numaproj · Sep 23, 2024 · 82ea731 · 82ea731
1 parent 8db67b1
commit 82ea731
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 6 deletions.
diff --git a/tests/registry/_mlflow_utils.py b/tests/registry/_mlflow_utils.py
@@ -385,7 +385,9 @@ def return_pyfunc_rundata():
             status="RUNNING",
             user_id="lol",
         ),
-        run_data=RunData(metrics={}, tags={}, params={}),
+        run_data=RunData(
+            metrics={}, tags={}, params=[mlflow.entities.Param("learning_rate", "0.01")]
+        ),
     )
 
 

diff --git a/tests/registry/test_mlflow_registry.py b/tests/registry/test_mlflow_registry.py
@@ -116,11 +116,11 @@ def test_load_multiple_models_when_pyfunc_model_exist(self):
         skeys = self.skeys
         dkeys_list = [["AE", "infer"], ["scaler", "infer"]]
         data = ml.load_multiple(skeys=skeys, dkeys_list=dkeys_list)
-        self.assertIsNotNone(data["AE:infer"].metadata)
-        self.assertIsNotNone(data["scaler:infer"].metadata)
-        self.assertIsInstance(data, dict)
-        self.assertIsInstance(data["AE:infer"].artifact, VanillaAE)
-        self.assertIsInstance(data["scaler:infer"].artifact, StandardScaler)
+        self.assertIsNotNone(data.metadata)
+        self.assertIsInstance(data, ArtifactData)
+        self.assertIsInstance(data.artifact, dict)
+        self.assertIsInstance(data.artifact["AE"].artifact, VanillaAE)
+        self.assertIsInstance(data.artifact["scaler"].artifact, StandardScaler)
 
     @patch("mlflow.sklearn.log_model", mock_log_model_sklearn)
     @patch("mlflow.log_param", mock_log_state_dict)