[Storage] Mocking Transport Tests & Fixes (#38948)

Author: weirongw23-msft

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

@@ -46,11 +46,11 @@ 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.read()  # Load the body in memory and close the socket
+            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.content))
+                       encode_base64(StorageContentValidation.get_content_md5(response.http_response.body()))
         if response.http_response.headers['content-md5'] != computed_md5:
             return True
     return False

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

@@ -46,8 +46,8 @@
 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.read()
-    content = cast(bytes, data.response.content)
+    await data.response.load_body()
+    content = cast(bytes, data.response.body())
     if encryption.get('key') is not None or encryption.get('resolver') is not None:
         try:
             return decrypt_blob(

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

@@ -53,7 +53,7 @@
 from devtools_testutils.aio import recorded_by_proxy_async
 from devtools_testutils.storage.aio import AsyncStorageRecordedTestCase
 from settings.testcase import BlobPreparer
-from test_helpers_async import AsyncStream
+from test_helpers_async import AsyncStream, MockStorageTransport
 
 # ------------------------------------------------------------------------------
 TEST_CONTAINER_PREFIX = 'container'
@@ -3455,4 +3455,59 @@ async def test_upload_blob_partial_stream_chunked(self, **kwargs):
         # Assert
         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_async.py

@@ -4,7 +4,12 @@
 # license information.
 # --------------------------------------------------------------------------
 from io import IOBase, UnsupportedOperation
-from typing import Optional
+from typing import Any, Dict, Optional
+
+from azure.core.pipeline.transport import AsyncHttpTransport
+from azure.core.rest import HttpRequest
+from azure.core.rest._aiohttp import RestAioHttpTransportResponse
+from aiohttp import ClientResponse
 
 
 class ProgressTracker:
@@ -60,3 +65,110 @@ async def read(self, size: int = -1) -> bytes:
         self._offset += len(data)
 
         return data
+
+
+class MockAioHttpClientResponse(ClientResponse):
+    def __init__(
+        self, url: str,
+        body_bytes: bytes,
+        headers: Dict[str, Any],
+        status: int = 200,
+        reason: str = "OK"
+    ) -> None:
+        super(MockAioHttpClientResponse).__init__()
+        self._url = url
+        self._body = body_bytes
+        self._headers = headers
+        self._cache = {}
+        self._loop = None
+        self.status = status
+        self.reason = reason
+
+
+class MockStorageTransport(AsyncHttpTransport):
+    """
+    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) -> RestAioHttpTransportResponse:
+        if request.method == 'GET':
+            # download_blob
+            headers = {
+                "Content-Type": "application/octet-stream",
+                "Content-Range": "bytes 0-17/18",
+                "Content-Length": "18",
+            }
+
+            if "x-ms-range-get-content-md5" in request.headers:
+                headers["Content-MD5"] = "I3pVbaOCUTom+G9F9uKFoA=="
+
+            rest_response = RestAioHttpTransportResponse(
+                request=request,
+                internal_response=MockAioHttpClientResponse(
+                    request.url,
+                    b"Hello Async World!",
+                    headers,
+                ),
+                decompress=False
+            )
+        elif request.method == 'HEAD':
+            # get_blob_properties
+            rest_response = RestAioHttpTransportResponse(
+                request=request,
+                internal_response=MockAioHttpClientResponse(
+                    request.url,
+                    b"",
+                    {
+                        "Content-Type": "application/octet-stream",
+                        "Content-Length": "1024",
+                    },
+                ),
+                decompress=False
+            )
+        elif request.method == 'PUT':
+            # upload_blob
+            rest_response = RestAioHttpTransportResponse(
+                request=request,
+                internal_response=MockAioHttpClientResponse(
+                    request.url,
+                    b"",
+                    {
+                        "Content-Length": "0",
+                    },
+                    201,
+                    "Created"
+                ),
+                decompress=False
+            )
+        elif request.method == 'DELETE':
+            # delete_blob
+            rest_response = RestAioHttpTransportResponse(
+                request=request,
+                internal_response=MockAioHttpClientResponse(
+                    request.url,
+                    b"",
+                    {
+                        "Content-Length": "0",
+                    },
+                    202,
+                    "Accepted"
+                ),
+                decompress=False
+            )
+        else:
+            raise ValueError("The request is not accepted as part of MockStorageTransport.")
+
+        await rest_response.read()
+        return rest_response
+
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, *args):
+        pass
+
+    async def open(self):
+        pass
+
+    async def close(self):
+        pass

sdk/storage/azure-storage-file-datalake/azure/storage/filedatalake/_deserialize.py

@@ -164,7 +164,7 @@ def process_storage_error(storage_error) -> NoReturn:  # pylint:disable=too-many
             error_dict = error_body.get('error', {})
         elif not error_code:
             _LOGGER.warning(
-                'Unexpected return type % from ContentDecodePolicy.deserialize_from_http_generics.', type(error_body))
+                'Unexpected return type %s from ContentDecodePolicy.deserialize_from_http_generics.', type(error_body))
             error_dict = {'message': str(error_body)}
 
         # If we extracted from a Json or XML response

sdk/storage/azure-storage-file-datalake/azure/storage/filedatalake/_shared/policies_async.py

@@ -46,11 +46,11 @@ 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.read()  # Load the body in memory and close the socket
+            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.content))
+                            encode_base64(StorageContentValidation.get_content_md5(response.http_response.body()))
         if response.http_response.headers['content-md5'] != computed_md5:
             return True
     return False

sdk/storage/azure-storage-file-datalake/tests/test_file_async.py

@@ -35,6 +35,7 @@
 from devtools_testutils.aio import recorded_by_proxy_async
 from devtools_testutils.storage.aio import AsyncStorageRecordedTestCase
 from settings.testcase import DataLakePreparer
+from test_helpers_async import AsyncStream, MockStorageTransport
 # ------------------------------------------------------------------------------
 
 TEST_DIRECTORY_PREFIX = 'directory'
@@ -1536,6 +1537,62 @@ async def test_bad_audience_file_client(self, **kwargs):
         await fc.get_file_properties()
         await fc.upload_data(data, overwrite=True)
 
+    @DataLakePreparer()
+    async def test_mock_transport_no_content_validation(self, **kwargs):
+        datalake_storage_account_name = kwargs.pop("datalake_storage_account_name")
+        datalake_storage_account_key = kwargs.pop("datalake_storage_account_key")
+
+        transport = MockStorageTransport()
+        file_client = DataLakeFileClient(
+            self.account_url(datalake_storage_account_name, 'dfs'),
+            "filesystem/",
+            "dir/file.txt",
+            credential=datalake_storage_account_key,
+            transport=transport,
+            retry_total=0
+        )
+
+        data = await file_client.download_file()
+        assert data is not None
+
+        props = await file_client.get_file_properties()
+        assert props is not None
+
+        data = b"Hello Async World!"
+        stream = AsyncStream(data)
+        resp = await file_client.upload_data(stream, overwrite=True)
+        assert resp is not None
+
+        file_data = await (await file_client.download_file()).read()
+        assert file_data == b"Hello Async World!"  # data is fixed by mock transport
+
+        resp = await file_client.delete_file()
+        assert resp is not None
+
+    @DataLakePreparer()
+    async def test_mock_transport_with_content_validation(self, **kwargs):
+        datalake_storage_account_name = kwargs.pop("datalake_storage_account_name")
+        datalake_storage_account_key = kwargs.pop("datalake_storage_account_key")
+
+        await self._setUp(datalake_storage_account_name, datalake_storage_account_key)
+
+        transport = MockStorageTransport()
+        file_client = DataLakeFileClient(
+            self.account_url(datalake_storage_account_name, 'dfs'),
+            "filesystem/",
+            "dir/file.txt",
+            credential=datalake_storage_account_key,
+            transport=transport,
+            retry_total=0
+        )
+
+        data = b"Hello Async World!"
+        stream = AsyncStream(data)
+        resp = await file_client.upload_data(stream, overwrite=True, validate_content=True)
+        assert resp is not None
+
+        file_data = await (await file_client.download_file(validate_content=True)).read()
+        assert file_data == b"Hello Async World!"  # data is fixed by mock transport
 
 # ------------------------------------------------------------------------------
 if __name__ == '__main__':