fixup! fixup! fixup! fixup! chore: Use SAMIA run artifact file endpoint to download artifacts

Author: nahua-aignx

src/aignostics/platform/resources/runs.py

@@ -368,10 +368,14 @@ def get_artifact_download_url(self, artifact_id: str) -> str:
             ServiceException: On 5xx server errors (retried automatically).
             RuntimeError: If the redirect Location header is missing.
         """
-        endpoint_url = f"{self._api.api_client.configuration.host}/v1/runs/{self.run_id}/artifacts/{artifact_id}/file"
+        endpoint_url = (
+            f"{self._api.api_client.configuration.host.rstrip('/')}"
+            f"/v1/runs/{self.run_id}/artifacts/{artifact_id}/file"
+        )
+        token_provider = getattr(self._api.api_client.configuration, "token_provider", None) or get_token
 
         def _resolve() -> str:
-            token = get_token()
+            token = token_provider()
             response = requests.get(
                 endpoint_url,
                 headers={"Authorization": f"Bearer {token}", "User-Agent": user_agent()},

tests/aignostics/application/download_test.py

@@ -1,12 +1,18 @@
 """Tests for download utility functions in the application module."""
 
+from collections.abc import Generator
 from pathlib import Path
-from unittest.mock import Mock, patch
+from unittest.mock import MagicMock, Mock, patch
 
 import pytest
 import requests
 
-from aignostics.application._download import download_url_to_file_with_progress, extract_filename_from_url
+from aignostics.application._download import (
+    download_available_items,
+    download_item_artifact,
+    download_url_to_file_with_progress,
+    extract_filename_from_url,
+)
 from aignostics.application._models import DownloadProgress, DownloadProgressState
 
 
@@ -399,168 +405,121 @@ def progress_callback(p: DownloadProgress) -> None:
         mock_get.assert_called_once_with(https_url, stream=True, timeout=60)
 
 
-@pytest.mark.unit
-def test_download_available_items_calls_url_resolver(tmp_path: Path) -> None:
-    """Test that download_available_items uses the get_artifact_download_url callback.
+@pytest.fixture
+def patched_item_and_run() -> Generator[tuple[MagicMock, MagicMock], None, None]:
+    """Fixture providing patched ItemState/ItemOutput enums and a mock run with one configurable artifact.
 
-    Verifies that when a callback is provided, it is called with (run_id, artifact_id)
-    and the resolved URL is passed to download_item_artifact.
+    Yields:
+        tuple[MagicMock, MagicMock]: The mock run and its single mock artifact.
+            Tests configure artifact attributes (output_artifact_id, download_url, metadata) as needed.
     """
-    from unittest.mock import MagicMock
-
-    from aignostics.application._download import download_available_items
-
-    # Build mock artifact
     mock_artifact = MagicMock()
-    mock_artifact.output_artifact_id = "artifact-xyz"
-    mock_artifact.name = "result"
-    mock_artifact.metadata = {"checksum_base64_crc32c": "AAAA", "checksum_crc32c": ""}
-
-    # Build mock item
     mock_item = MagicMock()
     mock_item.external_id = "slide-1"
     mock_item.state = "TERMINATED"
     mock_item.output = "FULL"
     mock_item.output_artifacts = [mock_artifact]
 
-    # Patch enums for comparison
     with (
         patch("aignostics.application._download.ItemState") as mock_item_state,
         patch("aignostics.application._download.ItemOutput") as mock_item_output,
     ):
         mock_item_state.TERMINATED = "TERMINATED"
         mock_item_output.FULL = "FULL"
 
-        # Build mock run
         mock_run = MagicMock()
         mock_run.run_id = "run-123"
         mock_run.results.return_value = [mock_item]
 
-        progress = DownloadProgress()
-        downloaded_items: set[str] = set()
-        resolved_url = "https://storage.googleapis.com/presigned"
-        mock_url_resolver = MagicMock(return_value=resolved_url)
-
-        with patch("aignostics.application._download.download_item_artifact") as mock_download:
-            download_available_items(
-                progress=progress,
-                application_run=mock_run,
-                destination_directory=tmp_path,
-                downloaded_items=downloaded_items,
-                get_artifact_download_url=mock_url_resolver,
-            )
+        yield mock_run, mock_artifact
+
+
+@pytest.mark.unit
+def test_download_available_items_calls_url_resolver(
+    tmp_path: Path, patched_item_and_run: tuple[MagicMock, MagicMock]
+) -> None:
+    """Test that download_available_items uses the get_artifact_download_url callback.
+
+    Verifies that when a callback is provided, it is called with (run_id, artifact_id)
+    and the resolved URL is passed to download_item_artifact.
+    """
+    mock_run, mock_artifact = patched_item_and_run
+    mock_artifact.output_artifact_id = "artifact-xyz"
+    mock_artifact.name = "result"
+    mock_artifact.metadata = {"checksum_base64_crc32c": "AAAA", "checksum_crc32c": ""}
+
+    resolved_url = "https://storage.googleapis.com/presigned"
+    mock_url_resolver = MagicMock(return_value=resolved_url)
+
+    with patch("aignostics.application._download.download_item_artifact") as mock_download:
+        download_available_items(
+            progress=DownloadProgress(),
+            application_run=mock_run,
+            destination_directory=tmp_path,
+            downloaded_items=set(),
+            get_artifact_download_url=mock_url_resolver,
+        )
 
