vellum-ai 0.14.65__py3-none-any.whl → 0.14.67__py3-none-any.whl

vellum_cli/tests/test_image_push_error_handling.py

@@ -0,0 +1,184 @@
+import subprocess
+from unittest.mock import patch
+from uuid import uuid4
+
+from click.testing import CliRunner
+
+from vellum.client.core.api_error import ApiError
+from vellum_cli import main as cli_main
+
+
+@patch("subprocess.run")
+@patch("docker.from_env")
+def test_image_push_docker_service_token_401_error(mock_docker_from_env, mock_run, vellum_client, monkeypatch):
+    monkeypatch.setenv("VELLUM_API_URL", "https://api.vellum.ai")
+    monkeypatch.setenv("VELLUM_API_KEY", "123456abcdef")
+
+    mock_docker_client = mock_docker_from_env.return_value
+    mock_docker_client.images.get.return_value.id = "test-image-id"
+
+    mock_run.side_effect = [
+        subprocess.CompletedProcess(
+            args="", returncode=0, stdout=b'{"manifests": [{"platform": {"architecture": "amd64"}}]}'
+        ),
+    ]
+
+    vellum_client.container_images.docker_service_token.side_effect = ApiError(status_code=401, body="Unauthorized")
+
+    runner = CliRunner()
+    result = runner.invoke(cli_main, ["image", "push", "myimage:latest"])
+
+    assert result.exit_code == 1
+    assert "Authentication failed" in result.output
+    assert "VELLUM_API_KEY" in result.output
+
+
+@patch("subprocess.run")
+@patch("docker.from_env")
+def test_image_push_docker_service_token_500_error(mock_docker_from_env, mock_run, vellum_client, monkeypatch):
+    monkeypatch.setenv("VELLUM_API_URL", "https://api.vellum.ai")
+    monkeypatch.setenv("VELLUM_API_KEY", "123456abcdef")
+
+    mock_docker_client = mock_docker_from_env.return_value
+    mock_docker_client.images.get.return_value.id = "test-image-id"
+
+    mock_run.side_effect = [
+        subprocess.CompletedProcess(
+            args="", returncode=0, stdout=b'{"manifests": [{"platform": {"architecture": "amd64"}}]}'
+        ),
+    ]
+
+    vellum_client.container_images.docker_service_token.side_effect = ApiError(
+        status_code=500, body="Internal Server Error"
+    )
+
+    runner = CliRunner()
+    result = runner.invoke(cli_main, ["image", "push", "myimage:latest"])
+
+    assert result.exit_code == 1
+    assert "Server error" in result.output
+    assert "try again later" in result.output
+
+
+@patch("subprocess.run")
+@patch("docker.from_env")
+def test_image_push_docker_service_token_other_error(mock_docker_from_env, mock_run, vellum_client, monkeypatch):
+    monkeypatch.setenv("VELLUM_API_URL", "https://api.vellum.ai")
+    monkeypatch.setenv("VELLUM_API_KEY", "123456abcdef")
+
+    mock_docker_client = mock_docker_from_env.return_value
+    mock_docker_client.images.get.return_value.id = "test-image-id"
+
+    mock_run.side_effect = [
+        subprocess.CompletedProcess(
+            args="", returncode=0, stdout=b'{"manifests": [{"platform": {"architecture": "amd64"}}]}'
+        ),
+    ]
+
+    vellum_client.container_images.docker_service_token.side_effect = ApiError(
+        status_code=429, body="Too Many Requests"
+    )
+
+    runner = CliRunner()
+    result = runner.invoke(cli_main, ["image", "push", "myimage:latest"])
+
+    assert result.exit_code == 1
+    assert "API request failed" in result.output
+    assert "HTTP 429" in result.output
+
+
+@patch("subprocess.run")
+@patch("docker.from_env")
+def test_image_push_container_image_401_error(mock_docker_from_env, mock_run, vellum_client, monkeypatch):
+    monkeypatch.setenv("VELLUM_API_URL", "https://api.vellum.ai")
+    monkeypatch.setenv("VELLUM_API_KEY", "123456abcdef")
+
+    mock_docker_client = mock_docker_from_env.return_value
+    mock_docker_client.images.get.return_value.id = "test-image-id"
+    mock_docker_client.images.push.return_value = ["pushed"]
+
+    mock_run.side_effect = [
+        subprocess.CompletedProcess(
+            args="", returncode=0, stdout=b'{"manifests": [{"platform": {"architecture": "amd64"}}]}'
+        ),
+        subprocess.CompletedProcess(args="", returncode=0, stdout=b'[{"RepoDigests": ["test-repo@sha256:abcd1234"]}]'),
+    ]
+
+    vellum_client.container_images.docker_service_token.return_value = type(
+        "obj", (object,), {"access_token": "345678mnopqr", "organization_id": str(uuid4()), "repository": "myrepo.net"}
+    )()
+
+    vellum_client.container_images.push_container_image.side_effect = ApiError(status_code=401, body="Unauthorized")
+
+    runner = CliRunner()
+    result = runner.invoke(cli_main, ["image", "push", "myimage:latest"])
+
+    assert result.exit_code == 1
+    assert "Authentication failed" in result.output
+    assert "VELLUM_API_KEY" in result.output
+
+
+@patch("subprocess.run")
+@patch("docker.from_env")
+def test_image_push_container_image_500_error(mock_docker_from_env, mock_run, vellum_client, monkeypatch):
+    monkeypatch.setenv("VELLUM_API_URL", "https://api.vellum.ai")
+    monkeypatch.setenv("VELLUM_API_KEY", "123456abcdef")
+
+    mock_docker_client = mock_docker_from_env.return_value
+    mock_docker_client.images.get.return_value.id = "test-image-id"
+    mock_docker_client.images.push.return_value = ["pushed"]
+
+    mock_run.side_effect = [
+        subprocess.CompletedProcess(
+            args="", returncode=0, stdout=b'{"manifests": [{"platform": {"architecture": "amd64"}}]}'
+        ),
+        subprocess.CompletedProcess(args="", returncode=0, stdout=b'[{"RepoDigests": ["test-repo@sha256:abcd1234"]}]'),
+    ]
+
+    vellum_client.container_images.docker_service_token.return_value = type(
+        "obj", (object,), {"access_token": "345678mnopqr", "organization_id": str(uuid4()), "repository": "myrepo.net"}
+    )()
+
+    vellum_client.container_images.push_container_image.side_effect = ApiError(
+        status_code=500, body="Internal Server Error"
+    )
+
+    runner = CliRunner()
+    result = runner.invoke(cli_main, ["image", "push", "myimage:latest"])
+
+    assert result.exit_code == 1
+    assert "Server error" in result.output
+    assert "try again later" in result.output
+
+
+@patch("subprocess.run")
+@patch("docker.from_env")
+def test_image_push_container_image_other_error(mock_docker_from_env, mock_run, vellum_client, monkeypatch):
+    monkeypatch.setenv("VELLUM_API_URL", "https://api.vellum.ai")
+    monkeypatch.setenv("VELLUM_API_KEY", "123456abcdef")
+
+    mock_docker_client = mock_docker_from_env.return_value
+    mock_docker_client.images.get.return_value.id = "test-image-id"
+    mock_docker_client.images.push.return_value = ["pushed"]
+
+    mock_run.side_effect = [
+        subprocess.CompletedProcess(
+            args="", returncode=0, stdout=b'{"manifests": [{"platform": {"architecture": "amd64"}}]}'
+        ),
+        subprocess.CompletedProcess(args="", returncode=0, stdout=b'[{"RepoDigests": ["test-repo@sha256:abcd1234"]}]'),
+    ]
+
+    vellum_client.container_images.docker_service_token.return_value = type(
+        "obj", (object,), {"access_token": "345678mnopqr", "organization_id": str(uuid4()), "repository": "myrepo.net"}
+    )()
+
+    vellum_client.container_images.push_container_image.side_effect = ApiError(
+        status_code=429, body="Too Many Requests"
+    )
+
+    runner = CliRunner()
+    result = runner.invoke(cli_main, ["image", "push", "myimage:latest"])
+
+    assert result.exit_code == 1
+    assert "API request failed" in result.output
+    assert "HTTP 429" in result.output

