fix test

Signed-off-by: Ann Zhang <ann.zhang@databricks.com>
mlflow · Apr 30, 2024 · c87ed7d · c87ed7d
1 parent 53f2d79
commit c87ed7d
Showing 1 changed file with 5 additions and 13 deletions.
diff --git a/tests/models/test_utils.py b/tests/models/test_utils.py
@@ -433,19 +433,11 @@ def test_enforce_array_with_errors():
 
 
 def test_model_code_validation():
-    invalid_code = """
-    dbutils.library.restartPython()
-    some_python_variable = 5
-    """
-
-    warning_code = """
-    # dbutils.library.restartPython()
-    # MAGIC %run ../wheel_installer
-    """
-
-    valid_code = """
-    some_valid_python_code = "valid"
-    """
+    invalid_code = "dbutils.library.restartPython()\nsome_python_variable = 5"
+
+    warning_code = "# dbutils.library.restartPython()\n# MAGIC %run ../wheel_installer"
+
+    valid_code = "some_valid_python_code = 'valid'"
 
     with pytest.raises(
         ValueError, match="The model file uses 'dbutils' command which is not supported."