diff --git a/pyrit/prompt_target/azure_blob_storage_target.py b/pyrit/prompt_target/azure_blob_storage_target.py index d94460246f..698199b6e4 100644 --- a/pyrit/prompt_target/azure_blob_storage_target.py +++ b/pyrit/prompt_target/azure_blob_storage_target.py @@ -276,7 +276,8 @@ def _parse_url(self) -> tuple[str, str]: "path, e.g. https://.blob.core.windows.net/." ) container_name = path_parts[1] - blob_prefix = "/".join(path_parts[2:]) + # Strip slashes so a trailing "/" on the URL does not add an empty segment to the blob path. + blob_prefix = "/".join(path_parts[2:]).strip("/") container_url = f"https://{parsed_url.netloc}/{container_name}" return container_url, blob_prefix @@ -303,7 +304,8 @@ async def _send_prompt_to_target_async(self, *, normalized_conversation: list[Me file_name = self._sanitize_file_name(str(request.prompt_metadata["file_name"])) data = str.encode(request.converted_value) - blob_url = self._container_url + "/" + file_name + # Drop a trailing "/" on the configured URL so the returned URL has no "//" before the file name. + blob_url = self._container_url.rstrip("/") + "/" + file_name await self._upload_blob_async(file_name=file_name, data=data, content_type=self._blob_content_type) diff --git a/tests/unit/prompt_target/target/test_prompt_target_azure_blob_storage.py b/tests/unit/prompt_target/target/test_prompt_target_azure_blob_storage.py index 9265225cba..56dd63c9b2 100644 --- a/tests/unit/prompt_target/target/test_prompt_target_azure_blob_storage.py +++ b/tests/unit/prompt_target/target/test_prompt_target_azure_blob_storage.py @@ -341,3 +341,40 @@ async def test_send_prompt_async_rejects_traversal_file_name( await azure_blob_storage_target.send_prompt_async(message=request) mock_upload.assert_not_called() + + +@pytest.mark.parametrize( + "container_url, expected_blob_path, expected_blob_url", + [ + ("https://test.blob.core.windows.net/xpia", "doc.txt", "https://test.blob.core.windows.net/xpia/doc.txt"), + ("https://test.blob.core.windows.net/xpia/", "doc.txt", "https://test.blob.core.windows.net/xpia/doc.txt"), + ( + "https://test.blob.core.windows.net/xpia/folder", + "folder/doc.txt", + "https://test.blob.core.windows.net/xpia/folder/doc.txt", + ), + ( + "https://test.blob.core.windows.net/xpia/folder/", + "folder/doc.txt", + "https://test.blob.core.windows.net/xpia/folder/doc.txt", + ), + ], +) +async def test_send_prompt_async_returns_url_of_uploaded_blob( + patch_central_database, container_url: str, expected_blob_path: str, expected_blob_url: str +): + target = AzureBlobStorageTarget(container_url=container_url, sas_token="valid_sas_token") + mock_blob_client = AsyncMock() + mock_blob_client.exists.return_value = False + mock_client = AsyncMock() + mock_client.get_blob_client = MagicMock(return_value=mock_blob_client) + message_piece = MessagePiece(role="user", original_value="payload", prompt_metadata={"file_name": "doc.txt"}) + + with patch.object(AsyncContainerClient, "from_container_url", return_value=mock_client) as mock_from_container_url: + response = await target.send_prompt_async(message=Message(message_pieces=[message_piece])) + + mock_from_container_url.assert_called_once_with( + container_url="https://test.blob.core.windows.net/xpia", credential="valid_sas_token" + ) + mock_client.get_blob_client.assert_called_once_with(blob=expected_blob_path) + assert response[0].get_value() == expected_blob_url