vellum_cli/tests/test_pull.py

@@ -745,7 +745,7 @@ def test_pull__same_pull_twice__one_entry_in_lockfile(vellum_client, mock_module
     zip_contents = _zip_file_map({"workflow.py": "print('hello')"})
     responses = iter([zip_contents, zip_contents])
 
-    def workflows_pull_side_effect(*args, **kwargs):
+    def workflows_pull_side_effect(*_args, **_kwargs):
         return iter([next(responses)])
 
     vellum_client.workflows.pull.side_effect = workflows_pull_side_effect
@@ -890,6 +890,7 @@ def test_pull__module_name_from_deployment_name(vellum_client):
         assert f.read() == "print('hello')"
 
 
+@pytest.mark.usefixtures("mock_module")
 def test_pull__invalid_zip_file(vellum_client):
     workflow_deployment = "test-workflow-deployment-id"
 
@@ -903,11 +904,12 @@ def test_pull__invalid_zip_file(vellum_client):
     # THEN the command returns an error
     assert result.exit_code == 1
     assert (
-        str(result.exception)
-        == "The API we tried to pull from returned an invalid zip file. Please make sure your `VELLUM_API_URL` environment variable is set correctly."  # noqa: E501
+        "Invalid response format" in result.output
+        or "Please verify your `VELLUM_API_URL` environment variable is set correctly" in result.output
     )
 
 
+@pytest.mark.usefixtures("mock_module")
 def test_pull__json_decode_error(vellum_client):
     workflow_deployment = "test-workflow-deployment-id"
 
@@ -927,11 +929,12 @@ def test_pull__json_decode_error(vellum_client):
     # THEN the command returns an error
     assert result.exit_code == 1
     assert (
-        str(result.exception)
-        == "The API we tried to pull is invalid. Please make sure your `VELLUM_API_URL` environment variable is set correctly."  # noqa: E501
+        "API request failed" in result.output
+        or "Please verify your `VELLUM_API_URL` environment variable is set correctly" in result.output
     )
 
 
+@pytest.mark.usefixtures("mock_module")
 def test_pull__unauthorized_error_path(vellum_client):
     workflow_deployment = "test-workflow-deployment-id"
 
@@ -948,9 +951,10 @@ def test_pull__unauthorized_error_path(vellum_client):
 
     # THEN the command returns an error
     assert result.exit_code == 1
-    assert str(result.exception) == "Please make sure your `VELLUM_API_KEY` environment variable is set correctly."
+    assert "Please make sure your `VELLUM_API_KEY` environment variable is set correctly" in result.output
 
 
+@pytest.mark.usefixtures("mock_module")
 def test_pull__unexpected_error_path(vellum_client):
     workflow_deployment = "test-workflow-deployment-id"
 
@@ -968,9 +972,8 @@ def test_pull__unexpected_error_path(vellum_client):
     # THEN the command returns an error
     assert result.exit_code == 1
     assert (
-        str(result.exception)
-        == """The Pull API failed with an unexpected error. \
-Please try again later and contact support if the problem persists."""
+        "Server error occurred" in result.output
+        or "Please try again in a few moments. If the problem persists, contact Vellum support" in result.output
     )