From 4292df06465e2bd12416835903656b56947080ec Mon Sep 17 00:00:00 2001
From: vincenttran-msft <vincenttran@microsoft.com>
Date: Fri, 14 Feb 2025 16:21:50 -0800
Subject: [PATCH] Iterate in _download_async, add live test for compressed data
 and tuning decompress knob

---
 .../azure/storage/blob/aio/_download_async.py |  7 ++++--
 .../tests/test_common_blob.py                 | 23 +++++++++++++++++
 .../tests/test_common_blob_async.py           | 25 +++++++++++++++++++
 3 files changed, 53 insertions(+), 2 deletions(-)

diff --git a/sdk/storage/azure-storage-blob/azure/storage/blob/aio/_download_async.py b/sdk/storage/azure-storage-blob/azure/storage/blob/aio/_download_async.py
index 4676c2e66754..ac866b4bae3d 100644
--- a/sdk/storage/azure-storage-blob/azure/storage/blob/aio/_download_async.py
+++ b/sdk/storage/azure-storage-blob/azure/storage/blob/aio/_download_async.py
@@ -46,8 +46,11 @@
 async def process_content(data: Any, start_offset: int, end_offset: int, encryption: Dict[str, Any]) -> bytes:
     if data is None:
         raise ValueError("Response cannot be None.")
-    await data.response.load_body()
-    content = cast(bytes, data.response.body())
+    if not data.response.is_stream_consumed:
+        content = b"".join([d async for d in data])
+    else:
+        await data.response.read()
+        content = cast(bytes, data.response.content)
     if encryption.get('key') is not None or encryption.get('resolver') is not None:
         try:
             return decrypt_blob(
diff --git a/sdk/storage/azure-storage-blob/tests/test_common_blob.py b/sdk/storage/azure-storage-blob/tests/test_common_blob.py
index b7e8936b3656..a82adf43f290 100644
--- a/sdk/storage/azure-storage-blob/tests/test_common_blob.py
+++ b/sdk/storage/azure-storage-blob/tests/test_common_blob.py
@@ -3531,4 +3531,27 @@ def test_upload_blob_partial_stream_chunked(self, **kwargs):
         result = blob.download_blob().readall()
         assert result == data[:length]
 
+    @pytest.mark.live_test_only
+    @BlobPreparer()
+    def test_download_blob_compressed_data_with_decompress(self, **kwargs):
+        storage_account_name = kwargs.pop("storage_account_name")
+        storage_account_key = kwargs.pop("storage_account_key")
+
+        # Arrange
+        self._setup(storage_account_name, storage_account_key)
+        blob_name = self._get_blob_reference()
+        blob = self.bsc.get_blob_client(self.container_name, blob_name)
+        compressed_data = b'\x1f\x8b\x08\x00\x00\x00\x00\x00\x00\xff\xcaH\xcd\xc9\xc9WH+\xca\xcfUH\xaf\xca,\x00\x00\x00\x00\xff\xff\x03\x00d\xaa\x8e\xb5\x0f\x00\x00\x00'
+        decompressed_data = b"hello from gzip"
+        content_settings = ContentSettings(content_encoding='gzip')
+
+        # Act / Assert
+        blob.upload_blob(data=compressed_data, content_settings=content_settings, overwrite=True)
+
+        result = blob.download_blob(decompress=True).readall()
+        assert result == decompressed_data
+
+        result = blob.download_blob(decompress=False).readall()
+        assert result == compressed_data
+
     # ------------------------------------------------------------------------------
\ No newline at end of file
diff --git a/sdk/storage/azure-storage-blob/tests/test_common_blob_async.py b/sdk/storage/azure-storage-blob/tests/test_common_blob_async.py
index aef6d9680603..e75a76f4a9ed 100644
--- a/sdk/storage/azure-storage-blob/tests/test_common_blob_async.py
+++ b/sdk/storage/azure-storage-blob/tests/test_common_blob_async.py
@@ -3510,4 +3510,29 @@ async def test_mock_transport_with_content_validation(self, **kwargs):
 
         blob_data = await (await blob_client.download_blob(validate_content=True)).read()
         assert blob_data == b"Hello Async World!"  # data is fixed by mock transport
+
+    @pytest.mark.live_test_only
+    @BlobPreparer()
+    async def test_download_blob_compressed_data_with_decompress(self, **kwargs):
+        storage_account_name = kwargs.pop("storage_account_name")
+        storage_account_key = kwargs.pop("storage_account_key")
+
+        # Arrange
+        await self._setup(storage_account_name, storage_account_key)
+        blob_name = self._get_blob_reference()
+        blob = self.bsc.get_blob_client(self.container_name, blob_name)
+        compressed_data = b'\x1f\x8b\x08\x00\x00\x00\x00\x00\x00\xff\xcaH\xcd\xc9\xc9WH+\xca\xcfUH\xaf\xca,\x00\x00\x00\x00\xff\xff\x03\x00d\xaa\x8e\xb5\x0f\x00\x00\x00'
+        decompressed_data = b"hello from gzip"
+        content_settings = ContentSettings(content_encoding='gzip')
+
+        # Act / Assert
+        await blob.upload_blob(data=compressed_data, content_settings=content_settings, overwrite=True)
+
+        downloaded = await blob.download_blob(decompress=True)
+        result = await downloaded.readall()
+        assert result == decompressed_data
+
+        downloaded = await blob.download_blob(decompress=False)
+        result = await downloaded.readall()
+        assert result == compressed_data
 # ------------------------------------------------------------------------------