-            mock_url_resolver.assert_called_once_with("run-123", "artifact-xyz")
-            mock_download.assert_called_once()
-            # artifact_download_url is the 3rd positional arg
-            assert mock_download.call_args[0][2] == resolved_url
+        mock_url_resolver.assert_called_once_with("run-123", "artifact-xyz")
+        mock_download.assert_called_once()
+        assert mock_download.call_args[0][2] == resolved_url  # artifact_download_url is 3rd positional arg
 
 
 @pytest.mark.unit
-def test_download_available_items_falls_back_to_download_url(tmp_path: Path) -> None:
+def test_download_available_items_falls_back_to_download_url(
+    tmp_path: Path, patched_item_and_run: tuple[MagicMock, MagicMock]
+) -> None:
     """Test that download_available_items falls back to artifact.download_url when no callback.
 
     Verifies the deprecated fallback path when get_artifact_download_url is None.
     """
-    from unittest.mock import MagicMock
-
-    from aignostics.application._download import download_available_items
-
-    mock_artifact = MagicMock()
+    mock_run, mock_artifact = patched_item_and_run
     mock_artifact.output_artifact_id = "artifact-xyz"
     mock_artifact.name = "result"
     mock_artifact.download_url = "https://old-signed-url.com/file"
     mock_artifact.metadata = {"checksum_base64_crc32c": "AAAA"}
 
-    mock_item = MagicMock()
-    mock_item.external_id = "slide-1"
-    mock_item.state = "TERMINATED"
-    mock_item.output = "FULL"
-    mock_item.output_artifacts = [mock_artifact]
-
-    with (
-        patch("aignostics.application._download.ItemState") as mock_item_state,
-        patch("aignostics.application._download.ItemOutput") as mock_item_output,
-    ):
-        mock_item_state.TERMINATED = "TERMINATED"
-        mock_item_output.FULL = "FULL"
-
-        mock_run = MagicMock()
-        mock_run.run_id = "run-123"
-        mock_run.results.return_value = [mock_item]
-
-        progress = DownloadProgress()
-        downloaded_items: set[str] = set()
-
-        with patch("aignostics.application._download.download_item_artifact") as mock_download:
-            download_available_items(
-                progress=progress,
-                application_run=mock_run,
-                destination_directory=tmp_path,
-                downloaded_items=downloaded_items,
-                get_artifact_download_url=None,
-            )
+    with patch("aignostics.application._download.download_item_artifact") as mock_download:
+        download_available_items(
+            progress=DownloadProgress(),
+            application_run=mock_run,
+            destination_directory=tmp_path,
+            downloaded_items=set(),
+            get_artifact_download_url=None,
+        )
 
-            mock_download.assert_called_once()
-            assert mock_download.call_args[0][2] == "https://old-signed-url.com/file"
+        mock_download.assert_called_once()
+        assert mock_download.call_args[0][2] == "https://old-signed-url.com/file"
 
 
 @pytest.mark.unit
-def test_download_available_items_skips_when_no_url_available(tmp_path: Path) -> None:
+def test_download_available_items_skips_when_no_url_available(
+    tmp_path: Path, patched_item_and_run: tuple[MagicMock, MagicMock]
+) -> None:
     """Test that download_available_items skips artifacts with no URL available.
 
     When get_artifact_download_url is None and artifact.download_url is also None,
     the artifact should be skipped.
     """
-    from unittest.mock import MagicMock
-
-    from aignostics.application._download import download_available_items
-
-    mock_artifact = MagicMock()
+    mock_run, mock_artifact = patched_item_and_run
     mock_artifact.output_artifact_id = None
-    mock_artifact.name = "result"
     mock_artifact.download_url = None
-    mock_artifact.metadata = {"checksum_base64_crc32c": "AAAA"}
-
-    mock_item = MagicMock()
-    mock_item.external_id = "slide-1"
-    mock_item.state = "TERMINATED"
-    mock_item.output = "FULL"
-    mock_item.output_artifacts = [mock_artifact]
-
-    with (
-        patch("aignostics.application._download.ItemState") as mock_item_state,
-        patch("aignostics.application._download.ItemOutput") as mock_item_output,
-    ):
-        mock_item_state.TERMINATED = "TERMINATED"
-        mock_item_output.FULL = "FULL"
 
-        mock_run = MagicMock()
-        mock_run.run_id = "run-123"
-        mock_run.results.return_value = [mock_item]
-
-        progress = DownloadProgress()
-        downloaded_items: set[str] = set()
-
-        with patch("aignostics.application._download.download_item_artifact") as mock_download:
-            download_available_items(
-                progress=progress,
-                application_run=mock_run,
-                destination_directory=tmp_path,
-                downloaded_items=downloaded_items,
-                get_artifact_download_url=None,
-            )
+    with patch("aignostics.application._download.download_item_artifact") as mock_download:
+        download_available_items(
+            progress=DownloadProgress(),
+            application_run=mock_run,
+            destination_directory=tmp_path,
+            downloaded_items=set(),
+            get_artifact_download_url=None,
+        )
 
-            mock_download.assert_not_called()
+        mock_download.assert_not_called()
 
 
 @pytest.mark.unit
 def test_download_item_artifact_uses_provided_url(tmp_path: Path) -> None:
     """Test that download_item_artifact uses the explicitly provided artifact_download_url."""
-    from unittest.mock import MagicMock
-
-    from aignostics.application._download import download_item_artifact
-
     mock_artifact = MagicMock()
     mock_artifact.name = "cell_classification"
     mock_artifact.metadata = {"checksum_base64_crc32c": "AAAA"}