vellum-ai 0.14.63__py3-none-any.whl → 1.13.5__py3-none-any.whl

vellum_cli/push.py

@@ -8,15 +8,16 @@ from uuid import UUID
 from typing import List, Optional
 
 from dotenv import load_dotenv
+from pydantic import ValidationError
 
 from vellum.client.core.api_error import ApiError
+from vellum.client.types import WorkflowPushDeploymentConfigRequest
 from vellum.resources.workflows.client import OMIT
-from vellum.types import WorkflowPushDeploymentConfigRequest
 from vellum.workflows.vellum_client import create_vellum_client
-from vellum.workflows.workflows.base import BaseWorkflow
 from vellum_cli.config import DEFAULT_WORKSPACE_CONFIG, WorkflowConfig, WorkflowDeploymentConfig, load_vellum_cli_config
 from vellum_cli.logger import handle_cli_error, load_cli_logger
-from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
+from vellum_ee.workflows.display.nodes.utils import to_kebab_case
+from vellum_ee.workflows.display.workflows.base_workflow_display import BaseWorkflowDisplay
 
 
 def push_command(
@@ -27,6 +28,7 @@ def push_command(
     deployment_name: Optional[str] = None,
     deployment_description: Optional[str] = None,
     release_tags: Optional[List[str]] = None,
+    release_description: Optional[str] = None,
     dry_run: Optional[bool] = None,
     strict: Optional[bool] = None,
     workspace: Optional[str] = None,
@@ -61,22 +63,45 @@ def push_command(
             config.workflows.append(new_config)
             workflow_configs = [new_config]
         elif not module:
-            raise ValueError("No Workflows found in project to push.")
+            handle_cli_error(
+                logger,
+                title="No workflows found in project to push",
+                message="No workflow configurations were found in vellum.lock.json.",
+                suggestion="If you pulled a workflow with a specific sandbox ID"
+                "(e.g., vellum workflows pull --workflow-sandbox-id=<id>), "
+                "you should push without specifying a different sandbox ID:\n\n"
+                "  vellum workflows push\n\n"
+                "Or, if you want to push to a specific module:\n\n"
+                "  vellum workflows push --module=<module-name>",
+            )
+            return
         else:
             raise ValueError(f"No workflow config for '{module}' found in project to push.")
 
     if len(workflow_configs) > 1:
-        raise ValueError("Multiple workflows found in project to push. Pushing only a single workflow is supported.")
+        raise ValueError(
+            "Multiple workflows found. Please specify a single workflow to push. "
+            f"Found: {', '.join([w.module for w in workflow_configs])}"
+        )
 
     workflow_config = workflow_configs[0]
 
     logger.info(f"Loading workflow from {workflow_config.module}")
     resolved_workspace = workspace or workflow_config.workspace or DEFAULT_WORKSPACE_CONFIG.name
-    workspace_config = (
-        next((w for w in config.workspaces if w.name == resolved_workspace), DEFAULT_WORKSPACE_CONFIG)
-        if workspace
-        else DEFAULT_WORKSPACE_CONFIG
-    )
+    workspace_config = next((w for w in config.workspaces if w.name == resolved_workspace), None)
+
+    if workspace_config is None:
+        if resolved_workspace == DEFAULT_WORKSPACE_CONFIG.name:
+            workspace_config = DEFAULT_WORKSPACE_CONFIG
+        else:
+            available_workspaces = sorted(set([w.name for w in config.workspaces] + [DEFAULT_WORKSPACE_CONFIG.name]))
+            handle_cli_error(
+                logger,
+                title=f"Workspace '{resolved_workspace}' not found in config",
+                message=f"Available workspaces: {', '.join(available_workspaces)}",
+            )
+            return
+
     api_key = os.getenv(workspace_config.api_key)
     if not api_key:
         raise ValueError(f"No API key value found in environment for workspace '{workspace_config.name}'.")
@@ -94,8 +119,9 @@ def push_command(
         and workflow_config.workflow_sandbox_id
         and workflow_config.workflow_sandbox_id != workflow_sandbox_id
     ):
-        raise ValueError(
-            f"Workflow sandbox id '{workflow_sandbox_id}' is already associated with '{workflow_config.module}'."
+        logger.warning(
+            f"Workflow sandbox id '{workflow_sandbox_id}' is already associated with '{workflow_config.module}'. "
+            f"Continuing with the provided workflow sandbox id '{workflow_sandbox_id}'."
         )
 
     client = create_vellum_client(
@@ -104,22 +130,34 @@ def push_command(
     )
     sys.path.insert(0, os.getcwd())
 
-    # Remove this once we could serialize using the artifact in Vembda
-    # https://app.shortcut.com/vellum/story/5585
-    workflow = BaseWorkflow.load_from_module(workflow_config.module)
-    workflow_display = get_workflow_display(
-        workflow_class=workflow,
-        client=client,
-        dry_run=dry_run or False,
-    )
-    exec_config = workflow_display.serialize()
+    try:
+        serialization_result = BaseWorkflowDisplay.serialize_module(
+            workflow_config.module,
+            client=client,
+            dry_run=dry_run or False,
+        )
+    except ValidationError as e:
+        handle_cli_error(
+            logger,
+            title=f"Validation error while trying to push {workflow_config.module}",
+            message=str(e),
+        )
+    except Exception as e:
+        handle_cli_error(
+            logger,
+            title=f"Error while trying to push {workflow_config.module}",
+            message=str(e),
+        )
+
+    exec_config = serialization_result.exec_config
 
     container_tag = workflow_config.container_image_tag
     if workflow_config.container_image_name and not workflow_config.container_image_tag:
         container_tag = "latest"
 
     exec_config["runner_config"] = {
-        "sdk_version": metadata.version("vellum-ai"),
+        "sdk_version": metadata.version("vellum-ai"),  # Deprecated in favor of codegen_version
+        "codegen_version": metadata.version("vellum-ai"),
         "container_image_tag": container_tag,
         "container_image_name": workflow_config.container_image_name,
     }
@@ -131,12 +169,28 @@ def push_command(
             workflow_config.deployments[0] if workflow_config.deployments else WorkflowDeploymentConfig()
         )
 
-        deployment_config = WorkflowPushDeploymentConfigRequest(
-            label=deployment_label or cli_deployment_config.label,
-            name=deployment_name or cli_deployment_config.name,
-            description=deployment_description or cli_deployment_config.description,
-            release_tags=release_tags or cli_deployment_config.release_tags,
-        )
+        try:
+            module_name = workflow_config.module.split(".")[-1]
+            deployment_config = WorkflowPushDeploymentConfigRequest(
+                label=deployment_label or cli_deployment_config.label,
+                name=deployment_name or cli_deployment_config.name or to_kebab_case(module_name),
+                description=deployment_description or cli_deployment_config.description,
+                release_tags=release_tags or cli_deployment_config.release_tags,
+                release_description=release_description or cli_deployment_config.release_description,
+            )
+        except ValidationError as e:
+            for error in e.errors():
+                if "release_tags" in str(error.get("loc", [])):
+                    handle_cli_error(
+                        logger,
+                        title="Invalid release tag format",
+                        message="Release tags must be provided as separate arguments. "
+                        "Use: --release-tag tag1 --release-tag tag2",
+                    )
+                    return
+
+            # Re-raise if it's not a release_tags validation error
+            raise e
 
         # We should check with fern if we could auto-serialize typed fields for us
         # https://app.shortcut.com/vellum/story/5568
@@ -147,7 +201,7 @@ def push_command(
         module_dir = workflow_config.module.replace(".", os.path.sep)
         for root, _, files in os.walk(module_dir):
             for filename in files:
-                if not filename.endswith(".py"):
+                if not BaseWorkflowDisplay.should_include_file(filename):
                     continue
 
                 file_path = os.path.join(root, filename)
@@ -164,20 +218,41 @@ def push_command(
     artifact.seek(0)
     artifact.name = f"{workflow_config.module.replace('.', '__')}.tar.gz"
 
+    provided_id = workflow_config.workflow_sandbox_id or workflow_sandbox_id
+
+    dataset_serialized = json.dumps(serialization_result.dataset) if serialization_result.dataset else OMIT
+
     try:
         response = client.workflows.push(
-            # Remove this once we could serialize using the artifact in Vembda
-            # https://app.shortcut.com/vellum/story/5585
             exec_config=json.dumps(exec_config),
-            workflow_sandbox_id=workflow_config.workflow_sandbox_id or workflow_sandbox_id,
+            workflow_sandbox_id=provided_id,
             artifact=artifact,
             # We should check with fern if we could auto-serialize typed object fields for us
             # https://app.shortcut.com/vellum/story/5568
             deployment_config=deployment_config_serialized,  # type: ignore[arg-type]
+            dataset=dataset_serialized,  # type: ignore[arg-type]
             dry_run=dry_run,
             strict=strict,
         )
     except ApiError as e:
+        if e.status_code == 404:
+            if provided_id:
+                handle_cli_error(
+                    logger,
+                    title="Workflow Sandbox not found",
+                    message=f"Could not find Workflow Sandbox with ID '{provided_id}' "
+                    f"in workspace '{resolved_workspace}'.",
+                )
+            else:
+                error_detail = e.body.get("detail") if isinstance(e.body, dict) else None
+                default_message = "The `/workflows/push` endpoint failed with a 404 response."
+                handle_cli_error(
+                    logger,
+                    title="Workflow Sandbox not found",
+                    message=error_detail or default_message,
+                )
+            return
+
         if e.status_code == 400 and isinstance(e.body, dict) and "diffs" in e.body:
             diffs: dict = e.body["diffs"]
             generated_only = diffs.get("generated_only", [])
@@ -229,7 +304,7 @@ def push_command(
         raise e
 
     if dry_run:
-        error_messages = [str(e) for e in workflow_display.errors]
+        error_messages = [error.message for error in serialization_result.errors]
         error_message = "\n".join(error_messages) if error_messages else "No errors found."
         logger.info(
             f"""\
@@ -242,6 +317,13 @@ def push_command(
 {json.dumps(response.proposed_diffs, indent=2)}
 """
         )  # type: ignore[attr-defined]
+
+        error_list = serialization_result.errors
+        has_errors = len(error_list) > 0
+        has_diffs = response.proposed_diffs is not None and response.proposed_diffs
+
+        if has_errors or has_diffs:
+            exit(1)
     else:
         default_api_url = client._client_wrapper._environment.default
         base_url = default_api_url.split("/v1")[0].replace("//api.", "//app.")
@@ -254,7 +336,15 @@ Visit at: {base_url}/workflow-sandboxes/{response.workflow_sandbox_id}"""
         workflow_config.workflow_sandbox_id = response.workflow_sandbox_id
 
     if not workflow_config.deployments and response.workflow_deployment_id:
-        workflow_config.deployments.append(WorkflowDeploymentConfig(id=UUID(response.workflow_deployment_id)))
+        stored_deployment_config = WorkflowDeploymentConfig(
+            id=UUID(response.workflow_deployment_id),
+            label=deployment_config.label if deploy else None,
+            name=deployment_config.name if deploy else None,
+            description=deployment_config.description if deploy else None,
+            release_tags=deployment_config.release_tags if deploy else None,
+            release_description=getattr(deployment_config, "release_description", None) if deploy else None,
+        )
+        workflow_config.deployments.append(stored_deployment_config)
 
     config.save()
     logger.info("Updated vellum.lock.json file.")

vellum_cli/tests/conftest.py

@@ -1,8 +1,6 @@
 import pytest
 from dataclasses import dataclass
-import os
-import shutil
-import tempfile
+import pathlib
 from uuid import uuid4
 from typing import Any, Callable, Dict, Generator
 
@@ -18,17 +16,17 @@ class MockModuleResult:
 
 
 @pytest.fixture
-def mock_module(request) -> Generator[MockModuleResult, None, None]:
-    current_dir = os.getcwd()
-    temp_dir = tempfile.mkdtemp()
-    os.chdir(temp_dir)
+def mock_module(
+    request, monkeypatch: pytest.MonkeyPatch, tmp_path: pathlib.Path
+) -> Generator[MockModuleResult, None, None]:
+    monkeypatch.chdir(tmp_path)
 
     # Use the test name to create a unique module path
     module = f"examples.mock.{request.node.name}"
     workflow_sandbox_id = str(uuid4())
 
     def set_pyproject_toml(vellum_config: Dict[str, Any]) -> None:
-        pyproject_toml_path = os.path.join(temp_dir, "pyproject.toml")
+        pyproject_toml_path = tmp_path / "pyproject.toml"
         with open(pyproject_toml_path, "wb") as f:
             tomli_w.dump(
                 {"tool": {"vellum": vellum_config}},
@@ -46,15 +44,18 @@ def mock_module(request) -> Generator[MockModuleResult, None, None]:
         }
     )
 
-    with open(os.path.join(temp_dir, ".env"), "w") as f:
+    with open(tmp_path / ".env", "w") as f:
         f.write("VELLUM_API_KEY=abcdef123456")
 
     yield MockModuleResult(
-        temp_dir=temp_dir,
+        temp_dir=str(tmp_path),
         module=module,
         set_pyproject_toml=set_pyproject_toml,
         workflow_sandbox_id=workflow_sandbox_id,
     )
 
-    os.chdir(current_dir)
-    shutil.rmtree(temp_dir)
+
+@pytest.fixture
+def info_log_level(monkeypatch):
+    """Set log level to INFO for tests that request this fixture"""
+    monkeypatch.setenv("LOG_LEVEL", "INFO")

vellum_cli/tests/test_image_push.py

@@ -1,34 +1,37 @@
 import pytest
 import json
 import os
-import shutil
+import pathlib
 import subprocess
-import tempfile
-from unittest.mock import MagicMock, patch
+from unittest.mock import MagicMock
 from uuid import uuid4
 from typing import Generator
 
 from click.testing import CliRunner
 from httpx import Response
 
+from vellum.client.types.docker_service_token import DockerServiceToken
 from vellum_cli import main as cli_main
 
 
 @pytest.fixture
-def mock_temp_dir() -> Generator[str, None, None]:
-    current_dir = os.getcwd()
-    temp_dir = tempfile.mkdtemp()
-    os.chdir(temp_dir)
+def mock_temp_dir(monkeypatch: pytest.MonkeyPatch, tmp_path: pathlib.Path) -> Generator[str, None, None]:
+    monkeypatch.chdir(tmp_path)
+    yield str(tmp_path)
 
-    yield temp_dir
 
-    os.chdir(current_dir)
-    shutil.rmtree(temp_dir)
+@pytest.fixture
+def mock_docker_from_env(mocker):
+    return mocker.patch("docker.from_env")
+
+
+@pytest.fixture
+def mock_subprocess_run(mocker):
+    return mocker.patch("subprocess.run")
 
 
-@patch("subprocess.run")
-@patch("docker.from_env")
-def test_image_push__self_hosted_happy_path(mock_docker_from_env, mock_run, vellum_client, monkeypatch):
+@pytest.mark.usefixtures("vellum_client", "info_log_level")
+def test_image_push__self_hosted_happy_path(mock_docker_from_env, mock_subprocess_run, monkeypatch):
     # GIVEN a self hosted vellum api URL env var
     monkeypatch.setenv("VELLUM_API_URL", "mycompany.api.com")
     monkeypatch.setenv("VELLUM_API_KEY", "123456abcdef")
@@ -37,7 +40,8 @@ def test_image_push__self_hosted_happy_path(mock_docker_from_env, mock_run, vell
     mock_docker_client = MagicMock()
     mock_docker_from_env.return_value = mock_docker_client
 
-    mock_run.side_effect = [
+    mock_subprocess_run.side_effect = [
+        subprocess.CompletedProcess(args="", returncode=0, stdout=b"Pruning successful"),
         subprocess.CompletedProcess(
             args="", returncode=0, stdout=b'{"manifests": [{"platform": {"architecture": "amd64"}}]}'
         ),
@@ -56,10 +60,9 @@ def test_image_push__self_hosted_happy_path(mock_docker_from_env, mock_run, vell
     assert "Image successfully pushed" in result.output
 
 
-@patch("subprocess.run")
-@patch("docker.from_env")
+@pytest.mark.usefixtures("info_log_level")
 def test_image_push__self_hosted_happy_path__workspace_option(
-    mock_docker_from_env, mock_run, mock_httpx_transport, mock_temp_dir
+    mock_docker_from_env, mock_subprocess_run, mock_httpx_transport, mock_temp_dir
 ):
     # GIVEN a workspace config with a new env for url
     with open(os.path.join(mock_temp_dir, "vellum.lock.json"), "w") as f:
@@ -90,7 +93,8 @@ def test_image_push__self_hosted_happy_path__workspace_option(
     mock_docker_client = MagicMock()
     mock_docker_from_env.return_value = mock_docker_client
 
-    mock_run.side_effect = [
+    mock_subprocess_run.side_effect = [
+        subprocess.CompletedProcess(args="", returncode=0, stdout=b"Pruning successful"),
         subprocess.CompletedProcess(
             args="", returncode=0, stdout=b'{"manifests": [{"platform": {"architecture": "amd64"}}]}'
         ),
@@ -144,9 +148,8 @@ def test_image_push__self_hosted_happy_path__workspace_option(
     assert str(request.url) == "https://api.vellum.mycompany.ai/v1/container-images/push"
 
 
-@patch("subprocess.run")
-@patch("docker.from_env")
-def test_image_push__self_hosted_blocks_repo(mock_docker_from_env, mock_run, vellum_client, monkeypatch):
+@pytest.mark.usefixtures("vellum_client", "mock_subprocess_run")
+def test_image_push__self_hosted_blocks_repo(mock_docker_from_env, monkeypatch):
     # GIVEN a self hosted vellum api URL env var
     monkeypatch.setenv("VELLUM_API_URL", "mycompany.api.com")
 
@@ -163,3 +166,171 @@ def test_image_push__self_hosted_blocks_repo(mock_docker_from_env, mock_run, vel
 
     # AND gives the error message for self hosted installs not including the repo
     assert "For adding images to your self hosted install you must include" in result.output
+
+
+@pytest.mark.usefixtures("info_log_level")
+def test_image_push_with_source_success(
+    mock_docker_from_env, mock_subprocess_run, vellum_client, monkeypatch, mock_temp_dir
+):
+    monkeypatch.setenv("VELLUM_API_URL", "https://api.vellum.ai")
+    monkeypatch.setenv("VELLUM_API_KEY", "123456abcdef")
+
+    dockerfile_path = os.path.join(mock_temp_dir, "Dockerfile")
+    with open(dockerfile_path, "w") as f:
+        f.write("FROM alpine:latest\n")
+
+    mock_docker_client = MagicMock()
+    mock_docker_from_env.return_value = mock_docker_client
+    mock_docker_client.images.push.return_value = [b'{"status": "Pushed"}']
+
+    mock_subprocess_run.side_effect = [
+        subprocess.CompletedProcess(args="", returncode=0, stdout=b"Pruning successful"),
+        subprocess.CompletedProcess(args="", returncode=0, stdout=b"Build successful"),
+        subprocess.CompletedProcess(
+            args="", returncode=0, stdout=b'{"manifests": [{"platform": {"architecture": "amd64"}}]}'
+        ),
+        subprocess.CompletedProcess(args="", returncode=0, stdout=b"sha256:hellosha"),
+    ]
+
+    vellum_client.container_images.docker_service_token.return_value = DockerServiceToken(
+        access_token="345678mnopqr", organization_id="test-org", repository="myrepo.net"
+    )
+
+    runner = CliRunner()
+    result = runner.invoke(cli_main, ["image", "push", "myimage:latest", "--source", dockerfile_path])
+
+    assert result.exit_code == 0, result.output
+
+    build_call = mock_subprocess_run.call_args_list[1]
+    assert build_call[0][0] == [
+        "docker",
+        "buildx",
+        "build",
+        "-f",
+        dockerfile_path,
+        "--platform=linux/amd64",
+        "-t",
+        "myimage:latest",
+        ".",
+    ]
+
+    assert "Docker build completed successfully" in result.output
+    assert "Image successfully pushed" in result.output
+
+
+@pytest.mark.usefixtures("mock_docker_from_env", "mock_subprocess_run", "vellum_client")
+def test_image_push_with_source_dockerfile_not_exists(monkeypatch, mock_temp_dir):
+    monkeypatch.setenv("VELLUM_API_URL", "https://api.vellum.ai")
+    monkeypatch.setenv("VELLUM_API_KEY", "123456abcdef")
+
+    nonexistent_dockerfile = os.path.join(mock_temp_dir, "nonexistent_dockerfile")
+
+    runner = CliRunner()
+    result = runner.invoke(cli_main, ["image", "push", "myimage:latest", "--source", nonexistent_dockerfile])
+
+    assert result.exit_code == 1
+    assert "Dockerfile does not exist" in result.output
+
+
+@pytest.mark.usefixtures("mock_docker_from_env", "vellum_client")
+def test_image_push_with_source_build_fails(mock_subprocess_run, monkeypatch, mock_temp_dir):
+    monkeypatch.setenv("VELLUM_API_URL", "https://api.vellum.ai")
+    monkeypatch.setenv("VELLUM_API_KEY", "123456abcdef")
+
+    dockerfile_path = os.path.join(mock_temp_dir, "Dockerfile")
+    with open(dockerfile_path, "w") as f:
+        f.write("FROM alpine:latest\n")
+
+    mock_subprocess_run.side_effect = [
+        subprocess.CompletedProcess(args="", returncode=0, stdout=b"Pruning successful"),
+        subprocess.CompletedProcess(args="", returncode=1, stderr=b"Build failed: missing dependency"),
+    ]
+
+    runner = CliRunner()
+    result = runner.invoke(cli_main, ["image", "push", "myimage:latest", "--source", dockerfile_path])
+
+    assert result.exit_code == 1
+    assert "Docker build failed" in result.output
+    assert "Build failed: missing dependency" in result.output
+
+
+@pytest.mark.usefixtures("vellum_client", "info_log_level")
+def test_image_push_includes_docker_prune(mock_docker_from_env, mock_subprocess_run, monkeypatch):
+    """
+    Tests that image_push_command calls docker image prune before validation.
+    """
+    # GIVEN a self hosted vellum api URL env var
+    monkeypatch.setenv("VELLUM_API_URL", "mycompany.api.com")
+    monkeypatch.setenv("VELLUM_API_KEY", "123456abcdef")
+
+    # AND a mock Docker client
+    mock_docker_client = MagicMock()
+    mock_docker_from_env.return_value = mock_docker_client
+
+    mock_subprocess_run.side_effect = [
+        subprocess.CompletedProcess(args="", returncode=0, stdout=b"Pruning successful"),
+        subprocess.CompletedProcess(
+            args="", returncode=0, stdout=b'{"manifests": [{"platform": {"architecture": "amd64"}}]}'
+        ),
+        subprocess.CompletedProcess(args="", returncode=0, stdout=b"manifest"),
+        subprocess.CompletedProcess(args="", returncode=0, stdout=b"sha256:hellosha"),
+    ]
+
+    # WHEN the user runs the image push command with any image
+    runner = CliRunner()
+    result = runner.invoke(cli_main, ["image", "push", "myrepo.net/myimage:latest"])
+
+    # THEN the command exits successfully
+    assert result.exit_code == 0, result.output
+
+    prune_call = mock_subprocess_run.call_args_list[0]
+    assert prune_call[0][0] == [
+        "docker",
+        "image",
+        "prune",
+        "--all",
+        "-f",
+        "--filter",
+        "label=image-type=python-workflow-runtime",
+    ]
+
+    # AND the success message includes pruning completion
+    assert "Docker image pruning completed successfully" in result.output
+
+
+@pytest.mark.usefixtures("vellum_client", "info_log_level")
+def test_image_push_continues_if_prune_fails(mock_docker_from_env, mock_subprocess_run, monkeypatch):
+    """
+    Tests that image_push_command continues if docker image prune fails.
+    """
+    # GIVEN a self hosted vellum api URL env var
+    monkeypatch.setenv("VELLUM_API_URL", "mycompany.api.com")
+    monkeypatch.setenv("VELLUM_API_KEY", "123456abcdef")
+
+    # AND a mock Docker client
+    mock_docker_client = MagicMock()
+    mock_docker_from_env.return_value = mock_docker_client
+
+    mock_subprocess_run.side_effect = [
+        subprocess.CalledProcessError(
+            1, ["docker", "image", "prune", "--all", "-f", "--filter", "label=image-type=python-workflow-runtime"]
+        ),
+        subprocess.CompletedProcess(
+            args="", returncode=0, stdout=b'{"manifests": [{"platform": {"architecture": "amd64"}}]}'
+        ),
+        subprocess.CompletedProcess(args="", returncode=0, stdout=b"manifest"),
+        subprocess.CompletedProcess(args="", returncode=0, stdout=b"sha256:hellosha"),
+    ]
+
+    # WHEN the user runs the image push command with any image
+    runner = CliRunner()
+    result = runner.invoke(cli_main, ["image", "push", "myrepo.net/myimage:latest"])
+
+    # THEN the command exits successfully despite pruning failure
+    assert result.exit_code == 0, result.output
+
+    # AND the warning message about pruning failure is shown
+    assert "Docker image pruning failed" in result.output
+
+    # AND the success message is still shown
+    assert "Image successfully pushed" in result.output