fix test case

sassoftware · Sep 9, 2021 · edeff5c · edeff5c
1 parent 28a1b70
commit edeff5c
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/src/sasctl/utils/models.py b/src/sasctl/utils/models.py
@@ -105,6 +105,7 @@ def __init__(self, model):
         self.properties = []
         self.sample_input = None
         self.sample_output = None
+        self.tool = None
 
         # Most models take DataFrame or array as input instead of individual variables.
         self.array_input = True

diff --git a/tests/unit/test_pymas.py b/tests/unit/test_pymas.py
@@ -721,7 +721,7 @@ def predict_proba(a, b):
 def test_from_model_info():
     from sasctl.utils.models import ModelInfo
     from sasctl.utils.pymas import from_model_info
-
+    from sasctl.utils.pymas.core import PyMAS2
     # Dummy model info
     info = ModelInfo(None)
 
@@ -737,5 +737,5 @@ def test_from_model_info():
 
     result = from_model_info(info)
 
-    assert isinstance(result, PyMAS)
+    assert isinstance(result, PyMAS2)