Update Modal step operator

tests/integration/functional/cli/test_model.py

@@ -16,24 +16,24 @@
 from uuid import uuid4
 
 import pytest
-from click.testing import CliRunner
 
 from tests.integration.functional.cli.conftest import NAME, PREFIX
+from tests.integration.functional.cli.utils import cli_runner
 from zenml.cli.cli import cli
 from zenml.client import Client
 
 
 def test_model_list(clean_client_with_models: "Client"):
     """Test that zenml model list does not fail."""
-    runner = CliRunner(mix_stderr=False)
+    runner = cli_runner()
     list_command = cli.commands["model"].commands["list"]
     result = runner.invoke(list_command)
     assert result.exit_code == 0, result.stderr
 
 
 def test_model_create_short_names(clean_client_with_models: "Client"):
     """Test that zenml model create does not fail with short names."""
-    runner = CliRunner(mix_stderr=False)
+    runner = cli_runner()
     create_command = cli.commands["model"].commands["register"]
     model_name = PREFIX + str(uuid4())
     result = runner.invoke(
@@ -82,7 +82,7 @@ def test_model_create_short_names(clean_client_with_models: "Client"):
 
 def test_model_create_full_names(clean_client_with_models: "Client"):
     """Test that zenml model create does not fail with full names."""
-    runner = CliRunner(mix_stderr=False)
+    runner = cli_runner()
     create_command = cli.commands["model"].commands["register"]
     model_name = PREFIX + str(uuid4())
     result = runner.invoke(
@@ -131,7 +131,7 @@ def test_model_create_full_names(clean_client_with_models: "Client"):
 
 def test_model_create_only_required(clean_client_with_models: "Client"):
     """Test that zenml model create does not fail."""
-    runner = CliRunner(mix_stderr=False)
+    runner = cli_runner()
     create_command = cli.commands["model"].commands["register"]
     model_name = PREFIX + str(uuid4())
     result = runner.invoke(
@@ -155,7 +155,7 @@ def test_model_create_only_required(clean_client_with_models: "Client"):
 
 def test_model_update(clean_client_with_models: "Client"):
     """Test that zenml model update does not fail."""
-    runner = CliRunner(mix_stderr=False)
+    runner = cli_runner()
     update_command = cli.commands["model"].commands["update"]
     result = runner.invoke(
         update_command,
@@ -185,7 +185,7 @@ def test_model_create_without_required_fails(
     clean_client_with_models: "Client",
 ):
     """Test that zenml model create fails."""
-    runner = CliRunner(mix_stderr=False)
+    runner = cli_runner()
     create_command = cli.commands["model"].commands["register"]
     result = runner.invoke(
         create_command,
@@ -195,7 +195,7 @@ def test_model_create_without_required_fails(
 
 def test_model_delete_found(clean_client_with_models: "Client"):
     """Test that zenml model delete does not fail."""
-    runner = CliRunner(mix_stderr=False)
+    runner = cli_runner()
     name = PREFIX + str(uuid4())
     create_command = cli.commands["model"].commands["register"]
     runner.invoke(
@@ -212,7 +212,7 @@ def test_model_delete_found(clean_client_with_models: "Client"):
 
 def test_model_delete_not_found(clean_client_with_models: "Client"):
     """Test that zenml model delete fail."""
-    runner = CliRunner(mix_stderr=False)
+    runner = cli_runner()
     name = PREFIX + str(uuid4())
     delete_command = cli.commands["model"].commands["delete"]
     result = runner.invoke(
@@ -224,15 +224,15 @@ def test_model_delete_not_found(clean_client_with_models: "Client"):
 
 def test_model_version_list(clean_client_with_models: "Client"):
     """Test that zenml model version list does not fail."""
-    runner = CliRunner(mix_stderr=False)
+    runner = cli_runner()
     list_command = cli.commands["model"].commands["version"].commands["list"]
     result = runner.invoke(list_command, args=[f"--model={NAME}"])
     assert result.exit_code == 0, result.stderr
 
 
 def test_model_version_delete_found(clean_client_with_models: "Client"):
     """Test that zenml model version delete does not fail."""
-    runner = CliRunner(mix_stderr=False)
+    runner = cli_runner()
     model_name = PREFIX + str(uuid4())
     model_version_name = PREFIX + str(uuid4())
     model = clean_client_with_models.create_model(
@@ -254,7 +254,7 @@ def test_model_version_delete_found(clean_client_with_models: "Client"):
 
 def test_model_version_delete_not_found(clean_client_with_models: "Client"):
     """Test that zenml model version delete fail."""
-    runner = CliRunner(mix_stderr=False)
+    runner = cli_runner()
     model_name = PREFIX + str(uuid4())
     model_version_name = PREFIX + str(uuid4())
     clean_client_with_models.create_model(
@@ -278,7 +278,7 @@ def test_model_version_links_list(
     command: str, clean_client_with_models: "Client"
 ):
     """Test that zenml model version artifacts list fails."""
-    runner = CliRunner(mix_stderr=False)
+    runner = cli_runner()
     list_command = cli.commands["model"].commands[command]
     result = runner.invoke(
         list_command,
@@ -289,7 +289,7 @@ def test_model_version_links_list(
 
 def test_model_version_update(clean_client_with_models: "Client"):
     """Test that zenml model version stage update pass."""
-    runner = CliRunner(mix_stderr=False)
+    runner = cli_runner()
     update_command = (
         cli.commands["model"].commands["version"].commands["update"]
     )

tests/integration/functional/cli/test_tag.py

@@ -14,9 +14,9 @@
 """Test zenml tag CLI commands."""
 
 import pytest
-from click.testing import CliRunner
 
 from tests.integration.functional.cli.utils import (
+    cli_runner,
     random_resource_name,
     tags_killer,
 )
@@ -29,7 +29,7 @@
 def test_tag_list():
     """Test that zenml tag list does not fail."""
     with tags_killer():
-        runner = CliRunner(mix_stderr=False)
+        runner = cli_runner()
         list_command = cli.commands["tag"].commands["list"]
         result = runner.invoke(list_command)
         assert result.exit_code == 0, result.stderr
@@ -38,7 +38,7 @@ def test_tag_list():
 def test_tag_create_short_names():
     """Test that zenml tag create does not fail with short names."""
     with tags_killer(0):
-        runner = CliRunner(mix_stderr=False)
+        runner = cli_runner()
         create_command = cli.commands["tag"].commands["register"]
         tag_name = random_resource_name()
         result = runner.invoke(
@@ -55,7 +55,7 @@ def test_tag_create_short_names():
 def test_tag_create_full_names():
     """Test that zenml tag create does not fail with full names."""
     with tags_killer(0):
-        runner = CliRunner(mix_stderr=False)
+        runner = cli_runner()
         create_command = cli.commands["tag"].commands["register"]
         tag_name = random_resource_name()
         result = runner.invoke(
@@ -72,7 +72,7 @@ def test_tag_create_full_names():
 def test_tag_create_only_required():
     """Test that zenml tag create does not fail."""
     with tags_killer(0):
-        runner = CliRunner(mix_stderr=False)
+        runner = cli_runner()
         create_command = cli.commands["tag"].commands["register"]
         tag_name = random_resource_name()
         result = runner.invoke(
@@ -93,7 +93,7 @@ def test_tag_update():
     """Test that zenml tag update does not fail."""
     with tags_killer(1) as tags:
         tag: TagResponse = tags[0]
-        runner = CliRunner(mix_stderr=False)
+        runner = cli_runner()
         update_command = cli.commands["tag"].commands["update"]
         color_to_set = "yellow" if tag.color.value != "yellow" else "grey"
         result = runner.invoke(
@@ -138,7 +138,7 @@ def test_tag_update():
 def test_tag_create_without_required_fails():
     """Test that zenml tag create fails."""
     with tags_killer(0):
-        runner = CliRunner(mix_stderr=False)
+        runner = cli_runner()
         create_command = cli.commands["tag"].commands["register"]
         result = runner.invoke(
             create_command,
@@ -150,7 +150,7 @@ def test_tag_delete_found():
     """Test that zenml tag delete does not fail."""
     with tags_killer(1) as tags:
         tag: TagResponse = tags[0]
-        runner = CliRunner(mix_stderr=False)
+        runner = cli_runner()
         delete_command = cli.commands["tag"].commands["delete"]
         result = runner.invoke(
             delete_command,
@@ -165,7 +165,7 @@ def test_tag_delete_found():
 def test_tag_delete_not_found():
     """Test that zenml tag delete fail."""
     with tags_killer(0):
-        runner = CliRunner(mix_stderr=False)
+        runner = cli_runner()
         delete_command = cli.commands["tag"].commands["delete"]
         result = runner.invoke(
             delete_command,

tests/integration/functional/cli/utils.py

@@ -11,9 +11,12 @@
 #  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
 #  or implied. See the License for the specific language governing
 #  permissions and limitations under the License.
+import inspect
 from contextlib import contextmanager
 from typing import Generator, Optional, Tuple
 
+from click.testing import CliRunner
+
 from tests.harness.harness import TestHarness
 from zenml.cli import cli
 from zenml.cli.utils import (
@@ -35,6 +38,16 @@
 ]
 
 
+def cli_runner(**kwargs) -> CliRunner:
+    """Return a Click runner that stays compatible across Click releases."""
+    runner_kwargs = dict(kwargs)
+    if "mix_stderr" not in runner_kwargs:
+        params = inspect.signature(CliRunner.__init__).parameters
+        if "mix_stderr" in params:
+            runner_kwargs["mix_stderr"] = False
+    return CliRunner(**runner_kwargs)
+
+
 # ----- #
 # USERS #
 # ----- #