fix parameter decorator

Author: TumanyanV

superannotate/parameter_decorators.py

@@ -68,14 +68,12 @@ def inner(*args, **kwargs):
             model_keyword = 'model'
         else:
             model_keyword = 'base_model'
-
-        new_model_arg = None
+        new_model_arg = new_kwargs[model_keyword]
         if isinstance(new_kwargs[model_keyword], str):
 
             all_models = search_models(
                 include_global=True, name=new_kwargs[model_keyword]
             )
-
             if len(all_models) > 2:
                 raise SABaseException(
                     0,

tests/test_ml_funcs.py

@@ -19,7 +19,7 @@
 
 
 @pytest.mark.skipif(
-    "SA_STRESS_TESTS" not in os.environ,
+    "SA_ML_FUNCTIONALITY" not in os.environ,
     reason="Requires env variable to be set"
 )
 def test_run_prediction():
@@ -77,7 +77,7 @@ def test_run_prediction():
 
 
 @pytest.mark.skipif(
-    "SA_STRESS_TESTS" not in os.environ,
+    "SA_ML_FUNCTIONALITY" not in os.environ,
     reason="Requires env variable to be set"
 )
 def test_run_segmentation():
@@ -127,5 +127,7 @@ def test_download_model(tmpdir):
 
     ml_model = sa.search_models(include_global=True)[0]
     Path.rmdir(Path(tmpdir) / './export')
-    assert (sa.download_model(ml_model, './export') == True)
-    assert (sa.download_model(ml_model['name'], './export') == True)
+    model = sa.download_model(ml_model, './export')
+    assert model['name']
+    model = sa.download_model(ml_model['name'], './export')
+    assert model['name']