diff --git a/azure-quantum/azure/quantum/job/base_job.py b/azure-quantum/azure/quantum/job/base_job.py index 96196453..9ffbd9ff 100644 --- a/azure-quantum/azure/quantum/job/base_job.py +++ b/azure-quantum/azure/quantum/job/base_job.py @@ -10,7 +10,7 @@ from datetime import datetime, timezone, timedelta from urllib.parse import urlparse, parse_qs from typing import Any, Dict, Optional, TYPE_CHECKING -from azure.storage.blob import BlobClient +from azure.storage.blob import BlobClient, BlobProperties from azure.quantum.storage import upload_blob, download_blob, download_blob_properties, ContainerClient from azure.quantum._client.models import JobDetails @@ -390,6 +390,25 @@ def download_attachment( return response + def list_attachments(self) -> list[BlobProperties]: + """ Lists the attachments in the job's linked storage container. Attachments are blobs of + data created as part of the Job's execution, or they can be uploaded directly from Python + using the upload_attachment method. + + :return: List of blobs in the job's linked storage container. + :rtype: list[~azure.storage.blob.BlobProperties] + """ + + # Use the job's linked storage container. + if self._details.container_uri is None: + container_uri = self.workspace.get_container_uri(job_id=self.id) + else: + container_uri = self._details.container_uri + + container_client = ContainerClient.from_container_url(container_uri) + return list(container_client.list_blobs()) + + def _get_blob_uri_with_sas_token(self, blob_uri: str) -> str: """Get Blob URI with SAS-token if one was not specified in blob_uri parameter :param blob_uri: Blob URI diff --git a/azure-quantum/tests/test_job_attachments.py b/azure-quantum/tests/test_job_attachments.py new file mode 100644 index 00000000..d537f61b --- /dev/null +++ b/azure-quantum/tests/test_job_attachments.py @@ -0,0 +1,54 @@ +## +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. +## + +from unittest.mock import Mock, patch +from azure.quantum import Job, JobDetails + + +CONTAINER_URI = "https://acct.blob.core.windows.net/job-id?sas" + + +def _job_with_container(container_uri=CONTAINER_URI, workspace=None) -> Job: + job_details = JobDetails( + id="job-id", + name="", + provider_id="", + target="", + container_uri=container_uri, + input_data_format="", + output_data_format="", + ) + return Job(workspace=workspace, job_details=job_details) + + +@patch("azure.quantum.job.base_job.ContainerClient") +def test_list_attachments_returns_container_blobs(mock_container_client): + job = _job_with_container() + + blob_a = Mock() + blob_b = Mock() + container = mock_container_client.from_container_url.return_value + container.list_blobs.return_value = [blob_a, blob_b] + + result = job.list_attachments() + + mock_container_client.from_container_url.assert_called_once_with(CONTAINER_URI) + assert result == [blob_a, blob_b] + + +@patch("azure.quantum.job.base_job.ContainerClient") +def test_list_attachments_uses_workspace_container_when_unset(mock_container_client): + workspace = Mock() + workspace.get_container_uri.return_value = CONTAINER_URI + job = _job_with_container(container_uri=None, workspace=workspace) + + container = mock_container_client.from_container_url.return_value + container.list_blobs.return_value = [] + + result = job.list_attachments() + + workspace.get_container_uri.assert_called_once_with(job_id="job-id") + mock_container_client.from_container_url.assert_called_once_with(CONTAINER_URI) + assert result == []