Update tests

Author: dima-aignostics

src/aignostics/platform/resources/runs.py

@@ -292,6 +292,9 @@ def get_artifact_download_url(self, artifact_id: str) -> str:
             RuntimeError: If the redirect ``Location`` header is missing or the
                 response status is unexpected.
         """
+
+        # Generated client follows redirect automatically which prevents us from getting the presigned URL,
+        # but we need the presigned URL to get query the metadata of the file (e.g. checksum) before downloading it.
         serialize = (
             self._api._get_artifact_url_v1_runs_run_id_artifacts_artifact_id_file_get_serialize  # noqa: SLF001
         )
@@ -303,24 +306,24 @@ def get_artifact_download_url(self, artifact_id: str) -> str:
             _headers={"User-Agent": user_agent()},
             _host_index=0,
         )
-        response = requests.get(
+        with requests.get(
             url,
             headers=dict(header_params),
             allow_redirects=False,
             timeout=settings().run_timeout,
-        )
-        if response.status_code == requests.codes.temporary_redirect:
-            location = response.headers.get("Location")
-            if not location:
-                msg = f"307 redirect received but Location header is absent for artifact {artifact_id!r}"
-                raise RuntimeError(msg)
-            return location
-        response.raise_for_status()
-        msg = (
-            f"Unexpected status {response.status_code} from artifact URL endpoint "
-            f"for artifact {artifact_id!r}; expected 307 redirect"
-        )
-        raise RuntimeError(msg)
+        ) as response:
+            if response.status_code == requests.codes.temporary_redirect:
+                location = response.headers.get("Location")
+                if not location:
+                    msg = f"307 redirect received but Location header is absent for artifact {artifact_id!r}"
+                    raise RuntimeError(msg)
+                return location
+            response.raise_for_status()
+            msg = (
+                f"Unexpected status {response.status_code} from artifact URL endpoint "
+                f"for artifact {artifact_id!r}; expected 307 redirect"
+            )
+            raise RuntimeError(msg)
 
     def download_to_folder(  # noqa: C901
         self,

tests/aignostics/application/download_test.py

@@ -403,87 +403,3 @@ def progress_callback(p: DownloadProgress) -> None:
 
         # Verify direct URL was used (no signed URL generation)
         mock_get.assert_called_once_with(https_url, stream=True, timeout=60)
-
-
-# ---------------------------------------------------------------------------
-# download_item_artifact tests
-# ---------------------------------------------------------------------------
-
-
-def _make_crc32c_checksum(data: bytes) -> str:
-    """Compute a base64-encoded CRC32C checksum matching the SDK's format."""
-    h = crc32c.CRC32CHash()
-    h.update(data)
-    return base64.b64encode(h.digest()).decode("ascii")
-
-
-@pytest.mark.unit
-def test_download_item_artifact_success(tmp_path: Path) -> None:
-    """Test that download_item_artifact fetches a URL from the run and downloads the file."""
-    file_content = b"artifact file content"
-    presigned_url = "https://storage.example.com/presigned/artifact.tiff"
-    checksum = _make_crc32c_checksum(file_content)
-
-    mock_run = Mock()
-    mock_run.get_artifact_download_url.return_value = presigned_url
-
-    artifact = Mock()
-    artifact.name = "result"
-    artifact.output_artifact_id = "artifact-uuid-123"
-    artifact.metadata = {"checksum_base64_crc32c": checksum, "media_type": "image/tiff"}
-
-    progress = DownloadProgress()
-
-    with patch("aignostics.application._download.requests.get") as mock_get:
-        mock_response = Mock()
-        mock_response.__enter__ = Mock(return_value=mock_response)
-        mock_response.__exit__ = Mock(return_value=False)
-        mock_response.raise_for_status = Mock()
-        mock_response.headers = {"content-length": str(len(file_content))}
-        mock_response.iter_content = Mock(return_value=[file_content])
-        mock_get.return_value = mock_response
-
-        with patch("aignostics.application._utils.get_file_extension_for_artifact", return_value=".tiff"):
-            download_item_artifact(progress, mock_run, artifact, tmp_path)
-
-    mock_run.get_artifact_download_url.assert_called_once_with("artifact-uuid-123")
-    assert (tmp_path / "result.tiff").exists()
-
-
-@pytest.mark.unit
-def test_download_item_artifact_no_checksum_raises(tmp_path: Path) -> None:
-    """Test that download_item_artifact raises ValueError when no checksum metadata is present."""
-    mock_run = Mock()
-    artifact = Mock()
-    artifact.name = "result"
-    artifact.metadata = {}  # no checksum
-
-    progress = DownloadProgress()
-
-    with pytest.raises(ValueError, match="No checksum metadata found"):
-        download_item_artifact(progress, mock_run, artifact, tmp_path)
-
-    mock_run.get_artifact_download_url.assert_not_called()
-
-
-@pytest.mark.unit
-def test_download_item_artifact_skips_existing_correct_checksum(tmp_path: Path) -> None:
-    """Test that download_item_artifact skips download when file exists with correct checksum."""
-    file_content = b"existing artifact content"
-    checksum = _make_crc32c_checksum(file_content)
-
-    mock_run = Mock()
-    artifact = Mock()
-    artifact.name = "result"
-    artifact.output_artifact_id = "artifact-uuid-456"
-    artifact.metadata = {"checksum_base64_crc32c": checksum, "media_type": "image/tiff"}
-
-    progress = DownloadProgress()
-
-    with patch("aignostics.application._utils.get_file_extension_for_artifact", return_value=".tiff"):
-        existing_file = tmp_path / "result.tiff"
-        existing_file.write_bytes(file_content)
-
-        download_item_artifact(progress, mock_run, artifact, tmp_path)
-
-    mock_run.get_artifact_download_url.assert_not_called()

