diff --git a/azure-quantum/azure/quantum/job/base_job.py b/azure-quantum/azure/quantum/job/base_job.py index 9ffbd9ff..41319784 100644 --- a/azure-quantum/azure/quantum/job/base_job.py +++ b/azure-quantum/azure/quantum/job/base_job.py @@ -42,6 +42,11 @@ class BaseJob(WorkspaceItem): :type details: ItemDetails """ + def __init__(self, workspace: "Workspace", details: JobDetails, **kwargs): + self._attachment_container_uri_cache: Optional[str] = None + self._attachment_container_uri_cache_container_name: Optional[str] = None + super().__init__(workspace=workspace, details=details, **kwargs) + @staticmethod def create_job_id() -> str: """Create a unique id for a new job.""" @@ -55,17 +60,20 @@ def details(self) -> JobDetails: @details.setter def details(self, value: JobDetails): self._details = value + self._attachment_container_uri_cache = None + self._attachment_container_uri_cache_container_name = None @property - def container_name(self): + def container_name(self) -> str: """Job input/output data container name""" if self._details.container_uri is None: return f"job-{self.id}" - else: - container_uri = self._details.container_uri - path = urlparse(container_uri).path - container_name = path.split("/")[1] + + path = urlparse(self._details.container_uri).path + container_name = path.lstrip("/").split("/", 1)[0] + if not container_name: + raise ValueError("Job container URI does not include a container name.") return container_name @classmethod @@ -344,12 +352,8 @@ def upload_attachment( :rtype: str """ - # Use Job's default container if not specified if container_uri is None: - 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_uri = self._get_attachment_container_uri() uploaded_blob_uri = self.upload_input_data( container_uri = container_uri, @@ -377,13 +381,9 @@ def download_attachment( :rtype: bytes """ - # Use Job's default container if not specified if container_uri is None: - 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_uri = self._get_attachment_container_uri() + container_client = ContainerClient.from_container_url(container_uri) blob_client = container_client.get_blob_client(name) response = blob_client.download_blob().readall() @@ -399,44 +399,94 @@ def list_attachments(self) -> list[BlobProperties]: :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_uri = self._get_attachment_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 - :type blob_uri: str - :return: Blob URI with SAS-token - :rtype: str + def _get_attachment_container_uri(self) -> str: + """Return a workspace-issued URI for the job's attachment container. + + Workspace-issued URIs are cached privately on the job and are never written back to job + details. A cached URI is reused while its container name is unchanged and its SAS expiry + remains outside the renewal window. """ - url = urlparse(blob_uri) - query_params = parse_qs(url.query) - token_expire_query_param = query_params.get("se") + container_uri = self._details.container_uri + container_name = self.container_name + cached_container_uri = self._attachment_container_uri_cache + if cached_container_uri: + if self._attachment_container_uri_cache_container_name != container_name: + self._attachment_container_uri_cache = None + self._attachment_container_uri_cache_container_name = None + elif self._is_attachment_container_uri_unexpired(cached_container_uri): + return cached_container_uri + + if container_uri is None: + refreshed_container_uri = self.workspace.get_container_uri(job_id=self.id) + else: + refreshed_container_uri = self.workspace.get_container_uri( + job_id=self.id, + container_name=container_name, + ) + + self._attachment_container_uri_cache = refreshed_container_uri + self._attachment_container_uri_cache_container_name = container_name + return refreshed_container_uri + - token_expire_time = None + # Buffer so a SAS token is never used a few seconds before its actual expiry. + _SAS_RENEWAL_BUFFER = timedelta(minutes=5) - if token_expire_query_param is not None: - token_expire_time_str = token_expire_query_param[0] + @staticmethod + def _get_sas_expiry_time(uri: str) -> Optional[datetime]: + """Parse the `se` (SAS expiry) query parameter from a URI. Returns None if missing or malformed.""" + + query_params = parse_qs(urlparse(uri).query) + token_expire_query_param = query_params.get("se") + if not token_expire_query_param: + return None - # Since python < 3.11 can not easily parse Z suffixed UTC timestamp and + try: + # Since python < 3.11 can not easily parse Z suffixed UTC timestamp and # assuming that the timestamp is always UTC, we replace that suffix with UTC offset. token_expire_time = datetime.fromisoformat( - token_expire_time_str.replace('Z', '+00:00') + token_expire_query_param[0].replace("Z", "+00:00") ) - - # Make an expiration time a little earlier, so there's no case where token is - # used a second or so before of its expiration. - token_expire_time = token_expire_time - timedelta(minutes=5) + except ValueError: + logger.debug("Unable to parse SAS expiry time.") + return None + + if token_expire_time.tzinfo is None: + token_expire_time = token_expire_time.replace(tzinfo=timezone.utc) + return token_expire_time + + + def _is_attachment_container_uri_unexpired( + self, + container_uri: str, + ) -> bool: + """Check whether a cached container SAS remains outside the renewal window.""" + + token_expire_time = self._get_sas_expiry_time(container_uri) + if token_expire_time is None: + return False + + current_utc_time = datetime.now(tz=timezone.utc) + return current_utc_time + self._SAS_RENEWAL_BUFFER < token_expire_time + + + 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 + :type blob_uri: str + :return: Blob URI with SAS-token + :rtype: str + """ + token_expire_time = self._get_sas_expiry_time(blob_uri) current_utc_time = datetime.now(tz=timezone.utc) - if token_expire_time is None or current_utc_time >= token_expire_time: + if token_expire_time is None or current_utc_time + self._SAS_RENEWAL_BUFFER >= token_expire_time: # blob_uri does not contains SAS token or it is expired, # get sas url from service blob_client = BlobClient.from_blob_url( diff --git a/azure-quantum/tests/test_job_attachments.py b/azure-quantum/tests/test_job_attachments.py index d537f61b..3fa6a034 100644 --- a/azure-quantum/tests/test_job_attachments.py +++ b/azure-quantum/tests/test_job_attachments.py @@ -3,16 +3,51 @@ # Licensed under the MIT License. ## -from unittest.mock import Mock, patch +from datetime import datetime, timedelta, timezone +from unittest.mock import Mock, call, patch + +import pytest + from azure.quantum import Job, JobDetails -CONTAINER_URI = "https://acct.blob.core.windows.net/job-id?sas" +JOB_ID = "job-id" +DEFAULT_CONTAINER_NAME = f"job-{JOB_ID}" +UNSIGNED_CONTAINER_URI = f"https://acct.blob.core.windows.net/{DEFAULT_CONTAINER_NAME}" +SIGNED_CONTAINER_URI = ( + f"{UNSIGNED_CONTAINER_URI}?sp=racwdl&se=2099-01-01T00%3A00%3A00Z&sig=signature" +) +EXPIRED_CONTAINER_URI = ( + f"{UNSIGNED_CONTAINER_URI}?sp=racwdl&se=2000-01-01T00%3A00%3A00Z&sig=signature" +) +FUTURE_START_CONTAINER_URI = ( + f"{UNSIGNED_CONTAINER_URI}?sp=racwdl&st=2098-01-01T00%3A00%3A00Z" + "&se=2099-01-01T00%3A00%3A00Z&sig=signature" +) +MALFORMED_EXPIRY_CONTAINER_URI = ( + f"{UNSIGNED_CONTAINER_URI}?sp=racwdl&se=not-a-date&sig=signature" +) +MINIMUM_EXPIRY_CONTAINER_URI = ( + f"{UNSIGNED_CONTAINER_URI}?sp=racwdl&se=0001-01-01T00%3A00%3A00Z&sig=signature" +) +NO_PERMISSIONS_CONTAINER_URI = ( + f"{UNSIGNED_CONTAINER_URI}?se=2099-01-01T00%3A00%3A00Z&sig=signature" +) +NO_SIGNATURE_CONTAINER_URI = ( + f"{UNSIGNED_CONTAINER_URI}?se=2099-01-01T00%3A00%3A00Z" +) +WRITE_ONLY_CONTAINER_URI = ( + f"{UNSIGNED_CONTAINER_URI}?sp=w&se=2099-01-01T00%3A00%3A00Z&sig=signature" +) +READ_LIST_CONTAINER_URI = ( + f"{UNSIGNED_CONTAINER_URI}?sp=rl&se=2099-01-01T00%3A00%3A00Z&sig=signature" +) +HTTP_SIGNED_CONTAINER_URI = SIGNED_CONTAINER_URI.replace("https://", "http://") -def _job_with_container(container_uri=CONTAINER_URI, workspace=None) -> Job: +def _job_with_container(container_uri=UNSIGNED_CONTAINER_URI, workspace=None) -> Job: job_details = JobDetails( - id="job-id", + id=JOB_ID, name="", provider_id="", target="", @@ -25,7 +60,9 @@ def _job_with_container(container_uri=CONTAINER_URI, workspace=None) -> Job: @patch("azure.quantum.job.base_job.ContainerClient") def test_list_attachments_returns_container_blobs(mock_container_client): - job = _job_with_container() + workspace = Mock() + workspace.get_container_uri.return_value = SIGNED_CONTAINER_URI + job = _job_with_container(workspace=workspace) blob_a = Mock() blob_b = Mock() @@ -34,14 +71,18 @@ def test_list_attachments_returns_container_blobs(mock_container_client): result = job.list_attachments() - mock_container_client.from_container_url.assert_called_once_with(CONTAINER_URI) + workspace.get_container_uri.assert_called_once_with( + job_id=JOB_ID, + container_name=DEFAULT_CONTAINER_NAME, + ) + mock_container_client.from_container_url.assert_called_once_with(SIGNED_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 + workspace.get_container_uri.return_value = SIGNED_CONTAINER_URI job = _job_with_container(container_uri=None, workspace=workspace) container = mock_container_client.from_container_url.return_value @@ -49,6 +90,435 @@ def test_list_attachments_uses_workspace_container_when_unset(mock_container_cli 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) + workspace.get_container_uri.assert_called_once_with(job_id=JOB_ID) + mock_container_client.from_container_url.assert_called_once_with(SIGNED_CONTAINER_URI) assert result == [] + + +def test_upload_attachment_uses_fresh_workspace_container_uri(): + workspace = Mock() + workspace.get_container_uri.return_value = SIGNED_CONTAINER_URI + job = _job_with_container(workspace=workspace) + job.upload_input_data = Mock(return_value="uploaded-uri") + + result = job.upload_attachment("attachment", b"data") + + workspace.get_container_uri.assert_called_once_with( + job_id=JOB_ID, + container_name=DEFAULT_CONTAINER_NAME, + ) + job.upload_input_data.assert_called_once_with( + container_uri=SIGNED_CONTAINER_URI, + blob_name="attachment", + input_data=b"data", + ) + assert result == "uploaded-uri" + + +@patch("azure.quantum.job.base_job.ContainerClient") +def test_download_attachment_uses_fresh_workspace_container_uri(mock_container_client): + workspace = Mock() + workspace.get_container_uri.return_value = SIGNED_CONTAINER_URI + job = _job_with_container(workspace=workspace) + blob_client = mock_container_client.from_container_url.return_value.get_blob_client.return_value + blob_client.download_blob.return_value.readall.return_value = b"data" + + result = job.download_attachment("attachment") + + workspace.get_container_uri.assert_called_once_with( + job_id=JOB_ID, + container_name=DEFAULT_CONTAINER_NAME, + ) + mock_container_client.from_container_url.assert_called_once_with(SIGNED_CONTAINER_URI) + assert result == b"data" + + +@patch("azure.quantum.job.base_job.ContainerClient") +def test_attachment_methods_honor_explicit_container_uri(mock_container_client): + workspace = Mock() + job = _job_with_container(workspace=workspace) + job.upload_input_data = Mock(return_value="uploaded-uri") + blob_client = mock_container_client.from_container_url.return_value.get_blob_client.return_value + blob_client.download_blob.return_value.readall.return_value = b"data" + explicit_uri = "https://custom.blob.core.windows.net/container?sas" + + job.upload_attachment("upload", b"data", container_uri=explicit_uri) + job.download_attachment("download", container_uri=explicit_uri) + + workspace.get_container_uri.assert_not_called() + job.upload_input_data.assert_called_once_with( + container_uri=explicit_uri, + blob_name="upload", + input_data=b"data", + ) + mock_container_client.from_container_url.assert_called_once_with(explicit_uri) + + +@patch("azure.quantum.job.base_job.ContainerClient") +def test_attachment_methods_refresh_signed_job_uri_once(mock_container_client): + workspace = Mock() + workspace.get_container_uri.return_value = SIGNED_CONTAINER_URI + job = _job_with_container(container_uri=SIGNED_CONTAINER_URI, workspace=workspace) + job.upload_input_data = Mock(return_value="uploaded-uri") + container_client = mock_container_client.from_container_url.return_value + container_client.list_blobs.return_value = [] + container_client.get_blob_client.return_value.download_blob.return_value.readall.return_value = b"data" + + job.upload_attachment("upload", b"data") + attachments = job.list_attachments() + downloaded = job.download_attachment("download") + + workspace.get_container_uri.assert_called_once_with( + job_id=JOB_ID, + container_name=DEFAULT_CONTAINER_NAME, + ) + job.upload_input_data.assert_called_once_with( + container_uri=SIGNED_CONTAINER_URI, + blob_name="upload", + input_data=b"data", + ) + assert mock_container_client.from_container_url.call_args_list == [ + call(SIGNED_CONTAINER_URI), + call(SIGNED_CONTAINER_URI), + ] + assert attachments == [] + assert downloaded == b"data" + + +def test_upload_attachment_refreshes_expired_job_uri(): + workspace = Mock() + workspace.get_container_uri.return_value = SIGNED_CONTAINER_URI + job = _job_with_container(container_uri=EXPIRED_CONTAINER_URI, workspace=workspace) + job.upload_input_data = Mock(return_value="uploaded-uri") + + job.upload_attachment("attachment", b"data") + + workspace.get_container_uri.assert_called_once_with( + job_id=JOB_ID, + container_name=DEFAULT_CONTAINER_NAME, + ) + job.upload_input_data.assert_called_once_with( + container_uri=SIGNED_CONTAINER_URI, + blob_name="attachment", + input_data=b"data", + ) + + +def test_cached_uri_is_reused_regardless_of_sas_start_time(): + workspace = Mock() + job = _job_with_container(workspace=workspace) + job._attachment_container_uri_cache = FUTURE_START_CONTAINER_URI + job._attachment_container_uri_cache_container_name = DEFAULT_CONTAINER_NAME + job.upload_input_data = Mock(return_value="uploaded-uri") + + job.upload_attachment("attachment", b"data") + + workspace.get_container_uri.assert_not_called() + job.upload_input_data.assert_called_once_with( + container_uri=FUTURE_START_CONTAINER_URI, + blob_name="attachment", + input_data=b"data", + ) + + +def test_upload_attachment_logs_and_refreshes_malformed_sas_expiry(caplog): + workspace = Mock() + workspace.get_container_uri.return_value = SIGNED_CONTAINER_URI + job = _job_with_container(workspace=workspace) + job._attachment_container_uri_cache = MALFORMED_EXPIRY_CONTAINER_URI + job._attachment_container_uri_cache_container_name = DEFAULT_CONTAINER_NAME + job.upload_input_data = Mock(return_value="uploaded-uri") + + with caplog.at_level("DEBUG", logger="azure.quantum.job.base_job"): + job.upload_attachment("attachment", b"data") + + assert "Unable to parse SAS expiry time" in caplog.text + workspace.get_container_uri.assert_called_once_with( + job_id=JOB_ID, + container_name=DEFAULT_CONTAINER_NAME, + ) + + +def test_cached_uri_inside_expiry_buffer_is_refreshed(): + near_expiry_container_uri = ( + f"{UNSIGNED_CONTAINER_URI}?se=" + f"{(datetime.now(timezone.utc) + timedelta(minutes=4)).isoformat().replace('+00:00', 'Z')}" + ) + workspace = Mock() + workspace.get_container_uri.return_value = SIGNED_CONTAINER_URI + job = _job_with_container(workspace=workspace) + job._attachment_container_uri_cache = near_expiry_container_uri + job._attachment_container_uri_cache_container_name = DEFAULT_CONTAINER_NAME + job.upload_input_data = Mock(return_value="uploaded-uri") + + job.upload_attachment("attachment", b"data") + + workspace.get_container_uri.assert_called_once_with( + job_id=JOB_ID, + container_name=DEFAULT_CONTAINER_NAME, + ) + job.upload_input_data.assert_called_once_with( + container_uri=SIGNED_CONTAINER_URI, + blob_name="attachment", + input_data=b"data", + ) + + +@pytest.mark.parametrize( + "cached_container_uri", + [UNSIGNED_CONTAINER_URI, EXPIRED_CONTAINER_URI, MINIMUM_EXPIRY_CONTAINER_URI], +) +def test_cached_uri_without_reusable_expiry_is_refreshed(cached_container_uri): + workspace = Mock() + workspace.get_container_uri.return_value = SIGNED_CONTAINER_URI + job = _job_with_container(workspace=workspace) + job._attachment_container_uri_cache = cached_container_uri + job._attachment_container_uri_cache_container_name = DEFAULT_CONTAINER_NAME + job.upload_input_data = Mock(return_value="uploaded-uri") + + job.upload_attachment("attachment", b"data") + + workspace.get_container_uri.assert_called_once_with( + job_id=JOB_ID, + container_name=DEFAULT_CONTAINER_NAME, + ) + job.upload_input_data.assert_called_once_with( + container_uri=SIGNED_CONTAINER_URI, + blob_name="attachment", + input_data=b"data", + ) + + +@pytest.mark.parametrize( + "cached_container_uri", + [ + FUTURE_START_CONTAINER_URI, + HTTP_SIGNED_CONTAINER_URI, + NO_PERMISSIONS_CONTAINER_URI, + NO_SIGNATURE_CONTAINER_URI, + WRITE_ONLY_CONTAINER_URI, + READ_LIST_CONTAINER_URI, + ], +) +def test_unexpired_cached_uri_is_reused_without_prevalidation(cached_container_uri): + workspace = Mock() + job = _job_with_container(workspace=workspace) + job._attachment_container_uri_cache = cached_container_uri + job._attachment_container_uri_cache_container_name = DEFAULT_CONTAINER_NAME + job.upload_input_data = Mock(return_value="uploaded-uri") + + job.upload_attachment("attachment", b"data") + + workspace.get_container_uri.assert_not_called() + job.upload_input_data.assert_called_once_with( + container_uri=cached_container_uri, + blob_name="attachment", + input_data=b"data", + ) + + +@pytest.mark.parametrize( + "refreshed_container_uri", + [ + UNSIGNED_CONTAINER_URI, + HTTP_SIGNED_CONTAINER_URI, + NO_PERMISSIONS_CONTAINER_URI, + WRITE_ONLY_CONTAINER_URI, + READ_LIST_CONTAINER_URI, + f"https://other-acct.blob.core.windows.net/{DEFAULT_CONTAINER_NAME}?sas", + ], +) +def test_workspace_issued_uri_is_used_without_prevalidation(refreshed_container_uri): + workspace = Mock() + workspace.get_container_uri.return_value = refreshed_container_uri + job = _job_with_container(workspace=workspace) + original_container_uri = job.details.container_uri + job.upload_input_data = Mock(return_value="uploaded-uri") + + job.upload_attachment("attachment", b"data") + + workspace.get_container_uri.assert_called_once_with( + job_id=JOB_ID, + container_name=DEFAULT_CONTAINER_NAME, + ) + job.upload_input_data.assert_called_once_with( + container_uri=refreshed_container_uri, + blob_name="attachment", + input_data=b"data", + ) + assert job._attachment_container_uri_cache == refreshed_container_uri + assert job._attachment_container_uri_cache_container_name == DEFAULT_CONTAINER_NAME + assert job.details.container_uri == original_container_uri + + +def test_upload_attachment_rejects_job_uri_without_container_name(): + pathless_uri = ( + "https://acct.blob.core.windows.net" + "?sp=racwdl&se=2099-01-01T00%3A00%3A00Z&sig=signature" + ) + workspace = Mock() + job = _job_with_container(container_uri=pathless_uri, workspace=workspace) + + with pytest.raises(ValueError, match="does not include a container name"): + job.upload_attachment("attachment", b"data") + + workspace.get_container_uri.assert_not_called() + + +@patch("azure.quantum.job.base_job.ContainerClient") +def test_attachment_methods_preserve_custom_container_name(mock_container_client): + custom_container_name = "custom-container" + custom_unsigned_uri = f"https://acct.blob.core.windows.net/{custom_container_name}" + custom_signed_uri = ( + f"{custom_unsigned_uri}?sp=racwdl&se=2099-01-01T00%3A00%3A00Z&sig=signature" + ) + workspace = Mock() + workspace.get_container_uri.return_value = custom_signed_uri + job = _job_with_container(container_uri=custom_unsigned_uri, workspace=workspace) + job.upload_input_data = Mock(return_value="uploaded-uri") + container_client = mock_container_client.from_container_url.return_value + container_client.list_blobs.return_value = [] + container_client.get_blob_client.return_value.download_blob.return_value.readall.return_value = b"data" + + job.upload_attachment("upload", b"data") + attachments = job.list_attachments() + downloaded = job.download_attachment("download") + + workspace.get_container_uri.assert_called_once_with( + job_id=JOB_ID, + container_name=custom_container_name, + ) + job.upload_input_data.assert_called_once_with( + container_uri=custom_signed_uri, + blob_name="upload", + input_data=b"data", + ) + assert mock_container_client.from_container_url.call_args_list == [ + call(custom_signed_uri), + call(custom_signed_uri), + ] + assert attachments == [] + assert downloaded == b"data" + + +@patch("azure.quantum.job.base_job.ContainerClient") +def test_replacing_job_details_invalidates_cached_container_uri(mock_container_client): + workspace = Mock() + workspace.get_container_uri.side_effect = [SIGNED_CONTAINER_URI, SIGNED_CONTAINER_URI] + job = _job_with_container(workspace=workspace) + mock_container_client.from_container_url.return_value.list_blobs.return_value = [] + + job.list_attachments() + job.details = JobDetails( + id=JOB_ID, + name="", + provider_id="", + target="", + container_uri=UNSIGNED_CONTAINER_URI, + input_data_format="", + output_data_format="", + ) + job.list_attachments() + + assert workspace.get_container_uri.call_count == 2 + + +@patch("azure.quantum.job.base_job.ContainerClient") +def test_cache_reuse_after_hostname_mutation_never_queries_by_hostname(mock_container_client): + other_account_same_name_uri = f"https://other-acct.blob.core.windows.net/{DEFAULT_CONTAINER_NAME}" + workspace = Mock() + workspace.get_container_uri.return_value = SIGNED_CONTAINER_URI + job = _job_with_container(workspace=workspace) + mock_container_client.from_container_url.return_value.list_blobs.return_value = [] + + job.list_attachments() + job.details.container_uri = other_account_same_name_uri + job.list_attachments() + + # container_name-only cache key reuses the cache despite the hostname mutation. + workspace.get_container_uri.assert_called_once_with( + job_id=JOB_ID, + container_name=DEFAULT_CONTAINER_NAME, + ) + # get_container_uri never received a hostname, so both calls used the same real account SAS. + assert mock_container_client.from_container_url.call_args_list == [ + call(SIGNED_CONTAINER_URI), + call(SIGNED_CONTAINER_URI), + ] + + +@patch("azure.quantum.job.base_job.ContainerClient") +def test_mutating_job_container_uri_invalidates_cached_container_uri(mock_container_client): + other_container_uri = "https://acct.blob.core.windows.net/other-container" + other_signed_uri = ( + f"{other_container_uri}?sp=racwdl&se=2099-01-01T00%3A00%3A00Z&sig=signature" + ) + workspace = Mock() + workspace.get_container_uri.side_effect = [SIGNED_CONTAINER_URI, other_signed_uri] + job = _job_with_container(workspace=workspace) + mock_container_client.from_container_url.return_value.list_blobs.return_value = [] + + job.list_attachments() + job.details.container_uri = other_container_uri + job.list_attachments() + + assert workspace.get_container_uri.call_args_list == [ + call(job_id=JOB_ID, container_name=DEFAULT_CONTAINER_NAME), + call(job_id=JOB_ID, container_name="other-container"), + ] + assert mock_container_client.from_container_url.call_args_list == [ + call(SIGNED_CONTAINER_URI), + call(other_signed_uri), + ] + + +BLOB_NAME = "outputData" +UNSIGNED_BLOB_URI = f"https://acct.blob.core.windows.net/{DEFAULT_CONTAINER_NAME}/{BLOB_NAME}" +SIGNED_BLOB_URI = ( + f"{UNSIGNED_BLOB_URI}?sp=r&se=2099-01-01T00%3A00%3A00Z&sig=signature" +) +EXPIRED_BLOB_URI = ( + f"{UNSIGNED_BLOB_URI}?sp=r&se=2000-01-01T00%3A00%3A00Z&sig=signature" +) +MALFORMED_EXPIRY_BLOB_URI = ( + f"{UNSIGNED_BLOB_URI}?sp=r&se=not-a-date&sig=signature" +) + + +def test_get_blob_uri_with_sas_token_reuses_unexpired_uri(): + workspace = Mock() + job = _job_with_container(workspace=workspace) + + result = job._get_blob_uri_with_sas_token(SIGNED_BLOB_URI) + + assert result == SIGNED_BLOB_URI + workspace._get_linked_storage_sas_uri.assert_not_called() + + +@pytest.mark.parametrize("blob_uri", [UNSIGNED_BLOB_URI, EXPIRED_BLOB_URI]) +def test_get_blob_uri_with_sas_token_refreshes_expired_or_unsigned_uri(blob_uri): + workspace = Mock() + workspace._get_linked_storage_sas_uri.return_value = SIGNED_BLOB_URI + job = _job_with_container(workspace=workspace) + + result = job._get_blob_uri_with_sas_token(blob_uri) + + assert result == SIGNED_BLOB_URI + workspace._get_linked_storage_sas_uri.assert_called_once_with( + DEFAULT_CONTAINER_NAME, BLOB_NAME + ) + + +def test_get_blob_uri_with_sas_token_logs_and_refreshes_malformed_expiry(caplog): + workspace = Mock() + workspace._get_linked_storage_sas_uri.return_value = SIGNED_BLOB_URI + job = _job_with_container(workspace=workspace) + + with caplog.at_level("DEBUG", logger="azure.quantum.job.base_job"): + result = job._get_blob_uri_with_sas_token(MALFORMED_EXPIRY_BLOB_URI) + + assert result == SIGNED_BLOB_URI + assert "Unable to parse SAS expiry time" in caplog.text + workspace._get_linked_storage_sas_uri.assert_called_once_with( + DEFAULT_CONTAINER_NAME, BLOB_NAME + )