[Storage] Fixed Download/Upload APIs to be compatible with all transports + responses (#40615)

Author: weirongw23-msft

sdk/storage/azure-storage-blob/assets.json

@@ -2,5 +2,5 @@
   "AssetsRepo": "Azure/azure-sdk-assets",
   "AssetsRepoPrefixPath": "python",
   "TagPrefix": "python/storage/azure-storage-blob",
-  "Tag": "python/storage/azure-storage-blob_e1b13301f6"
+  "Tag": "python/storage/azure-storage-blob_2bfcc41daa"
 }

sdk/storage/azure-storage-blob/azure/storage/blob/_shared/policies_async.py

@@ -39,10 +39,11 @@ async def retry_hook(settings, **kwargs):
 async def is_checksum_retry(response):
     # retry if invalid content md5
     if response.context.get("validate_content", False) and response.http_response.headers.get("content-md5"):
-        try:
-            await response.http_response.load_body()  # Load the body in memory and close the socket
-        except (StreamClosedError, StreamConsumedError):
-            pass
+        if hasattr(response.http_response, "load_body"):
+            try:
+                await response.http_response.load_body()  # Load the body in memory and close the socket
+            except (StreamClosedError, StreamConsumedError):
+                pass
         computed_md5 = response.http_request.headers.get("content-md5", None) or encode_base64(
             StorageContentValidation.get_content_md5(response.http_response.body())
         )

sdk/storage/azure-storage-blob/azure/storage/blob/aio/_download_async.py

@@ -46,8 +46,10 @@
 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 hasattr(data.response, "is_stream_consumed") and data.response.is_stream_consumed:
+        content = data.response.content
+    else:
+        content = b"".join([d async for d in data])
     if encryption.get('key') is not None or encryption.get('resolver') is not None:
         try:
             return decrypt_blob(
@@ -57,12 +59,14 @@ async def process_content(data: Any, start_offset: int, end_offset: int, encrypt
                 content,
                 start_offset,
                 end_offset,
-                data.response.headers)
+                data.response.headers
+            )
         except Exception as error:
             raise HttpResponseError(
                 message="Decryption failed.",
                 response=data.response,
-                error=error) from error
+                error=error
+            ) from error
     return content
 
 

sdk/storage/azure-storage-blob/tests/test_common_blob.py

@@ -53,11 +53,7 @@
 from devtools_testutils import FakeTokenCredential, recorded_by_proxy
 from devtools_testutils.storage import StorageRecordedTestCase
 from settings.testcase import BlobPreparer
-from test_helpers import (
-    MockStorageTransport,
-    _build_base_file_share_headers,
-    _create_file_share_oauth,
-)
+from test_helpers import _build_base_file_share_headers, _create_file_share_oauth
 
 # ------------------------------------------------------------------------------
 SMALL_BLOB_SIZE = 1024
@@ -3597,57 +3593,4 @@ def test_upload_blob_partial_stream_chunked(self, **kwargs):
         result = blob.download_blob().readall()
         assert result == data[:length]
 
-    @BlobPreparer()
-    def test_mock_transport_no_content_validation(self, **kwargs):
-        storage_account_name = kwargs.pop("storage_account_name")
-        storage_account_key = kwargs.pop("storage_account_key")
-
-        transport = MockStorageTransport()
-        blob_client = BlobClient(
-            self.account_url(storage_account_name, "blob"),
-            container_name='test_cont',
-            blob_name='test_blob',
-            credential=storage_account_key,
-            transport=transport,
-            retry_total=0
-        )
-
-        content = blob_client.download_blob()
-        assert content is not None
-
-        props = blob_client.get_blob_properties()
-        assert props is not None
-
-        data = b"Hello World!"
-        resp = blob_client.upload_blob(data, overwrite=True)
-        assert resp is not None
-
-        blob_data = blob_client.download_blob().read()
-        assert blob_data == b"Hello World!"  # data is fixed by mock transport
-
-        resp = blob_client.delete_blob()
-        assert resp is None
-
-    @BlobPreparer()
-    def test_mock_transport_with_content_validation(self, **kwargs):
-        storage_account_name = kwargs.pop("storage_account_name")
-        storage_account_key = kwargs.pop("storage_account_key")
-
-        transport = MockStorageTransport()
-        blob_client = BlobClient(
-            self.account_url(storage_account_name, "blob"),
-            container_name='test_cont',
-            blob_name='test_blob',
-            credential=storage_account_key,
-            transport=transport,
-            retry_total=0
-        )
-
-        data = b"Hello World!"
-        resp = blob_client.upload_blob(data, overwrite=True, validate_content=True)
-        assert resp is not None
-
-        blob_data = blob_client.download_blob(validate_content=True).read()
-        assert blob_data == b"Hello World!"  # data is fixed by mock transport
-
     # ------------------------------------------------------------------------------

