diff --git a/src/aignostics/application/_download.py b/src/aignostics/application/_download.py index acdd4b312..7a441af65 100644 --- a/src/aignostics/application/_download.py +++ b/src/aignostics/application/_download.py @@ -10,7 +10,7 @@ import requests from loguru import logger -from aignostics.platform import ItemOutput, ItemState, Run, generate_signed_url +from aignostics.platform import ItemOutput, ItemState, OutputArtifactElement, Run, generate_signed_url from aignostics.utils import sanitize_path_component from ._models import DownloadProgress, DownloadProgressState @@ -146,6 +146,44 @@ def update_progress( download_progress_queue.put_nowait(progress) +def _resolve_artifact_url( + artifact: OutputArtifactElement, + run_id: str, + get_artifact_download_url: Callable[[str, str], str] | None, +) -> str | None: + """Resolve the download URL for an artifact. + + Tries the new file endpoint first (via get_artifact_download_url), then falls back to + the deprecated download_url field on the artifact if the endpoint fails. + + Args: + artifact: The artifact object with optional output_artifact_id and download_url. + run_id: The run ID, passed to get_artifact_download_url. + get_artifact_download_url: Callable for the new endpoint, or None to skip. + + Returns: + str | None: The resolved download URL, or None if unavailable. + + Raises: + Exception: Re-raises if the new endpoint fails and no fallback download_url exists. + """ + if get_artifact_download_url and artifact.output_artifact_id: + try: + return get_artifact_download_url(run_id, artifact.output_artifact_id) + except Exception as e: + fallback_url: str | None = getattr(artifact, "download_url", None) + if fallback_url: + logger.warning( + "Failed to resolve download URL via file endpoint for artifact {} ({}). " + "Falling back to deprecated download_url field.", + artifact.output_artifact_id, + e, + ) + return fallback_url + raise + return getattr(artifact, "download_url", None) + + def download_available_items( # noqa: PLR0913, PLR0917 progress: DownloadProgress, application_run: Run, @@ -154,6 +192,7 @@ def download_available_items( # noqa: PLR0913, PLR0917 create_subdirectory_per_item: bool = False, download_progress_queue: Any | None = None, # noqa: ANN401 download_progress_callable: Callable | None = None, # type: ignore[type-arg] + get_artifact_download_url: Callable[[str, str], str] | None = None, ) -> None: """Download items that are available and not yet downloaded. @@ -165,6 +204,9 @@ def download_available_items( # noqa: PLR0913, PLR0917 create_subdirectory_per_item (bool): Whether to create a subdirectory for each item. download_progress_queue (Queue | None): Queue for GUI progress updates. download_progress_callable (Callable | None): Callback for CLI progress updates. + get_artifact_download_url (Callable[[str, str], str] | None): Callback that takes + (run_id, artifact_id) and returns a presigned download URL. If None, falls back + to artifact.download_url (deprecated). """ items = list(application_run.results()) progress.item_count = len(items) @@ -201,9 +243,16 @@ def download_available_items( # noqa: PLR0913, PLR0917 progress.artifact = artifact update_progress(progress, download_progress_callable, download_progress_queue) + # Resolve artifact download URL via the new file endpoint, with fallback + artifact_url = _resolve_artifact_url(artifact, application_run.run_id, get_artifact_download_url) + if not artifact_url: + logger.warning("No download URL available for artifact {}", artifact.output_artifact_id) + continue + download_item_artifact( progress, artifact, + artifact_url, item_directory, item.external_id if not create_subdirectory_per_item else "", download_progress_queue, @@ -216,6 +265,7 @@ def download_available_items( # noqa: PLR0913, PLR0917 def download_item_artifact( # noqa: PLR0913, PLR0917 progress: DownloadProgress, artifact: Any, # noqa: ANN401 + artifact_download_url: str, destination_directory: Path, prefix: str = "", download_progress_queue: Any | None = None, # noqa: ANN401 @@ -226,6 +276,7 @@ def download_item_artifact( # noqa: PLR0913, PLR0917 Args: progress (DownloadProgress): Progress tracking object for GUI or CLI updates. artifact (Any): The artifact to download. + artifact_download_url (str): The presigned URL to download the artifact from. destination_directory (Path): Directory to save the file. prefix (str): Prefix for the file name, if needed. download_progress_queue (Queue | None): Queue for GUI progress updates. @@ -260,7 +311,7 @@ def download_item_artifact( # noqa: PLR0913, PLR0917 download_file_with_progress( progress, - artifact.download_url, + artifact_download_url, artifact_path, metadata_checksum, download_progress_queue, diff --git a/src/aignostics/application/_gui/_page_application_run_describe.py b/src/aignostics/application/_gui/_page_application_run_describe.py index 740aae23a..370743b7d 100644 --- a/src/aignostics/application/_gui/_page_application_run_describe.py +++ b/src/aignostics/application/_gui/_page_application_run_describe.py @@ -14,8 +14,9 @@ ui, # noq ) from nicegui import run as nicegui_run +from nicegui.events import ClickEventArguments -from aignostics.platform import ItemOutput, ItemResult, ItemState, RunState +from aignostics.platform import ItemOutput, ItemResult, ItemState, Run, RunState from aignostics.third_party.showinfm.showinfm import show_in_file_manager from aignostics.utils import GUILocalFilePicker, get_user_data_directory @@ -338,9 +339,9 @@ async def start_download() -> None: # Activate the timer now that download is starting progress_timer.activate() + download_button.disable() + download_button.props(add="loading") try: - download_button.disable() - download_button.props(add="loading") results_folder = await nicegui_run.cpu_bound( Service.application_run_download_static, run_id=run.run_id, @@ -364,19 +365,24 @@ async def start_download() -> None: else: ui.notify("Download completed.", type="positive") show_in_file_manager(str(results_folder)) - except ValueError as e: + except Exception as e: + logger.exception( + "Download failed for run {} (qupath_project={}, marimo={}, folder={})", + run.run_id, + current_qupath_project, + current_marimo, + current_folder, + ) ui.notify(f"Download failed: {e}", type="negative", multi_line=True) + finally: progress_timer.deactivate() progress_state["queue"] = None - return - progress_timer.deactivate() - progress_state["queue"] = None - download_button.props(remove="loading") - download_button.enable() - download_item_status.set_visibility(False) - download_item_progress.set_visibility(False) - download_artifact_status.set_visibility(False) - download_artifact_progress.set_visibility(False) + download_button.props(remove="loading") + download_button.enable() + download_item_status.set_visibility(False) + download_item_progress.set_visibility(False) + download_artifact_status.set_visibility(False) + download_artifact_progress.set_visibility(False) ui.separator() with ui.row(align_items="end").classes("w-full justify-end"): @@ -781,29 +787,85 @@ def render_item(item: ItemResult) -> None: # noqa: C901, PLR0912, PLR0915 icon=mime_type_to_icon(mime_type), group="artifacts", ).classes("w-full"): - if artifact.download_url: - url = artifact.download_url + if artifact.output_artifact_id: + artifact_id = artifact.output_artifact_id title = artifact.name metadata = artifact.metadata + with ui.button_group(): if mime_type == "image/tiff": - ui.button( + preview_button = ui.button( "Preview", icon=mime_type_to_icon(mime_type), - on_click=lambda _, url=url, title=title: tiff_dialog_open(title, url), ) + + async def _preview_tiff( + _: ClickEventArguments, + aid: str = artifact_id, + t: str = title, + _run: Run = run, + _btn: ui.button = preview_button, + ) -> None: + try: + _btn.props(add="loading") + url = await nicegui_run.io_bound( + _run.get_artifact_download_url, aid + ) + tiff_dialog_open(t, url) + except Exception as e: + ui.notify(f"Failed to resolve preview URL: {e}", type="warning") + finally: + _btn.props(remove="loading") + + preview_button.on_click(_preview_tiff) + if mime_type == "text/csv": - ui.button( + preview_button = ui.button( "Preview", icon=mime_type_to_icon(mime_type), - on_click=lambda _, url=url, title=title: csv_dialog_open(title, url), - ) - if url: - ui.button( - text="Download", - icon="cloud_download", - on_click=lambda _, url=url: webbrowser.open(url), ) + + async def _preview_csv( + _: ClickEventArguments, + aid: str = artifact_id, + t: str = title, + _run: Run = run, + _btn: ui.button = preview_button, + ) -> None: + try: + _btn.props(add="loading") + url = await nicegui_run.io_bound( + _run.get_artifact_download_url, aid + ) + csv_dialog_open(t, url) + except Exception as e: + ui.notify(f"Failed to resolve preview URL: {e}", type="warning") + finally: + _btn.props(remove="loading") + + preview_button.on_click(_preview_csv) + + artifact_dl_button = ui.button( + text="Download", + icon="cloud_download", + ) + + async def _download_artifact( + _: ClickEventArguments, + aid: str = artifact_id, + _run: Run = run, + _btn: ui.button = artifact_dl_button, + ) -> None: + try: + _btn.props(add="loading") + url = await nicegui_run.io_bound(_run.get_artifact_download_url, aid) + webbrowser.open(url) + except Exception as e: + ui.notify(f"Failed to resolve download URL: {e}", type="warning") + finally: + _btn.props(remove="loading") + + artifact_dl_button.on_click(_download_artifact) if metadata: ui.button( text="Schema", diff --git a/src/aignostics/application/_service.py b/src/aignostics/application/_service.py index 4c6c8556d..3b2531746 100644 --- a/src/aignostics/application/_service.py +++ b/src/aignostics/application/_service.py @@ -1483,6 +1483,19 @@ def update_qupath_add_input_progress(qupath_add_input_progress: QuPathAddProgres update_progress(progress, download_progress_callable, download_progress_queue) downloaded_items: set[str] = set() # Track downloaded items to avoid re-downloading + + def _get_artifact_download_url(run_id: str, artifact_id: str) -> str: + """Resolve artifact download URL via the new API endpoint. + + Args: + run_id (str): The run ID. + artifact_id (str): The artifact ID. + + Returns: + str: The presigned download URL. + """ + return self._get_platform_client().run(run_id).get_artifact_download_url(artifact_id) + while True: run_details = application_run.details() # (Re)load current run details progress.run = run_details @@ -1496,6 +1509,7 @@ def update_qupath_add_input_progress(qupath_add_input_progress: QuPathAddProgres create_subdirectory_per_item, download_progress_queue, download_progress_callable, + get_artifact_download_url=_get_artifact_download_url, ) if run_details.state == RunState.TERMINATED: diff --git a/src/aignostics/platform/__init__.py b/src/aignostics/platform/__init__.py index 9e4cc47fc..f369e2b56 100644 --- a/src/aignostics/platform/__init__.py +++ b/src/aignostics/platform/__init__.py @@ -91,7 +91,7 @@ get_mime_type_for_artifact, mime_type_to_file_ending, ) -from .resources.runs import LIST_APPLICATION_RUNS_MAX_PAGE_SIZE, LIST_APPLICATION_RUNS_MIN_PAGE_SIZE, Run +from .resources.runs import LIST_APPLICATION_RUNS_MAX_PAGE_SIZE, LIST_APPLICATION_RUNS_MIN_PAGE_SIZE, Artifact, Run __all__ = [ "API_ROOT_DEV", @@ -146,6 +146,7 @@ "Application", "ApplicationSummary", "ApplicationVersion", + "Artifact", "Client", "InputArtifact", "InputArtifactData", diff --git a/src/aignostics/platform/resources/runs.py b/src/aignostics/platform/resources/runs.py index dc6dca005..505265311 100644 --- a/src/aignostics/platform/resources/runs.py +++ b/src/aignostics/platform/resources/runs.py @@ -8,12 +8,14 @@ import time import typing as t from collections.abc import Iterator +from http import HTTPStatus from pathlib import Path from time import sleep from typing import Any, cast +import requests from aignx.codegen.api.public_api import PublicApi -from aignx.codegen.exceptions import NotFoundException, ServiceException +from aignx.codegen.exceptions import ApiException, NotFoundException, ServiceException from aignx.codegen.models import ( CustomMetadataUpdateRequest, ItemCreationRequest, @@ -49,6 +51,7 @@ from urllib3.exceptions import IncompleteRead, PoolError, ProtocolError, ProxyError from urllib3.exceptions import TimeoutError as Urllib3TimeoutError +from aignostics.platform._authentication import get_token from aignostics.platform._operation_cache import cached_operation, operation_cache_clear from aignostics.platform._sdk_metadata import ( build_item_sdk_metadata, @@ -105,6 +108,145 @@ class DownloadTimeoutError(RuntimeError): """Exception raised when the download operation exceeds its timeout.""" +class Artifact: + """Represents a single output artifact belonging to a run. + + Provides operations to download the artifact or resolve its presigned download URL. + """ + + def __init__(self, api: PublicApi, run_id: str, artifact_id: str) -> None: + """Initializes an Artifact instance. + + Args: + api (PublicApi): The configured API client. + run_id (str): The ID of the parent run. + artifact_id (str): The ID of the artifact. + """ + self._api = api + self.run_id = run_id + self.artifact_id = artifact_id + + @staticmethod + def _fetch_redirect_url( + endpoint_url: str, + token: str, + proxy: str | None, + ssl_verify: bool | str, + artifact_id: str, + ) -> str: + """Execute the HTTP request to the artifact file endpoint and return the redirect URL. + + Args: + endpoint_url (str): Full URL of the artifact file endpoint. + token (str): Bearer token for authorization. + proxy (str | None): Optional proxy URL. + ssl_verify (bool | str): SSL verification setting (True/False or CA bundle path). + artifact_id (str): Artifact ID used in error messages. + + Returns: + str: The presigned URL from the redirect Location header. + + Raises: + NotFoundException: If the artifact is not found (404). + ServiceException: On network errors or 5xx responses. + ApiException: On 4xx client errors. + RuntimeError: If the redirect has no Location header or returns an unexpected status. + """ + try: + response = requests.get( + endpoint_url, + headers={"Authorization": f"Bearer {token}", "User-Agent": user_agent()}, + allow_redirects=False, + timeout=settings().run_timeout, + proxies={"http": proxy, "https": proxy} if proxy else None, + verify=ssl_verify, + stream=True, + ) + response.close() + except requests.Timeout as e: + raise ServiceException(status=HTTPStatus.SERVICE_UNAVAILABLE, reason="Request timed out") from e + except requests.ConnectionError as e: + raise ServiceException(status=HTTPStatus.SERVICE_UNAVAILABLE, reason="Connection failed") from e + except requests.RequestException as e: + raise ServiceException(status=HTTPStatus.SERVICE_UNAVAILABLE, reason=f"Request failed: {e}") from e + + if response.status_code in { + HTTPStatus.MOVED_PERMANENTLY, + HTTPStatus.FOUND, + HTTPStatus.TEMPORARY_REDIRECT, + HTTPStatus.PERMANENT_REDIRECT, + }: + location = response.headers.get("Location") + if not location: + msg = f"Redirect response {response.status_code} missing Location header for artifact {artifact_id}" + raise RuntimeError(msg) + return location + if response.status_code == HTTPStatus.NOT_FOUND: + raise NotFoundException(status=HTTPStatus.NOT_FOUND, reason="Artifact not found") + if response.status_code >= HTTPStatus.INTERNAL_SERVER_ERROR: + raise ServiceException(status=response.status_code, reason=response.reason) + if response.status_code >= HTTPStatus.BAD_REQUEST: + raise ApiException(status=response.status_code, reason=response.reason) + msg = f"Unexpected status {response.status_code} from artifact file endpoint for artifact {artifact_id}" + raise RuntimeError(msg) + + def get_download_url(self) -> str: + """Get a presigned download URL for this artifact via the file endpoint. + + Calls GET /api/v1/runs/{run_id}/artifacts/{artifact_id}/file with + allow_redirects=False and extracts the presigned URL from the Location + header of the 307 redirect response. Using the codegen method is not + viable here because urllib3 follows the redirect automatically, fetching + the binary file body instead of the URL. + + Returns: + str: The presigned download URL from the redirect Location header. + + Raises: + NotFoundException: If the artifact is not found (HTTP 404). + ApiException: For other client-error responses (HTTP 4xx) such as invalid parameters or + insufficient permissions. + ServiceException: For server-error responses (HTTP 5xx) or network errors (retried + automatically where transient). + RuntimeError: If the redirect Location header is missing. + """ + configuration = self._api.api_client.configuration + host = configuration.host.rstrip("/") + endpoint_url = f"{host}/api/v1/runs/{self.run_id}/artifacts/{self.artifact_id}/file" + token_provider = getattr(configuration, "token_provider", None) or get_token + proxy = getattr(configuration, "proxy", None) + ssl_ca_cert = getattr(configuration, "ssl_ca_cert", None) + verify_ssl = getattr(configuration, "verify_ssl", True) + ssl_verify: bool | str = ssl_ca_cert or verify_ssl + logger.trace("Resolving download URL for artifact {} via {}", self.artifact_id, endpoint_url) + + url = Retrying( + retry=retry_if_exception_type(exception_types=RETRYABLE_EXCEPTIONS), + stop=stop_after_attempt(settings().run_retry_attempts), + wait=wait_exponential_jitter(initial=settings().run_retry_wait_min, max=settings().run_retry_wait_max), + before_sleep=_log_retry_attempt, + reraise=True, + )(lambda: self._fetch_redirect_url(endpoint_url, token_provider(), proxy, ssl_verify, self.artifact_id)) + logger.trace("Resolved download URL for artifact {}", self.artifact_id) + return url + + def download(self) -> object: + """Download the artifact content via the API client. + + Uses the generated API client which automatically follows the redirect to + the presigned URL and returns the artifact content. + + Returns: + object: The artifact content. + """ + return self._api.get_artifact_url_v1_runs_run_id_artifacts_artifact_id_file_get( + run_id=self.run_id, + artifact_id=self.artifact_id, + _request_timeout=settings().run_timeout, + _headers={"User-Agent": user_agent()}, + ) + + class Run: """Represents a single application run. @@ -345,8 +487,36 @@ def download_to_folder( # noqa: C901 msg = f"Download operation failed unexpectedly for run {self.run_id}: {e}" raise RuntimeError(msg) from e - @staticmethod + def artifact(self, artifact_id: str) -> Artifact: + """Get an Artifact instance for the given artifact ID. + + Args: + artifact_id (str): The output artifact ID. + + Returns: + Artifact: An Artifact instance for downloading content or resolving the presigned URL. + """ + return Artifact(self._api, self.run_id, artifact_id) + + def get_artifact_download_url(self, artifact_id: str) -> str: + """Get a presigned download URL for an artifact via the file endpoint. + + Args: + artifact_id (str): The output artifact ID. + + Returns: + str: The presigned download URL from the redirect Location header. + + Raises: + NotFoundException: If the artifact is not found (HTTP 404). + ApiException: For other client-error responses (HTTP 4xx). + ServiceException: For server-error responses (HTTP 5xx) or network errors. + RuntimeError: If the redirect Location header is missing. + """ + return self.artifact(artifact_id).get_download_url() + def ensure_artifacts_downloaded( + self, base_folder: Path, item: ItemResultReadResponse, checksum_attribute_key: str = "checksum_base64_crc32c", @@ -370,34 +540,41 @@ def ensure_artifacts_downloaded( downloaded_at_least_one_artifact = False for artifact in item.output_artifacts: - if artifact.download_url: - item_dir.mkdir(exist_ok=True, parents=True) - file_ending = mime_type_to_file_ending(get_mime_type_for_artifact(artifact)) - file_path = item_dir / f"{artifact.name}{file_ending}" - if not artifact.metadata: - logger.error( - "Skipping artifact %s for item %s, no metadata present", artifact.name, item.external_id - ) + if not artifact.output_artifact_id: + logger.warning( + "Skipping artifact {} for item {}: missing output_artifact_id", artifact.name, item.external_id + ) + if print_status: print( - f"> Skipping artifact {artifact.name} for item {item.external_id}, no metadata present" - ) if print_status else None - continue - checksum = artifact.metadata[checksum_attribute_key] - - if file_path.exists(): - file_checksum = calculate_file_crc32c(file_path) - if file_checksum != checksum: - logger.trace("Resume download for {} to {}", artifact.name, file_path) - print(f"> Resume download for {artifact.name} to {file_path}") if print_status else None - else: - continue + f"> Skipping artifact {artifact.name} for item {item.external_id}: missing output_artifact_id" + ) + continue + item_dir.mkdir(exist_ok=True, parents=True) + file_ending = mime_type_to_file_ending(get_mime_type_for_artifact(artifact)) + file_path = item_dir / f"{artifact.name}{file_ending}" + if not artifact.metadata: + logger.error("Skipping artifact %s for item %s, no metadata present", artifact.name, item.external_id) + print( + f"> Skipping artifact {artifact.name} for item {item.external_id}, no metadata present" + ) if print_status else None + continue + checksum = artifact.metadata[checksum_attribute_key] + + if file_path.exists(): + file_checksum = calculate_file_crc32c(file_path) + if file_checksum != checksum: + logger.trace("Resume download for {} to {}", artifact.name, file_path) + print(f"> Resume download for {artifact.name} to {file_path}") if print_status else None else: - downloaded_at_least_one_artifact = True - logger.trace("Download for {} to {}", artifact.name, file_path) - print(f"> Download for {artifact.name} to {file_path}") if print_status else None + continue + else: + logger.trace("Download for {} to {}", artifact.name, file_path) + print(f"> Download for {artifact.name} to {file_path}") if print_status else None - # if file is not there at all or only partially downloaded yet - download_file(artifact.download_url, str(file_path), checksum) + # Resolve presigned URL via the artifact file endpoint and download + downloaded_at_least_one_artifact = True + artifact_url = self.get_artifact_download_url(artifact.output_artifact_id) + download_file(artifact_url, str(file_path), checksum) if downloaded_at_least_one_artifact: logger.trace("Downloaded results for item: {} to {}", item.external_id, item_dir) diff --git a/tests/aignostics/application/download_test.py b/tests/aignostics/application/download_test.py index e5537ef62..f97427dbf 100644 --- a/tests/aignostics/application/download_test.py +++ b/tests/aignostics/application/download_test.py @@ -1,14 +1,28 @@ """Tests for download utility functions in the application module.""" +import base64 +from collections.abc import Generator from pathlib import Path -from unittest.mock import Mock, patch +from unittest.mock import MagicMock, Mock, patch +import crc32c import pytest import requests -from aignostics.application._download import download_url_to_file_with_progress, extract_filename_from_url +from aignostics.application._download import ( + _resolve_artifact_url, + download_available_items, + download_file_with_progress, + download_item_artifact, + download_url_to_file_with_progress, + extract_filename_from_url, +) from aignostics.application._models import DownloadProgress, DownloadProgressState +_PATCH_DOWNLOAD_ITEM_ARTIFACT = "aignostics.application._download.download_item_artifact" +_OLD_SIGNED_URL = "https://old-signed-url.com/file" +_ERROR_403_FORBIDDEN = "403 Forbidden" + @pytest.mark.unit def test_extract_filename_from_url_gs() -> None: @@ -397,3 +411,355 @@ 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) + + +@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. + + Yields: + tuple[MagicMock, MagicMock]: The mock run and its single mock artifact. + Tests configure artifact attributes (output_artifact_id, download_url, metadata) as needed. + """ + mock_artifact = MagicMock() + 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] + + 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(_PATCH_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() + 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, 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. + """ + mock_run, mock_artifact = patched_item_and_run + mock_artifact.output_artifact_id = "artifact-xyz" + mock_artifact.name = "result" + mock_artifact.download_url = _OLD_SIGNED_URL + mock_artifact.metadata = {"checksum_base64_crc32c": "AAAA"} + + with patch(_PATCH_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] == _OLD_SIGNED_URL + + +@pytest.mark.unit +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. + """ + mock_run, mock_artifact = patched_item_and_run + mock_artifact.output_artifact_id = None + mock_artifact.download_url = None + + with patch(_PATCH_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() + + +@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.""" + mock_artifact = MagicMock() + mock_artifact.name = "cell_classification" + mock_artifact.metadata = {"checksum_base64_crc32c": "AAAA"} + + progress = DownloadProgress() + artifact_url = "https://storage.googleapis.com/presigned-url" + + with ( + patch("aignostics.application._download.get_file_extension_for_artifact", return_value=".csv"), + patch("aignostics.application._download.download_file_with_progress") as mock_download, + ): + download_item_artifact( + progress=progress, + artifact=mock_artifact, + artifact_download_url=artifact_url, + destination_directory=tmp_path, + ) + + mock_download.assert_called_once() + call_args = mock_download.call_args[0] + assert call_args[0] is progress + assert call_args[1] == artifact_url + assert call_args[3] == "AAAA" + + +@pytest.mark.unit +def test_resolve_artifact_url_uses_new_endpoint() -> None: + """Test that _resolve_artifact_url calls get_artifact_download_url when output_artifact_id is set.""" + mock_artifact = MagicMock() + mock_artifact.output_artifact_id = "artifact-abc" + mock_artifact.download_url = None + + resolver = MagicMock(return_value="https://new-endpoint.com/presigned") + result = _resolve_artifact_url(mock_artifact, "run-123", resolver) + + assert result == "https://new-endpoint.com/presigned" + resolver.assert_called_once_with("run-123", "artifact-abc") + + +@pytest.mark.unit +def test_resolve_artifact_url_falls_back_to_download_url_on_error() -> None: + """Test that _resolve_artifact_url falls back to download_url when the new endpoint raises.""" + mock_artifact = MagicMock() + mock_artifact.output_artifact_id = "artifact-abc" + mock_artifact.download_url = _OLD_SIGNED_URL + + resolver = MagicMock(side_effect=RuntimeError(_ERROR_403_FORBIDDEN)) + result = _resolve_artifact_url(mock_artifact, "run-123", resolver) + + assert result == _OLD_SIGNED_URL + + +@pytest.mark.unit +def test_resolve_artifact_url_reraises_when_no_fallback_available() -> None: + """Test that _resolve_artifact_url re-raises if the endpoint fails and download_url is absent.""" + mock_artifact = MagicMock() + mock_artifact.output_artifact_id = "artifact-abc" + mock_artifact.download_url = None + + resolver = MagicMock(side_effect=RuntimeError(_ERROR_403_FORBIDDEN)) + + with pytest.raises(RuntimeError, match=_ERROR_403_FORBIDDEN): + _resolve_artifact_url(mock_artifact, "run-123", resolver) + + +@pytest.mark.unit +def test_download_available_items_skips_already_downloaded_items( + tmp_path: Path, patched_item_and_run: tuple[MagicMock, MagicMock] +) -> None: + """Test that items already in downloaded_items are skipped without downloading.""" + mock_run, mock_artifact = patched_item_and_run + mock_artifact.output_artifact_id = "artifact-xyz" + + with patch(_PATCH_DOWNLOAD_ITEM_ARTIFACT) as mock_download: + download_available_items( + progress=DownloadProgress(), + application_run=mock_run, + destination_directory=tmp_path, + downloaded_items={"slide-1"}, # "slide-1" is already downloaded + get_artifact_download_url=None, + ) + + mock_download.assert_not_called() + + +@pytest.mark.unit +def test_download_available_items_with_create_subdirectory_per_item( + tmp_path: Path, patched_item_and_run: tuple[MagicMock, MagicMock] +) -> None: + """Test that a subdirectory is created per item when create_subdirectory_per_item=True.""" + mock_run, mock_artifact = patched_item_and_run + mock_artifact.output_artifact_id = "artifact-xyz" + + with patch(_PATCH_DOWNLOAD_ITEM_ARTIFACT): + download_available_items( + progress=DownloadProgress(), + application_run=mock_run, + destination_directory=tmp_path, + downloaded_items=set(), + create_subdirectory_per_item=True, + get_artifact_download_url=None, + ) + + # Item external_id is "slide-1", stem is "slide-1" + assert (tmp_path / "slide-1").is_dir() + + +@pytest.mark.unit +def test_download_url_to_file_with_progress_network_error(tmp_path: Path) -> None: + """Test that network (non-HTTP) errors are wrapped in RuntimeError.""" + url = "https://example.com/slide.tiff" + destination = tmp_path / "slide.tiff" + progress = DownloadProgress() + + with patch("aignostics.application._download.requests.get") as mock_get: + mock_get.side_effect = requests.ConnectionError("Network unreachable") + + with pytest.raises(RuntimeError, match="Network error downloading"): + download_url_to_file_with_progress(progress, url, destination) + + assert not destination.exists() + + +@pytest.mark.unit +def test_download_item_artifact_no_checksum_raises(tmp_path: Path) -> None: + """Test that ValueError is raised when no checksum metadata is found for an artifact.""" + mock_artifact = MagicMock() + mock_artifact.name = "result.csv" + mock_artifact.metadata = {} # No checksum fields + + with pytest.raises(ValueError, match="No checksum metadata found"): + download_item_artifact( + progress=DownloadProgress(), + artifact=mock_artifact, + artifact_download_url="https://example.com/result.csv", + destination_directory=tmp_path, + ) + + +@pytest.mark.unit +def test_download_item_artifact_file_exists_correct_checksum_skips_download(tmp_path: Path) -> None: + """Test that download is skipped when the file already exists with the correct CRC32C checksum.""" + test_content = b"test file content for checksum" + artifact_path = tmp_path / "result.csv" + artifact_path.write_bytes(test_content) + + hasher = crc32c.CRC32CHash() + hasher.update(test_content) + correct_checksum = base64.b64encode(hasher.digest()).decode("ascii") + + mock_artifact = MagicMock() + mock_artifact.name = "result" + mock_artifact.metadata = {"checksum_base64_crc32c": correct_checksum} + + with ( + patch("aignostics.application._download.get_file_extension_for_artifact", return_value=".csv"), + patch("aignostics.application._download.download_file_with_progress") as mock_download, + ): + download_item_artifact( + progress=DownloadProgress(), + artifact=mock_artifact, + artifact_download_url="https://example.com/result.csv", + destination_directory=tmp_path, + ) + + mock_download.assert_not_called() + + +@pytest.mark.unit +def test_download_file_with_progress_success(tmp_path: Path) -> None: + """Test that download_file_with_progress downloads a file and verifies its CRC32C checksum.""" + signed_url = "https://example.com/artifact.csv" + artifact_path = tmp_path / "result.csv" + file_content = b"test artifact content for crc32c" + + hasher = crc32c.CRC32CHash() + hasher.update(file_content) + expected_checksum = base64.b64encode(hasher.digest()).decode("ascii") + + progress = DownloadProgress() + progress.artifact = MagicMock() + progress.artifact.name = "result" + progress.item_external_id = "slide-1" + + with patch("aignostics.application._download.requests.get") as mock_get: + mock_response = MagicMock() + 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 + + download_file_with_progress( + progress=progress, + signed_url=signed_url, + artifact_path=artifact_path, + metadata_checksum=expected_checksum, + ) + + assert artifact_path.exists() + assert artifact_path.read_bytes() == file_content + + +@pytest.mark.unit +def test_download_file_with_progress_checksum_mismatch_removes_file_and_raises(tmp_path: Path) -> None: + """Test that a CRC32C checksum mismatch causes the file to be removed and raises ValueError.""" + signed_url = "https://example.com/artifact.csv" + artifact_path = tmp_path / "result.csv" + file_content = b"test artifact content" + wrong_checksum = "AAAAAAAAAAAAAAAA==" + + progress = DownloadProgress() + progress.artifact = MagicMock() + progress.artifact.name = "result" + progress.item_external_id = "slide-1" + + with patch("aignostics.application._download.requests.get") as mock_get: + mock_response = MagicMock() + 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 pytest.raises(ValueError, match="Checksum mismatch"): + download_file_with_progress( + progress=progress, + signed_url=signed_url, + artifact_path=artifact_path, + metadata_checksum=wrong_checksum, + ) + + assert not artifact_path.exists() diff --git a/tests/aignostics/application/service_test.py b/tests/aignostics/application/service_test.py index d358554f9..6e4cda7ca 100644 --- a/tests/aignostics/application/service_test.py +++ b/tests/aignostics/application/service_test.py @@ -1,6 +1,7 @@ """Tests to verify the service functionality of the application module.""" from datetime import UTC, datetime, timedelta +from pathlib import Path from unittest.mock import MagicMock, patch import pytest @@ -503,3 +504,57 @@ def test_application_run_update_item_custom_metadata_not_found(mock_get_client: with pytest.raises(NotFoundException, match="not found"): service.application_run_update_item_custom_metadata("run-123", "invalid-item-id", {"key": "value"}) + + +@pytest.mark.unit +@patch("aignostics.application._service.Service._get_platform_client") +@patch("aignostics.application._service.download_available_items") +def test_application_run_download_artifact_url_closure_delegates_to_platform_client( + mock_download_available_items: MagicMock, + mock_get_client: MagicMock, + tmp_path: Path, +) -> None: + """Test that the _get_artifact_download_url closure delegates to the platform client.""" + mock_run_details = MagicMock() + mock_run_details.run_id = "run-123" + mock_run_details.state = MagicMock() # Non-TERMINATED, triggers wait_for_completion=False break + mock_run_details.error_message = None + mock_run_details.error_code = None + + mock_platform_run = MagicMock() + mock_platform_run.run_id = "run-123" + mock_platform_run.details.return_value = mock_run_details + mock_platform_run.results.return_value = [] + mock_platform_run.get_artifact_download_url.return_value = "https://presigned.example.com/artifact.tiff" + + mock_client = MagicMock() + mock_client.run.return_value = mock_platform_run + mock_get_client.return_value = mock_client + + # Capture the get_artifact_download_url callback passed to download_available_items + captured_callback: list[object] = [] + + def capture_fn(*args: object, **kwargs: object) -> None: + fn = kwargs.get("get_artifact_download_url") + if fn is not None: + captured_callback.append(fn) + + mock_download_available_items.side_effect = capture_fn + + service = ApplicationService() + service.application_run_download( + run_id="run-123", + destination_directory=tmp_path, + wait_for_completion=False, + create_subdirectory_for_run=False, + ) + + # Verify the callback was passed to download_available_items + assert len(captured_callback) == 1 + callback = captured_callback[0] + + # Call the closure to exercise line 1497 and verify it delegates correctly + result = callback("run-123", "artifact-456") # type: ignore[operator] + + mock_platform_run.get_artifact_download_url.assert_called_with("artifact-456") + assert result == "https://presigned.example.com/artifact.tiff" diff --git a/tests/aignostics/platform/resources/runs_test.py b/tests/aignostics/platform/resources/runs_test.py index c8367196e..5ec5f4de3 100644 --- a/tests/aignostics/platform/resources/runs_test.py +++ b/tests/aignostics/platform/resources/runs_test.py @@ -4,7 +4,7 @@ verifying their functionality for listing, creating, and managing application runs. """ -from unittest.mock import Mock +from unittest.mock import Mock, patch import pytest from aignx.codegen.api.public_api import PublicApi @@ -16,9 +16,18 @@ RunReadResponse, ) -from aignostics.platform.resources.runs import LIST_APPLICATION_RUNS_MAX_PAGE_SIZE, Run, Runs +from aignostics.platform.resources.runs import LIST_APPLICATION_RUNS_MAX_PAGE_SIZE, Artifact, Run, Runs from aignostics.platform.resources.utils import PAGE_SIZE +_PLATFORM_HOST = "https://platform.aignostics.com" +_PATCH_GET_AUTH = "aignostics.platform.resources.runs.get_token" +_PATCH_REQUESTS_GET = "aignostics.platform.resources.runs.requests.get" +_PATCH_SETTINGS = "aignostics.platform.resources.runs.settings" +_PATCH_DOWNLOAD_FILE = "aignostics.platform.resources.runs.download_file" +_PRESIGNED_URL = "https://storage.googleapis.com/presigned-url" +_PROXY_URL = "https://corporate-proxy:8080" +_MIME_TYPE_CSV = "text/csv" + @pytest.fixture def mock_api() -> Mock: @@ -56,6 +65,19 @@ def app_run(mock_api) -> Run: return Run(mock_api, "test-run-id") +@pytest.fixture +def artifact_instance(mock_api) -> Artifact: + """Create an Artifact instance with a mock API for testing. + + Args: + mock_api: A mock instance of ExternalsApi. + + Returns: + Artifact: An Artifact instance using the mock API. + """ + return Artifact(mock_api, "test-run-id", "artifact-123") + + @pytest.mark.unit def test_runs_list_with_pagination(runs, mock_api) -> None: """Test that Runs.list() correctly handles pagination. @@ -722,3 +744,466 @@ def test_run_details_does_not_retry_other_exceptions(app_run, mock_api) -> None: app_run.details() assert mock_api.get_run_v1_runs_run_id_get.call_count == 1 + + +@pytest.mark.unit +def test_artifact_get_download_url_returns_presigned_url(artifact_instance, mock_api) -> None: + """Test that Artifact.get_download_url extracts the presigned URL from a 307 redirect. + + The implementation calls GET /api/v1/runs/{run_id}/artifacts/{artifact_id}/file with + allow_redirects=False and returns the Location header from the 307 response. + + Args: + artifact_instance: Artifact instance with mock API. + mock_api: Mock ExternalsApi instance. + """ + expected_url = "https://storage.googleapis.com/presigned-url?token=abc" + mock_api.api_client = Mock() + mock_api.api_client.configuration.host = _PLATFORM_HOST + mock_api.api_client.configuration.token_provider = None + + mock_response = Mock() + mock_response.status_code = 307 + mock_response.headers = {"Location": expected_url} + + with ( + patch(_PATCH_GET_AUTH, return_value="test-token"), + patch(_PATCH_REQUESTS_GET, return_value=mock_response) as mock_get, + ): + result = artifact_instance.get_download_url() + + assert result == expected_url + mock_get.assert_called_once() + call_kwargs = mock_get.call_args + assert call_kwargs.args[0] == "https://platform.aignostics.com/api/v1/runs/test-run-id/artifacts/artifact-123/file" + assert call_kwargs.kwargs["allow_redirects"] is False + assert call_kwargs.kwargs["headers"]["Authorization"] == "Bearer test-token" + + +@pytest.mark.unit +def test_artifact_get_download_url_retries_on_service_exception(artifact_instance, mock_api) -> None: + """Test that Artifact.get_download_url retries when a 5xx response is received. + + Args: + artifact_instance: Artifact instance with mock API. + mock_api: Mock ExternalsApi instance. + """ + mock_api.api_client = Mock() + mock_api.api_client.configuration.host = _PLATFORM_HOST + mock_api.api_client.configuration.token_provider = None + + error_response = Mock() + error_response.status_code = 503 + error_response.reason = "Service Unavailable" + + success_response = Mock() + success_response.status_code = 307 + success_response.headers = {"Location": _PRESIGNED_URL} + + mock_settings = Mock() + mock_settings.run_retry_attempts = 2 + mock_settings.run_retry_wait_min = 0.0 + mock_settings.run_retry_wait_max = 0.0 + + with ( + patch(_PATCH_GET_AUTH, return_value="test-token"), + patch(_PATCH_SETTINGS, return_value=mock_settings), + patch( + _PATCH_REQUESTS_GET, + side_effect=[error_response, success_response], + ) as mock_get, + ): + result = artifact_instance.get_download_url() + + assert result == _PRESIGNED_URL + assert mock_get.call_count == 2 + + +@pytest.mark.unit +def test_artifact_get_download_url_raises_not_found(artifact_instance, mock_api) -> None: + """Test that Artifact.get_download_url raises NotFoundException for a 404 response. + + Args: + artifact_instance: Artifact instance with mock API. + mock_api: Mock ExternalsApi instance. + """ + from aignx.codegen.exceptions import NotFoundException + + mock_api.api_client = Mock() + mock_api.api_client.configuration.host = _PLATFORM_HOST + mock_api.api_client.configuration.token_provider = None + + not_found_response = Mock() + not_found_response.status_code = 404 + not_found_response.reason = "Not Found" + + with ( + patch(_PATCH_GET_AUTH, return_value="test-token"), + patch(_PATCH_REQUESTS_GET, return_value=not_found_response), + pytest.raises(NotFoundException), + ): + artifact_instance.get_download_url() + + +@pytest.mark.unit +def test_artifact_get_download_url_raises_service_exception_on_timeout(artifact_instance, mock_api) -> None: + """Test that Artifact.get_download_url wraps requests.Timeout in ServiceException. + + Args: + artifact_instance: Artifact instance with mock API. + mock_api: Mock ExternalsApi instance. + """ + import requests as requests_lib + from aignx.codegen.exceptions import ServiceException + + mock_api.api_client = Mock() + mock_api.api_client.configuration.host = _PLATFORM_HOST + mock_api.api_client.configuration.token_provider = None + + with ( + patch(_PATCH_GET_AUTH, return_value="test-token"), + patch( + _PATCH_REQUESTS_GET, + side_effect=requests_lib.Timeout("Connection timed out"), + ), + pytest.raises(ServiceException), + ): + artifact_instance.get_download_url() + + +@pytest.mark.unit +def test_artifact_get_download_url_raises_service_exception_on_connection_error(artifact_instance, mock_api) -> None: + """Test that Artifact.get_download_url wraps requests.ConnectionError in ServiceException. + + Args: + artifact_instance: Artifact instance with mock API. + mock_api: Mock ExternalsApi instance. + """ + import requests as requests_lib + from aignx.codegen.exceptions import ServiceException + + mock_api.api_client = Mock() + mock_api.api_client.configuration.host = _PLATFORM_HOST + mock_api.api_client.configuration.token_provider = None + + with ( + patch(_PATCH_GET_AUTH, return_value="test-token"), + patch( + _PATCH_REQUESTS_GET, + side_effect=requests_lib.ConnectionError("Connection refused"), + ), + pytest.raises(ServiceException), + ): + artifact_instance.get_download_url() + + +@pytest.mark.unit +def test_artifact_get_download_url_raises_runtime_error_on_missing_location_header(artifact_instance, mock_api) -> None: + """Test that Artifact.get_download_url raises RuntimeError when redirect has no Location header. + + Args: + artifact_instance: Artifact instance with mock API. + mock_api: Mock ExternalsApi instance. + """ + mock_api.api_client = Mock() + mock_api.api_client.configuration.host = _PLATFORM_HOST + mock_api.api_client.configuration.token_provider = None + + redirect_response = Mock() + redirect_response.status_code = 307 + redirect_response.headers = {} # No Location header + + with ( + patch(_PATCH_GET_AUTH, return_value="test-token"), + patch(_PATCH_REQUESTS_GET, return_value=redirect_response), + pytest.raises(RuntimeError, match="missing Location header"), + ): + artifact_instance.get_download_url() + + +@pytest.mark.unit +def test_artifact_get_download_url_raises_api_exception_on_4xx(artifact_instance, mock_api) -> None: + """Test that Artifact.get_download_url raises ApiException for 4xx client errors. + + Args: + artifact_instance: Artifact instance with mock API. + mock_api: Mock ExternalsApi instance. + """ + from aignx.codegen.exceptions import ApiException + + mock_api.api_client = Mock() + mock_api.api_client.configuration.host = _PLATFORM_HOST + mock_api.api_client.configuration.token_provider = None + + bad_request_response = Mock() + bad_request_response.status_code = 400 + bad_request_response.reason = "Bad Request" + + with ( + patch(_PATCH_GET_AUTH, return_value="test-token"), + patch(_PATCH_REQUESTS_GET, return_value=bad_request_response), + pytest.raises(ApiException), + ): + artifact_instance.get_download_url() + + +@pytest.mark.unit +def test_artifact_get_download_url_passes_proxy_and_ssl_config(artifact_instance, mock_api) -> None: + """Test that Artifact.get_download_url forwards proxy and SSL config from the API client. + + Args: + artifact_instance: Artifact instance with mock API. + mock_api: Mock ExternalsApi instance. + """ + mock_api.api_client = Mock() + mock_api.api_client.configuration.host = _PLATFORM_HOST + mock_api.api_client.configuration.token_provider = None + mock_api.api_client.configuration.proxy = _PROXY_URL + mock_api.api_client.configuration.ssl_ca_cert = "/path/to/ca-bundle.crt" + mock_api.api_client.configuration.verify_ssl = True + + mock_response = Mock() + mock_response.status_code = 307 + mock_response.headers = {"Location": _PRESIGNED_URL} + + with ( + patch(_PATCH_GET_AUTH, return_value="test-token"), + patch(_PATCH_REQUESTS_GET, return_value=mock_response) as mock_get, + ): + result = artifact_instance.get_download_url() + + assert result == _PRESIGNED_URL + call_kwargs = mock_get.call_args.kwargs + assert call_kwargs["proxies"] == { + "http": _PROXY_URL, + "https": _PROXY_URL, + } + assert call_kwargs["verify"] == "/path/to/ca-bundle.crt" + + +@pytest.mark.unit +def test_artifact_download_uses_codegen_client(artifact_instance, mock_api) -> None: + """Test that Artifact.download delegates to the codegen API client. + + Args: + artifact_instance: Artifact instance with mock API. + mock_api: Mock ExternalsApi instance. + """ + mock_api.get_artifact_url_v1_runs_run_id_artifacts_artifact_id_file_get.return_value = b"content" + + result = artifact_instance.download() + + assert result == b"content" + mock_api.get_artifact_url_v1_runs_run_id_artifacts_artifact_id_file_get.assert_called_once() + call_kwargs = mock_api.get_artifact_url_v1_runs_run_id_artifacts_artifact_id_file_get.call_args.kwargs + assert call_kwargs["run_id"] == "test-run-id" + assert call_kwargs["artifact_id"] == "artifact-123" + + +@pytest.mark.unit +def test_run_get_artifact_download_url_delegates_to_artifact(app_run, mock_api) -> None: + """Test that Run.get_artifact_download_url delegates to Artifact.get_download_url. + + Args: + app_run: Run instance with mock API. + mock_api: Mock ExternalsApi instance. + """ + mock_api.api_client = Mock() + mock_api.api_client.configuration.host = _PLATFORM_HOST + mock_api.api_client.configuration.token_provider = None + + mock_response = Mock() + mock_response.status_code = 307 + mock_response.headers = {"Location": _PRESIGNED_URL} + + with ( + patch(_PATCH_GET_AUTH, return_value="test-token"), + patch(_PATCH_REQUESTS_GET, return_value=mock_response), + ): + result = app_run.get_artifact_download_url("artifact-123") + + assert result == _PRESIGNED_URL + + +@pytest.mark.unit +def test_ensure_artifacts_downloaded_uses_output_artifact_id(app_run, mock_api, tmp_path) -> None: + """Test that ensure_artifacts_downloaded resolves URLs via get_artifact_download_url. + + Verifies that the method uses output_artifact_id to resolve presigned URLs + instead of using the deprecated download_url field. + + Args: + app_run: Run instance with mock API. + mock_api: Mock ExternalsApi instance. + tmp_path: Temporary directory for test files. + """ + from aignx.codegen.models import ArtifactOutput, ArtifactState, OutputArtifactResultReadResponse + + artifact = OutputArtifactResultReadResponse.model_construct( + output_artifact_id="artifact-abc", + name="cell_classification", + metadata={"checksum_base64_crc32c": "AAAA", "media_type": _MIME_TYPE_CSV}, + state=ArtifactState.TERMINATED, + output=ArtifactOutput.AVAILABLE, + ) + item = ItemResultReadResponse.model_construct( + external_id="slide-1", + output_artifacts=[artifact], + ) + + presigned_url = _PRESIGNED_URL + + with ( + patch.object(app_run, "get_artifact_download_url", return_value=presigned_url) as mock_get_url, + patch(_PATCH_DOWNLOAD_FILE) as mock_download_file, + patch("aignostics.platform.resources.runs.get_mime_type_for_artifact", return_value=_MIME_TYPE_CSV), + patch("aignostics.platform.resources.runs.mime_type_to_file_ending", return_value=".csv"), + ): + app_run.ensure_artifacts_downloaded(tmp_path, item) + + mock_get_url.assert_called_once_with("artifact-abc") + mock_download_file.assert_called_once_with( + presigned_url, + str(tmp_path / "slide-1" / "cell_classification.csv"), + "AAAA", + ) + + +@pytest.mark.unit +def test_ensure_artifacts_downloaded_skips_artifact_without_id(app_run, tmp_path) -> None: + """Test that ensure_artifacts_downloaded skips artifacts with no output_artifact_id. + + Args: + app_run: Run instance with mock API. + tmp_path: Temporary directory for test files. + """ + from aignx.codegen.models import ArtifactOutput, ArtifactState, OutputArtifactResultReadResponse + + artifact = OutputArtifactResultReadResponse.model_construct( + output_artifact_id=None, + name="cell_classification", + metadata={"checksum_base64_crc32c": "AAAA"}, + state=ArtifactState.TERMINATED, + output=ArtifactOutput.AVAILABLE, + ) + item = ItemResultReadResponse.model_construct( + external_id="slide-1", + output_artifacts=[artifact], + ) + + with ( + patch.object(app_run, "get_artifact_download_url") as mock_get_url, + patch(_PATCH_DOWNLOAD_FILE) as mock_download_file, + ): + app_run.ensure_artifacts_downloaded(tmp_path, item, print_status=False) + + mock_get_url.assert_not_called() + mock_download_file.assert_not_called() + + +@pytest.mark.unit +def test_ensure_artifacts_downloaded_skips_existing_file_with_matching_checksum(app_run, tmp_path) -> None: + """Test that ensure_artifacts_downloaded skips files that already exist with correct checksum. + + Args: + app_run: Run instance with mock API. + tmp_path: Temporary directory for test files. + """ + from aignx.codegen.models import ArtifactOutput, ArtifactState, OutputArtifactResultReadResponse + + artifact = OutputArtifactResultReadResponse.model_construct( + output_artifact_id="artifact-abc", + name="result", + metadata={"checksum_base64_crc32c": "test_checksum"}, + state=ArtifactState.TERMINATED, + output=ArtifactOutput.AVAILABLE, + ) + item = ItemResultReadResponse.model_construct( + external_id="slide-1", + output_artifacts=[artifact], + ) + + # Create the file so it "already exists" + item_dir = tmp_path / "slide-1" + item_dir.mkdir() + existing_file = item_dir / "result.csv" + existing_file.write_bytes(b"existing content") + + with ( + patch.object(app_run, "get_artifact_download_url") as mock_get_url, + patch(_PATCH_DOWNLOAD_FILE) as mock_download_file, + patch("aignostics.platform.resources.runs.get_mime_type_for_artifact", return_value=_MIME_TYPE_CSV), + patch("aignostics.platform.resources.runs.mime_type_to_file_ending", return_value=".csv"), + patch("aignostics.platform.resources.runs.calculate_file_crc32c", return_value="test_checksum"), + ): + app_run.ensure_artifacts_downloaded(tmp_path, item, print_status=False) + + mock_get_url.assert_not_called() + mock_download_file.assert_not_called() + + +@pytest.mark.unit +def test_artifact_get_download_url_raises_service_exception_on_request_exception(artifact_instance, mock_api) -> None: + """Test that Artifact.get_download_url wraps generic requests.RequestException in ServiceException. + + Covers the fallback except clause for non-Timeout/ConnectionError request failures, + e.g. TooManyRedirects or other requests.RequestException subclasses. + + Args: + artifact_instance: Artifact instance with mock API. + mock_api: Mock ExternalsApi instance. + """ + import requests as requests_lib + from aignx.codegen.exceptions import ServiceException + + mock_api.api_client = Mock() + mock_api.api_client.configuration.host = _PLATFORM_HOST + mock_api.api_client.configuration.token_provider = None + + mock_settings = Mock() + mock_settings.run_retry_attempts = 1 + mock_settings.run_retry_wait_min = 0.0 + mock_settings.run_retry_wait_max = 0.0 + + with ( + patch(_PATCH_GET_AUTH, return_value="test-token"), + patch(_PATCH_SETTINGS, return_value=mock_settings), + patch( + _PATCH_REQUESTS_GET, + side_effect=requests_lib.TooManyRedirects("Too many redirects"), + ), + pytest.raises(ServiceException), + ): + artifact_instance.get_download_url() + + +@pytest.mark.unit +def test_artifact_get_download_url_raises_runtime_error_on_unexpected_status(artifact_instance, mock_api) -> None: + """Test that Artifact.get_download_url raises RuntimeError for unexpected 2xx/3xx status codes. + + The endpoint should only return 307 redirects or error codes; any other status + (e.g. 200 OK) is unexpected and should raise a RuntimeError. + + Args: + artifact_instance: Artifact instance with mock API. + mock_api: Mock ExternalsApi instance. + """ + mock_api.api_client = Mock() + mock_api.api_client.configuration.host = _PLATFORM_HOST + mock_api.api_client.configuration.token_provider = None + + unexpected_response = Mock() + unexpected_response.status_code = 200 # Not a redirect, not an error + + mock_settings = Mock() + mock_settings.run_retry_attempts = 1 + mock_settings.run_retry_wait_min = 0.0 + mock_settings.run_retry_wait_max = 0.0 + + with ( + patch(_PATCH_GET_AUTH, return_value="test-token"), + patch(_PATCH_SETTINGS, return_value=mock_settings), + patch(_PATCH_REQUESTS_GET, return_value=unexpected_response), + pytest.raises(RuntimeError, match="Unexpected status 200"), + ): + artifact_instance.get_download_url()