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

vellum_cli/tests/test_push.py

@@ -13,6 +13,7 @@ from vellum.client.core.api_error import ApiError
 from vellum.client.types.workflow_push_response import WorkflowPushResponse
 from vellum.utils.uuid import is_valid_uuid
 from vellum_cli import main as cli_main
+from vellum_ee.workflows.display.nodes.utils import to_kebab_case
 
 
 def _extract_tar_gz(tar_gz_bytes: bytes) -> dict[str, str]:
@@ -61,8 +62,41 @@ def test_push__no_config(mock_module):
 
     # THEN it should fail
     assert result.exit_code == 1
-    assert result.exception
-    assert str(result.exception) == "No Workflows found in project to push."
+    assert "No workflows found in project to push" in result.output
+    assert "No workflow configurations were found in vellum.lock.json" in result.output
+
+
+def test_push__no_matching_config_with_workflow_sandbox_id(mock_module):
+    # GIVEN a workflow was pulled with a specific sandbox ID
+    temp_dir = mock_module.temp_dir
+    module = mock_module.module
+    original_workflow_sandbox_id = mock_module.workflow_sandbox_id
+
+    # AND we have a valid project config from a pulled workflow
+    mock_module.set_pyproject_toml(
+        {
+            "workflows": [
+                {
+                    "module": module,
+                    "workflow_sandbox_id": original_workflow_sandbox_id,
+                }
+            ]
+        }
+    )
+
+    # AND a workflow exists in the module successfully
+    _ensure_workflow_py(temp_dir, module)
+
+    # WHEN the user tries to push with a different workflow sandbox ID via --workflow-sandbox-id
+    different_workflow_sandbox_id = str(uuid4())
+    runner = CliRunner()
+    result = runner.invoke(cli_main, ["workflows", "push", "--workflow-sandbox-id", different_workflow_sandbox_id])
+
+    # THEN it should fail with a helpful error message
+    assert result.exit_code == 1
+    assert "No workflows found in project to push" in result.output
+    assert "No workflow configurations were found in vellum.lock.json" in result.output
+    assert "vellum workflows push" in result.output
 
 
 def test_push__multiple_workflows_configured__no_module_specified(mock_module):
@@ -76,9 +110,8 @@ def test_push__multiple_workflows_configured__no_module_specified(mock_module):
     # THEN it should fail
     assert result.exit_code == 1
     assert result.exception