sdk/storage/azure-storage-blob/tests/test_common_blob_async.py

@@ -55,7 +55,6 @@
 from settings.testcase import BlobPreparer
 from test_helpers_async import (
     AsyncStream,
-    MockStorageTransport,
     _build_base_file_share_headers,
     _create_file_share_oauth
 )
@@ -3525,59 +3524,4 @@ async def test_upload_blob_partial_stream_chunked(self, **kwargs):
         result = await (await blob.download_blob()).readall()
         assert result == data[:length]
 
-    @BlobPreparer()
-    async def test_mock_transport_no_content_validation(self, **kwargs):
-        storage_account_name = kwargs.pop("storage_account_name")
-        storage_account_key = kwargs.pop("storage_account_key")
-
-        transport = MockStorageTransport()
-        blob_client = BlobClient(
-            self.account_url(storage_account_name, "blob"),
-            container_name='test_cont',
-            blob_name='test_blob',
-            credential=storage_account_key,
-            transport=transport,
-            retry_total=0
-        )
-
-        content = await blob_client.download_blob()
-        assert content is not None
-
-        props = await blob_client.get_blob_properties()
-        assert props is not None
-
-        data = b"Hello Async World!"
-        stream = AsyncStream(data)
-        resp = await blob_client.upload_blob(stream, overwrite=True)
-        assert resp is not None
-
-        blob_data = await (await blob_client.download_blob()).read()
-        assert blob_data == b"Hello Async World!"  # data is fixed by mock transport
-
-        resp = await blob_client.delete_blob()
-        assert resp is None
-
-    @BlobPreparer()
-    async def test_mock_transport_with_content_validation(self, **kwargs):
-        storage_account_name = kwargs.pop("storage_account_name")
-        storage_account_key = kwargs.pop("storage_account_key")
-
-        transport = MockStorageTransport()
-        blob_client = BlobClient(
-            self.account_url(storage_account_name, "blob"),
-            container_name='test_cont',
-            blob_name='test_blob',
-            credential=storage_account_key,
-            transport=transport,
-            retry_total=0
-        )
-
-        data = b"Hello Async World!"
-        stream = AsyncStream(data)
-        resp = await blob_client.upload_blob(stream, overwrite=True, validate_content=True)
-        assert resp is not None
-
-        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
-
 # ------------------------------------------------------------------------------

sdk/storage/azure-storage-blob/tests/test_helpers.py

@@ -10,7 +10,7 @@
 from typing import Any, Dict, Optional, Tuple
 from typing_extensions import Self
 
-from azure.core.pipeline.transport import HttpTransport, RequestsTransportResponse
+from azure.core.pipeline.transport import RequestsTransport, RequestsTransportResponse
 from azure.core.rest import HttpRequest
 from azure.storage.blob._serialize import get_api_version
 from requests import Response
@@ -92,15 +92,15 @@ def tell(self):
         return self.wrapped_stream.tell()
 
 
-class MockHttpClientResponse(Response):
+class MockClientResponse(Response):
     def __init__(
         self, url: str,
         body_bytes: bytes,
         headers: Dict[str, Any],
         status: int = 200,
         reason: str = "OK"
     ) -> None:
-        super(MockHttpClientResponse).__init__()
+        super(MockClientResponse).__init__()
         self._url = url
         self._body = body_bytes
         self._content = body_bytes