tests/aignostics/platform/resources/runs_test.py

@@ -4,9 +4,10 @@
 verifying their functionality for listing, creating, and managing application runs.
 """
 
-from unittest.mock import Mock
+from unittest.mock import MagicMock, Mock, patch
 
 import pytest
+import requests
 from aignx.codegen.api.public_api import PublicApi
 from aignx.codegen.models import (
     InputArtifactCreationRequest,
@@ -703,6 +704,141 @@ def test_run_details_raises_not_found_after_timeout(app_run, mock_api) -> None:
     assert mock_api.get_run_v1_runs_run_id_get.call_count > 1
 
 
+@pytest.fixture
+def mock_serialize(mock_api) -> Mock:
+    """Configure the serialize method on mock_api to return a valid (method, url, headers) tuple.
+
+    Returns:
+        Mock: The configured serialize mock.
+    """
+    serialize = Mock(return_value=("GET", "https://api.example.com/v1/runs/test-run-id/artifacts/art-1/file", {}, None))
+    mock_api._get_artifact_url_v1_runs_run_id_artifacts_artifact_id_file_get_serialize = serialize
+    return serialize
+
+
+@pytest.mark.unit
+def test_get_artifact_download_url_returns_location_on_307(app_run, mock_serialize) -> None:
+    """Test that get_artifact_download_url returns the Location header value on a 307 redirect.
+
+    Args:
+        app_run: Run instance with mock API.
+        mock_serialize: Mock serializer configured on the API.
+    """
+    # Arrange
+    presigned_url = "https://storage.example.com/artifact?sig=abc123"
+    mock_response = MagicMock()
+    mock_response.__enter__ = Mock(return_value=mock_response)
+    mock_response.__exit__ = Mock(return_value=False)
+    mock_response.status_code = requests.codes.temporary_redirect
+    mock_response.headers = {"Location": presigned_url}
+
+    with patch("aignostics.platform.resources.runs.requests.get", return_value=mock_response) as mock_get:
+        # Act
+        result = app_run.get_artifact_download_url("art-1")
+
+    # Assert
+    assert result == presigned_url
+    mock_get.assert_called_once_with(
+        "https://api.example.com/v1/runs/test-run-id/artifacts/art-1/file",
+        headers={},
+        allow_redirects=False,
+        timeout=mock_get.call_args[1]["timeout"],
+    )
+
+@pytest.mark.parametrize(
+    ("status_code", "expected_message"),
+    [
+        (200, "Unexpected status 200 from artifact URL endpoint"),
+        (307, "307 redirect received but Location header is absent"),
+        (404, "Unexpected status 404 from artifact URL endpoint for artifact 'art-1'; expected 307 redirect"),
+    ]
+)
+@pytest.mark.unit
+def test_get_artifact_download_url_errors(app_run, mock_serialize, status_code, expected_message) -> None:
+    """Test that get_artifact_download_url raises RuntimeError when 307 response has no Location header.
+
+    Args:
+        app_run: Run instance with mock API.
+        mock_serialize: Mock serializer configured on the API.
+    """
+    # Arrange
+    mock_response = MagicMock()
+    mock_response.__enter__ = Mock(return_value=mock_response)
+    mock_response.__exit__ = Mock(return_value=False)
+    mock_response.status_code = status_code
+    mock_response.headers = {}  # No Location header
+
+    with patch("aignostics.platform.resources.runs.requests.get", return_value=mock_response):
+        # Act & Assert
+        with pytest.raises(RuntimeError, match=expected_message):
+            app_run.get_artifact_download_url("art-1")
+
+
+@pytest.mark.unit
+def test_get_artifact_download_url_passes_correct_artifact_id(app_run, mock_api) -> None:
+    """Test that get_artifact_download_url passes the artifact_id to the serializer.
+
+    Args:
+        app_run: Run instance with mock API.
+        mock_api: Mock ExternalsApi instance.
+    """
+    # Arrange
+    artifact_id = "specific-artifact-xyz"
+    serialize = Mock(
+        return_value=("GET", "https://api.example.com/v1/runs/test-run-id/artifacts/specific-artifact-xyz/file", {}, None)
+    )
+    mock_api._get_artifact_url_v1_runs_run_id_artifacts_artifact_id_file_get_serialize = serialize
+
+    presigned_url = "https://storage.example.com/file?sig=xyz"
+    mock_response = MagicMock()
+    mock_response.__enter__ = Mock(return_value=mock_response)
+    mock_response.__exit__ = Mock(return_value=False)
+    mock_response.status_code = requests.codes.temporary_redirect
+    mock_response.headers = {"Location": presigned_url}
+
+    with patch("aignostics.platform.resources.runs.requests.get", return_value=mock_response):
+        # Act
+        result = app_run.get_artifact_download_url(artifact_id)
+
+    # Assert
+    assert result == presigned_url
+    serialize.assert_called_once()
+    call_kwargs = serialize.call_args[1]
+    assert call_kwargs["run_id"] == app_run.run_id
+    assert call_kwargs["artifact_id"] == artifact_id
+
+
+@pytest.mark.unit
+def test_get_artifact_download_url_uses_headers_from_serializer(app_run, mock_api) -> None:
+    """Test that get_artifact_download_url passes serializer-provided headers to requests.get.
+
+    Args:
+        app_run: Run instance with mock API.
+        mock_api: Mock ExternalsApi instance.
+    """
+    # Arrange
+    auth_headers = [("Authorization", "Bearer token123"), ("X-Custom", "header-value")]
+    serialize = Mock(
+        return_value=("GET", "https://api.example.com/v1/runs/test-run-id/artifacts/art-1/file", auth_headers, None)
+    )
+    mock_api._get_artifact_url_v1_runs_run_id_artifacts_artifact_id_file_get_serialize = serialize
+
+    mock_response = MagicMock()
+    mock_response.__enter__ = Mock(return_value=mock_response)
+    mock_response.__exit__ = Mock(return_value=False)
+    mock_response.status_code = requests.codes.temporary_redirect
+    mock_response.headers = {"Location": "https://storage.example.com/file"}
+
+    with patch("aignostics.platform.resources.runs.requests.get", return_value=mock_response) as mock_get:
+        # Act
+        app_run.get_artifact_download_url("art-1")
+
+    # Assert
+    call_kwargs = mock_get.call_args[1]
+    assert call_kwargs["headers"] == dict(auth_headers)
+    assert call_kwargs["allow_redirects"] is False
+
+
 @pytest.mark.unit
 def test_run_details_does_not_retry_other_exceptions(app_run, mock_api) -> None:
     """Test that the outer retry does not catch non-NotFoundException errors.