-    assert (
-        str(result.exception)
-        == "Multiple workflows found in project to push. Pushing only a single workflow is supported."
+    assert str(result.exception) == (
+        "Multiple workflows found. Please specify a single workflow to push. Found: examples.mock, examples.mock2"
     )
 
 
@@ -140,6 +173,7 @@ def test_push__happy_path(mock_module, vellum_client, base_command):
     assert extracted_files["workflow.py"] == workflow_py_file_content
 
 
+@pytest.mark.usefixtures("info_log_level")
 def test_push__verify_default_url_in_raw_httpx_transport(mock_module, mock_httpx_transport):
     # GIVEN a single workflow configured
     module = mock_module.module
@@ -347,9 +381,8 @@ def test_push__workflow_sandbox_option__existing_id_different_module(mock_module
     # THEN it should fail
     assert result.exit_code == 1
     assert result.exception
-    assert (
-        str(result.exception)
-        == "Multiple workflows found in project to push. Pushing only a single workflow is supported."
+    assert str(result.exception) == (
+        f"Multiple workflows found. Please specify a single workflow to push. Found: {module}, {second_module}"
     )
 
 
@@ -390,12 +423,15 @@ def test_push__deployment(mock_module, vellum_client, base_command):
     assert json.loads(call_args["exec_config"])["workflow_raw_data"]["definition"]["name"] == "ExampleWorkflow"
     assert is_valid_uuid(call_args["workflow_sandbox_id"])
     assert call_args["artifact"].name == expected_artifact_name
-    assert call_args["deployment_config"] == "{}"
+    expected_deployment_name = to_kebab_case(module.split(".")[-1])
+    deployment_config = json.loads(call_args["deployment_config"])
+    assert deployment_config["name"] == expected_deployment_name
 
     extracted_files = _extract_tar_gz(call_args["artifact"].read())
     assert extracted_files["workflow.py"] == workflow_py_file_content
 
 
+@pytest.mark.usefixtures("info_log_level")
 def test_push__dry_run_option_returns_report(mock_module, vellum_client):
     # GIVEN a single workflow configured
     temp_dir = mock_module.temp_dir
@@ -444,8 +480,8 @@ class ExampleWorkflow(BaseWorkflow):
     runner = CliRunner(mix_stderr=True)
     result = runner.invoke(cli_main, ["push", module, "--dry-run"])
 
-    # THEN it should succeed
-    assert result.exit_code == 0
+    # THEN it should fail with exit code 1 due to errors
+    assert result.exit_code == 1
 
     # AND we should have called the push API with the dry-run option
     vellum_client.workflows.push.assert_called_once()
@@ -459,6 +495,39 @@ class ExampleWorkflow(BaseWorkflow):
     assert "iterable_item_added" in result.output
 
 
+@pytest.mark.usefixtures("info_log_level")
+def test_push__dry_run_option_no_errors_returns_success(mock_module, vellum_client):
+    """Test that dry-run returns exit code 0 when there are no errors or diffs"""
+    # GIVEN a workflow module with a valid workflow (using the same pattern as happy path test)
+    temp_dir = mock_module.temp_dir
+    module = mock_module.module
+    _ensure_workflow_py(temp_dir, module)
+
+    # AND the push API call returns successfully with no errors and no diffs
+    vellum_client.workflows.push.return_value = WorkflowPushResponse(
+        workflow_sandbox_id=str(uuid4()),
+        proposed_diffs=None,
+    )
+
+    # WHEN calling `vellum push` with dry-run
+    runner = CliRunner(mix_stderr=True)
+    result = runner.invoke(cli_main, ["push", module, "--dry-run"])
+
+    # THEN it should succeed with exit code 0
+    assert result.exit_code == 0
+
+    # AND we should have called the push API with the dry-run option
+    vellum_client.workflows.push.assert_called_once()
+    call_args = vellum_client.workflows.push.call_args.kwargs
+    assert call_args["dry_run"] is True
+
+    # AND the report should be in the output
+    assert "## Errors" in result.output
+    assert "No errors found." in result.output
+    assert "## Proposed Diffs" in result.output
+
+
+@pytest.mark.usefixtures("info_log_level")
 def test_push__strict_option_returns_diffs(mock_module, vellum_client):
     # GIVEN a single workflow configured
     temp_dir = mock_module.temp_dir
@@ -556,6 +625,84 @@ def test_push__push_fails_due_to_400_error(mock_module, vellum_client):
     assert "Traceback" not in result.output
 
 
+def test_push__push_fails_due_to_404_error_with_id(mock_module, vellum_client):
+    """
+    Tests that a 404 error with a workflow sandbox ID shows a helpful error message.
+    """
+
+    # GIVEN a single workflow configured with a workflow_sandbox_id
+    temp_dir = mock_module.temp_dir
+    module = mock_module.module
+
+    # AND a workflow exists in the module successfully
+    _ensure_workflow_py(temp_dir, module)
+
+    # AND the push API call returns a 404 response
+    vellum_client.workflows.push.side_effect = ApiError(
+        status_code=404,
+        body={},
+    )
+
+    # WHEN calling `vellum push`
+    runner = CliRunner()
+    result = runner.invoke(cli_main, ["push", module])
+
+    # THEN it should fail with a user error code
+    assert result.exit_code == 1
+
+    # AND the error message should be in the error message
+    assert "Workflow Sandbox not found" in result.output
+    assert mock_module.workflow_sandbox_id in result.output
+    assert "workspace" in result.output
+
+    # AND the stack trace should not be present
+    assert "Traceback" not in result.output
+
+
+def test_push__push_fails_due_to_404_error_without_id(tmp_path, vellum_client, monkeypatch):
+    """
+    Tests that a 404 error without a workflow sandbox ID shows a helpful error message.
+    """
+
+    # GIVEN a workflow module exists but has no workflow_sandbox_id configured
+    module = "examples.mock.test_push__404_without_id"
+    module_dir = tmp_path / "examples" / "mock" / "test_push__404_without_id"
+    module_dir.mkdir(parents=True)
+
+    # AND a workflow.py file exists
+    workflow_file = module_dir / "workflow.py"
+    workflow_file.write_text(
+        "from vellum.workflows import BaseWorkflow\n\nclass ExampleWorkflow(BaseWorkflow):\n    pass\n"
+    )
+
+    # AND the pyproject.toml file exists without workflow_sandbox_id
+    pyproject_file = tmp_path / "pyproject.toml"
+    pyproject_file.write_text(f'[tool.vellum]\n[[tool.vellum.workflows]]\nmodule = "{module}"\n')
+
+    monkeypatch.chdir(tmp_path)
+
+    # AND the push API call returns a 404 response
+    vellum_client.workflows.push.side_effect = ApiError(
+        status_code=404,
+        body={},
+    )
+
+    # WHEN calling `vellum push`
+    runner = CliRunner()
+    result = runner.invoke(cli_main, ["push", module])
+
+    # THEN it should fail with a user error code
+    assert result.exit_code == 1
+
+    # AND the error message should be in the output
+    assert "Workflow Sandbox not found" in result.output
+    assert "/workflows/push" in result.output
+    assert "404 response" in result.output
+
+    # AND the stack trace should not be present
+    assert "Traceback" not in result.output
+
+
 @pytest.mark.parametrize(
     "file_data",
     [
@@ -685,6 +832,7 @@ MY_OTHER_VELLUM_API_KEY=aaabbbcccddd
     vellum_client_class.assert_called_once_with(
         api_key="aaabbbcccddd",
         environment=mock.ANY,
+        api_version=None,
     )
 
     # AND the vellum lock file should have been updated with the correct workspace
@@ -702,6 +850,7 @@ MY_OTHER_VELLUM_API_KEY=aaabbbcccddd
         }
 
 
+@pytest.mark.usefixtures("info_log_level")
 def test_push__workspace_option__uses_different_api_url_env(mock_module, mock_httpx_transport):
     # GIVEN a single workflow configured
     temp_dir = mock_module.temp_dir
@@ -845,6 +994,7 @@ MY_OTHER_VELLUM_API_KEY=aaabbbcccddd
     vellum_client_class.assert_called_once_with(
         api_key="aaabbbcccddd",
         environment=mock.ANY,
+        api_version=None,
     )
 
     # AND the vellum lock file should have the same two workflows
@@ -943,3 +1093,285 @@ def test_push__use_default_workspace_if_not_specified__multiple_workflows_config
         config = configs[0]
         assert config["workflow_sandbox_id"] == workflow_sandbox_id
         assert config["workspace"] == "default"
+
+
+def test_push__deploy_with_malformed_release_tags_shows_friendly_validation_error(mock_module, vellum_client):
+    # GIVEN a single workflow configured
+    temp_dir = mock_module.temp_dir
+    module = mock_module.module
+
+    # AND a workflow exists in the module successfully
+    _ensure_workflow_py(temp_dir, module)
+
+    # AND the push API call would return successfully
+    vellum_client.workflows.push.return_value = WorkflowPushResponse(
+        workflow_sandbox_id=str(uuid4()),
+    )
+
+    # WHEN calling `vellum workflows push` with --deploy and --release-tag
+    runner = CliRunner()
+    result = runner.invoke(cli_main, ["workflows", "push", module, "--deploy", "--release-tag", None])  # type: ignore
+
+    # THEN it should show the friendly error message instead of a raw Pydantic traceback
+    assert "Invalid release tag format" in result.output
+    assert "Release tags must be provided as separate arguments" in result.output
+    assert "--release-tag tag1 --release-tag tag2" in result.output
+
+
+@pytest.mark.usefixtures("info_log_level")
+def test_push__deploy_with_release_tags_success(mock_module, vellum_client):
+    # GIVEN a single workflow configured
+    temp_dir = mock_module.temp_dir
+    module = mock_module.module
+
+    # AND a workflow exists in the module successfully
+    _ensure_workflow_py(temp_dir, module)
+
+    # AND the push API call returns successfully
+    workflow_deployment_id = str(uuid4())
+    vellum_client.workflows.push.return_value = WorkflowPushResponse(
+        workflow_sandbox_id=str(uuid4()),
+        workflow_deployment_id=workflow_deployment_id,
+    )
+
+    # WHEN calling `vellum workflows push` with --deploy and --release-tag
+    runner = CliRunner()
+    result = runner.invoke(cli_main, ["workflows", "push", module, "--deploy", "--release-tag", "v1.0.0"])
+
+    # THEN it should succeed
+    assert result.exit_code == 0, result.output
+
+    # AND we should have called the push API with the correct deployment config
+    vellum_client.workflows.push.assert_called_once()
+    call_args = vellum_client.workflows.push.call_args.kwargs
+
+    # AND the deployment_config should contain the release tags
+    deployment_config_str = call_args["deployment_config"]
+    deployment_config = json.loads(deployment_config_str)
+    assert deployment_config["release_tags"] == ["v1.0.0"]
+
+    # AND should show success message
+    assert "Successfully pushed" in result.output
+    assert "Updated vellum.lock.json file." in result.output
+
+
+@pytest.mark.usefixtures("info_log_level")
+def test_push__deploy_with_release_description_success(mock_module, vellum_client):
+    # GIVEN a single workflow configured
+    temp_dir = mock_module.temp_dir
+    module = mock_module.module
+
+    # AND a workflow exists in the module successfully
+    _ensure_workflow_py(temp_dir, module)
+
+    # AND the push API call returns successfully
+    workflow_deployment_id = str(uuid4())
+    vellum_client.workflows.push.return_value = WorkflowPushResponse(
+        workflow_sandbox_id=str(uuid4()),
+        workflow_deployment_id=workflow_deployment_id,
+    )
+
+    # WHEN calling `vellum workflows push` with --deploy and --release-description
+    runner = CliRunner()
+    result = runner.invoke(
+        cli_main, ["workflows", "push", module, "--deploy", "--release-description", "This is a test release"]
+    )
+
+    # THEN it should succeed
+    assert result.exit_code == 0, result.output
+
+    # AND we should have called the push API with the correct deployment config
+    vellum_client.workflows.push.assert_called_once()
+    call_args = vellum_client.workflows.push.call_args.kwargs
+
+    # AND the deployment_config should contain the release description
+    deployment_config_str = call_args["deployment_config"]
+    deployment_config = json.loads(deployment_config_str)
+    assert deployment_config["release_description"] == "This is a test release"
+
+    # AND should show success message
+    assert "Successfully pushed" in result.output
+    assert "Updated vellum.lock.json file." in result.output
+
+
+def test_push__deploy_stores_deployment_config_in_lock_file(mock_module, vellum_client):
+    # GIVEN a single workflow
+    temp_dir = mock_module.temp_dir
+    module = mock_module.module
+
+    # AND a workflow exists in the module successfully
+    _ensure_workflow_py(temp_dir, module)
+
+    # AND the push API call returns successfully with a deployment
+    workflow_deployment_id = str(uuid4())
+    vellum_client.workflows.push.return_value = WorkflowPushResponse(
+        workflow_sandbox_id=str(uuid4()),
+        workflow_deployment_id=workflow_deployment_id,
+    )
+
+    # WHEN calling `vellum workflows push --deploy` for the first time
+    runner = CliRunner()
+    result = runner.invoke(cli_main, ["workflows", "push", module, "--deploy"])
+
+    # THEN it should succeed
+    assert result.exit_code == 0, result.output
+
+    # AND the deployment config should be stored in the lock file with the deployment ID and module name
+    with open(os.path.join(temp_dir, "vellum.lock.json")) as f:
+        lock_data = json.loads(f.read())
+        assert len(lock_data["workflows"][0]["deployments"]) == 1
+        deployment = lock_data["workflows"][0]["deployments"][0]
+        assert str(deployment["id"]) == workflow_deployment_id
+        assert deployment["name"] == "test-push-deploy-stores-deployment-config-in-lock-file"
+        assert deployment.get("label") is None
+        assert deployment.get("description") is None
+        assert deployment.get("release_tags") is None
+
+    # AND when we do a second push
+    vellum_client.workflows.push.reset_mock()
+    vellum_client.workflows.push.return_value = WorkflowPushResponse(
+        workflow_sandbox_id=str(uuid4()),
+        workflow_deployment_id=workflow_deployment_id,
+    )
+
+    result = runner.invoke(cli_main, ["workflows", "push", module, "--deploy"])
+
+    # THEN it should succeed
+    assert result.exit_code == 0, result.output
+
+    # AND we should have called the push API with the module name as deployment name
+    vellum_client.workflows.push.assert_called_once()
+    call_args = vellum_client.workflows.push.call_args.kwargs
+    deployment_config_str = call_args["deployment_config"]
+    deployment_config = json.loads(deployment_config_str)
+    assert deployment_config["name"] == "test-push-deploy-stores-deployment-config-in-lock-file"
+
+
+def test_push__custom_workspace_in_lockfile__uses_custom_workspace_without_flag(mock_module, vellum_client_class):
+    """
+    Tests that push respects the custom workspace from the lockfile when no --workspace flag is provided.
+    """
+
+    # GIVEN a workflow already configured in the lockfile with a custom workspace
+    temp_dir = mock_module.temp_dir
+    module = mock_module.module
+    workflow_sandbox_id = mock_module.workflow_sandbox_id
+    set_pyproject_toml = mock_module.set_pyproject_toml
+
+    # AND the lockfile has the workflow with a custom workspace
+    with open(os.path.join(temp_dir, "vellum.lock.json"), "w") as f:
+        json.dump(
+            {
+                "version": "1.0",
+                "workflows": [
+                    {
+                        "module": module,
+                        "workflow_sandbox_id": workflow_sandbox_id,
+                        "workspace": "my_custom_workspace",
+                        "container_image_name": None,
+                        "container_image_tag": None,
+                        "deployments": [],
+                        "ignore": None,
+                        "target_directory": None,
+                    }
+                ],
+                "workspaces": [],
+            },
+            f,
+            indent=2,
+        )
+
+    # AND the custom workspace is defined in pyproject.toml
+    set_pyproject_toml(
+        {
+            "workflows": [],
+            "workspaces": [
+                {
+                    "name": "my_custom_workspace",
+                    "api_key": "MY_CUSTOM_VELLUM_API_KEY",
+                }
+            ],
+        }
+    )
+
+    # AND the .env file has the custom api key
+    with open(os.path.join(temp_dir, ".env"), "w") as f:
+        f.write(
+            """\
+VELLUM_API_KEY=abcdef123456
+MY_CUSTOM_VELLUM_API_KEY=custom-key-xyz
+"""
+        )
+
+    # AND a workflow exists in the module
+    _ensure_workflow_py(temp_dir, module)
+
+    # AND the push API returns successfully
+    vellum_client_class.return_value.workflows.push.return_value = WorkflowPushResponse(
+        workflow_sandbox_id=workflow_sandbox_id,
+    )
+    vellum_client_class.return_value._client_wrapper._environment.default = "https://api.vellum.ai/v1"
+
+    # WHEN calling `vellum push` WITHOUT the --workspace flag
+    runner = CliRunner()
+    result = runner.invoke(cli_main, ["push", module])
+
+    # THEN it should succeed
+    assert result.exit_code == 0, result.output
+
+    # AND the custom workspace API key should have been used
+    vellum_client_class.assert_called_once_with(
+        api_key="custom-key-xyz",
+        environment=mock.ANY,
+        api_version=None,
+    )
+
+    with open(os.path.join(temp_dir, "vellum.lock.json")) as f:
+        lock_file_content = json.load(f)
+        assert len(lock_file_content["workflows"]) == 1
+        assert lock_file_content["workflows"][0]["workspace"] == "my_custom_workspace"
+        assert lock_file_content["workflows"][0]["workflow_sandbox_id"] == workflow_sandbox_id
+
+
+def test_push__workspace_option__nonexistent_workspace_should_fail(mock_module):
+    """
+    Tests that pushing with a nonexistent workspace that isn't default should fail.
+    """
+
+    # GIVEN a single workflow configured
+    temp_dir = mock_module.temp_dir
+    module = mock_module.module
+    set_pyproject_toml = mock_module.set_pyproject_toml
+
+    # AND a custom workspace is configured
+    set_pyproject_toml(
+        {
+            "workflows": [
+                {
+                    "module": module,
+                }
+            ],
+            "workspaces": [
+                {
+                    "name": "my_custom_workspace",
+                    "api_key": "MY_CUSTOM_VELLUM_API_KEY",
+                }
+            ],
+        }
+    )
+
+    # AND a workflow exists in the module successfully
+    _ensure_workflow_py(temp_dir, module)
+
+    # WHEN calling `vellum push` with a nonexistent workspace
+    runner = CliRunner()
+    result = runner.invoke(cli_main, ["push", module, "--workspace", "nonexistent_workspace"])
+
+    # THEN it should exit with an error
+    assert result.exit_code == 1
+
+    # AND the error message should indicate the workspace doesn't exist
+    assert "Workspace 'nonexistent_workspace' not found in config" in result.output
+    assert "Available workspaces:" in result.output
+    assert "default" in result.output
+    assert "my_custom_workspace" in result.output