@@ -113,9 +113,9 @@ def __init__(
         self.raw = HTTPResponse()
 
 
-class MockStorageTransport(HttpTransport):
+class MockLegacyTransport(RequestsTransport):
     """
-    This transport returns legacy http response objects from azure core and is
+    This transport returns http response objects from azure core pipelines and is
     intended only to test our backwards compatibility support.
     """
     def send(self, request: HttpRequest, **kwargs: Any) -> RequestsTransportResponse:
@@ -132,7 +132,7 @@ def send(self, request: HttpRequest, **kwargs: Any) -> RequestsTransportResponse
 
             rest_response = RequestsTransportResponse(
                 request=request,
-                requests_response=MockHttpClientResponse(
+                requests_response=MockClientResponse(
                     request.url,
                     b"Hello World!",
                     headers,
@@ -142,7 +142,7 @@ def send(self, request: HttpRequest, **kwargs: Any) -> RequestsTransportResponse
             # get_blob_properties
             rest_response = RequestsTransportResponse(
                 request=request,
-                requests_response=MockHttpClientResponse(
+                requests_response=MockClientResponse(
                     request.url,
                     b"",
                     {
@@ -155,7 +155,7 @@ def send(self, request: HttpRequest, **kwargs: Any) -> RequestsTransportResponse
             # upload_blob
             rest_response = RequestsTransportResponse(
                 request=request,
-                requests_response=MockHttpClientResponse(
+                requests_response=MockClientResponse(
                     request.url,
                     b"",
                     {
@@ -169,7 +169,7 @@ def send(self, request: HttpRequest, **kwargs: Any) -> RequestsTransportResponse
             # delete_blob
             rest_response = RequestsTransportResponse(
                 request=request,
-                requests_response=MockHttpClientResponse(
+                requests_response=MockClientResponse(
                     request.url,
                     b"",
                     {
@@ -180,7 +180,7 @@ def send(self, request: HttpRequest, **kwargs: Any) -> RequestsTransportResponse
                 )
             )
         else:
-            raise ValueError("The request is not accepted as part of MockStorageTransport.")
+            raise ValueError("The request is not accepted as part of MockLegacyTransport.")
         return rest_response
 
     def __enter__(self) -> Self:

sdk/storage/azure-storage-blob/tests/test_helpers_async.py

@@ -3,16 +3,20 @@
 # Licensed under the MIT License. See License.txt in the project root for
 # license information.
 # --------------------------------------------------------------------------
-
+import asyncio
 import aiohttp
+from collections import deque
 from datetime import datetime, timezone
 from io import IOBase, UnsupportedOperation
 from typing import Any, Dict, Optional, Tuple
+from unittest.mock import Mock, AsyncMock
 
 from azure.core.pipeline.transport import AioHttpTransportResponse, AsyncHttpTransport
 from azure.core.rest import HttpRequest
 from azure.storage.blob._serialize import get_api_version
 from aiohttp import ClientResponse
+from aiohttp.streams import StreamReader
+from aiohttp.client_proto import ResponseHandler
 
 
 def _build_base_file_share_headers(bearer_token_string: str, content_length: int = 0) -> Dict[str, Any]:
@@ -126,11 +130,15 @@ def __init__(
         self._loop = None
         self.status = status
         self.reason = reason
+        self.content = StreamReader(ResponseHandler(asyncio.get_event_loop()), 65535)
+        self.content.total_bytes = len(body_bytes)
+        self.content._buffer = deque([body_bytes])
+        self.content._eof = True
 
 
-class MockStorageTransport(AsyncHttpTransport):
+class MockLegacyTransport(AsyncHttpTransport):
     """
-    This transport returns legacy http response objects from azure core and is 
+    This transport returns legacy http response objects from azure core and is
     intended only to test our backwards compatibility support.
     """
     async def send(self, request: HttpRequest, **kwargs: Any) -> AioHttpTransportResponse:
@@ -199,7 +207,7 @@ async def send(self, request: HttpRequest, **kwargs: Any) -> AioHttpTransportRes
                 decompress=False
             )
         else:
-            raise ValueError("The request is not accepted as part of MockStorageTransport.")
+            raise ValueError("The request is not accepted as part of MockLegacyTransport.")
 
         await rest_response.load_body()
         return rest_response