From 55e8b7f9e732f3adecfe69bb5c994463d59426a0 Mon Sep 17 00:00:00 2001 From: Krista Pratico Date: Fri, 18 Oct 2019 16:29:23 -0700 Subject: [PATCH] fixes tests, adds docstrings to wrapper classes --- .../azure/storage/blob/_shared/base_client.py | 25 +++++++++++++++++- .../storage/blob/_shared/base_client_async.py | 25 ++++++++++++++++++ .../tests/test_common_blob.py | 16 ++++++++++++ .../tests/test_common_blob_async.py | 21 ++++++++++++++- .../azure/storage/file/_shared/base_client.py | 25 ++++++++++++++++++ .../storage/file/_shared/base_client_async.py | 25 ++++++++++++++++++ .../azure-storage-file/tests/test_share.py | 18 +++++++++++++ .../tests/test_share_async.py | 22 ++++++++++++++++ .../storage/queue/_shared/base_client.py | 26 ++++++++++++++++++- .../queue/_shared/base_client_async.py | 25 ++++++++++++++++++ .../azure-storage-queue/tests/test_queue.py | 16 ++++++++++++ .../tests/test_queue_async.py | 17 ++++++++++++ 12 files changed, 258 insertions(+), 3 deletions(-) diff --git a/sdk/storage/azure-storage-blob/azure/storage/blob/_shared/base_client.py b/sdk/storage/azure-storage-blob/azure/storage/blob/_shared/base_client.py index 32f38cbd9b5b..cb8114aafae0 100644 --- a/sdk/storage/azure-storage-blob/azure/storage/blob/_shared/base_client.py +++ b/sdk/storage/azure-storage-blob/azure/storage/blob/_shared/base_client.py @@ -28,7 +28,7 @@ from azure.core.configuration import Configuration from azure.core.exceptions import HttpResponseError from azure.core.pipeline import Pipeline -from azure.core.pipeline.transport import RequestsTransport +from azure.core.pipeline.transport import RequestsTransport, HttpTransport from azure.core.pipeline.policies import ( RedirectPolicy, ContentDecodePolicy, @@ -221,6 +221,29 @@ def _batch_send( except StorageErrorException as error: process_storage_error(error) +class TransportWrapper(HttpTransport): + """Wrapper class that ensures that an inner client created + by a `get_client` method does not close the outer transport for the parent + when used in a context manager. + """ + def __init__(self, transport): + self._transport = transport + + def send(self, request, **kwargs): + return self._transport.send(request, **kwargs) + + def open(self): + pass + + def close(self): + pass + + def __enter__(self): + pass + + def __exit__(self, *args): # pylint: disable=arguments-differ + pass + def format_shared_key_credential(account, credential): if isinstance(credential, six.string_types): diff --git a/sdk/storage/azure-storage-blob/azure/storage/blob/_shared/base_client_async.py b/sdk/storage/azure-storage-blob/azure/storage/blob/_shared/base_client_async.py index 0fe9eaa90785..215292a64f3e 100644 --- a/sdk/storage/azure-storage-blob/azure/storage/blob/_shared/base_client_async.py +++ b/sdk/storage/azure-storage-blob/azure/storage/blob/_shared/base_client_async.py @@ -18,6 +18,7 @@ AsyncRedirectPolicy, DistributedTracingPolicy ) +from azure.core.pipeline.transport import AsyncHttpTransport from .constants import STORAGE_OAUTH_SCOPE, DEFAULT_SOCKET_TIMEOUT from .authentication import SharedKeyCredentialPolicy @@ -126,3 +127,27 @@ async def _batch_send( return response.parts() # Return an AsyncIterator except StorageErrorException as error: process_storage_error(error) + + +class AsyncTransportWrapper(AsyncHttpTransport): + """Wrapper class that ensures that an inner client created + by a `get_client` method does not close the outer transport for the parent + when used in a context manager. + """ + def __init__(self, async_transport): + self._transport = async_transport + + async def send(self, request, **kwargs): + return await self._transport.send(request, **kwargs) + + async def open(self): + pass + + async def close(self): + pass + + async def __aenter__(self): + pass + + async def __aexit__(self, *args): # pylint: disable=arguments-differ + pass 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 92f66abad986..d9b09a11d593 100644 --- a/sdk/storage/azure-storage-blob/tests/test_common_blob.py +++ b/sdk/storage/azure-storage-blob/tests/test_common_blob.py @@ -17,6 +17,7 @@ ResourceNotFoundError, ResourceExistsError, ClientAuthenticationError) +from azure.core.pipeline.transport import RequestsTransport from azure.storage.blob import ( upload_blob_to_url, download_blob_from_url, @@ -1833,6 +1834,21 @@ def test_set_blob_permission(self): self.assertEqual(permission.delete, True) self.assertEqual(permission.write, True) self.assertEqual(permission._str, 'wrdx') + + def test_transport_closed_only_once(self): + if TestMode.need_recording_file(self.test_mode): + return + transport = RequestsTransport() + url = self._get_account_url() + credential = self._get_shared_key_credential() + blob_name = self._get_blob_reference() + with BlobServiceClient(url, credential=credential, transport=transport) as bsc: + bsc.get_service_properties() + assert transport.session is not None + with bsc.get_blob_client(self.container_name, blob_name) as bc: + assert transport.session is not None + bsc.get_service_properties() + assert transport.session is not None #------------------------------------------------------------------------------ if __name__ == '__main__': 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 dd732030b4de..ee1118a21cb2 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 @@ -18,7 +18,7 @@ ResourceNotFoundError, ResourceExistsError, ClientAuthenticationError) - +from azure.core.pipeline.transport import AsyncioRequestsTransport from azure.core.pipeline.transport import AioHttpTransport from multidict import CIMultiDict, CIMultiDictProxy @@ -2281,6 +2281,25 @@ def test_upload_to_url_file_with_credential(self): loop = asyncio.get_event_loop() loop.run_until_complete(self._test_upload_to_url_file_with_credential()) + async def _test_transport_closed_only_once(self): + if TestMode.need_recording_file(self.test_mode): + return + transport = AioHttpTransport() + url = self._get_account_url() + credential = self._get_shared_key_credential() + blob_name = self._get_blob_reference() + async with BlobServiceClient(url, credential=credential, transport=transport) as bsc: + await bsc.get_service_properties() + assert transport.session is not None + async with bsc.get_blob_client(self.container_name, blob_name) as bc: + assert transport.session is not None + await bsc.get_service_properties() + assert transport.session is not None + + @record + def test_transport_closed_only_once(self): + loop = asyncio.get_event_loop() + loop.run_until_complete(self._test_transport_closed_only_once()) # ------------------------------------------------------------------------------ if __name__ == '__main__': unittest.main() diff --git a/sdk/storage/azure-storage-file/azure/storage/file/_shared/base_client.py b/sdk/storage/azure-storage-file/azure/storage/file/_shared/base_client.py index 85111a60c90a..6cdf780498d6 100644 --- a/sdk/storage/azure-storage-file/azure/storage/file/_shared/base_client.py +++ b/sdk/storage/azure-storage-file/azure/storage/file/_shared/base_client.py @@ -28,6 +28,7 @@ from azure.core.configuration import Configuration from azure.core.exceptions import HttpResponseError from azure.core.pipeline import Pipeline +from azure.core.pipeline.transport import RequestsTransport, HttpTransport from azure.core.pipeline.transport import RequestsTransport from azure.core.pipeline.policies import ( RedirectPolicy, @@ -221,6 +222,30 @@ def _batch_send( process_storage_error(error) +class TransportWrapper(HttpTransport): + """Wrapper class that ensures that an inner client created + by a `get_client` method does not close the outer transport for the parent + when used in a context manager. + """ + def __init__(self, transport): + self._transport = transport + + def send(self, request, **kwargs): + return self._transport.send(request, **kwargs) + + def open(self): + pass + + def close(self): + pass + + def __enter__(self): + pass + + def __exit__(self, *args): # pylint: disable=arguments-differ + pass + + def format_shared_key_credential(account, credential): if isinstance(credential, six.string_types): if len(account) < 2: diff --git a/sdk/storage/azure-storage-file/azure/storage/file/_shared/base_client_async.py b/sdk/storage/azure-storage-file/azure/storage/file/_shared/base_client_async.py index 32a03fcf4f7a..77f6e48ef5ab 100644 --- a/sdk/storage/azure-storage-file/azure/storage/file/_shared/base_client_async.py +++ b/sdk/storage/azure-storage-file/azure/storage/file/_shared/base_client_async.py @@ -19,6 +19,7 @@ DistributedTracingPolicy ) +from azure.core.pipeline.transport import AsyncHttpTransport from .constants import STORAGE_OAUTH_SCOPE, DEFAULT_SOCKET_TIMEOUT from .authentication import SharedKeyCredentialPolicy from .base_client import create_configuration @@ -123,3 +124,27 @@ async def _batch_send( return response.parts() # Return an AsyncIterator except StorageErrorException as error: process_storage_error(error) + + +class AsyncTransportWrapper(AsyncHttpTransport): + """Wrapper class that ensures that an inner client created + by a `get_client` method does not close the outer transport for the parent + when used in a context manager. + """ + def __init__(self, async_transport): + self._transport = async_transport + + async def send(self, request, **kwargs): + return await self._transport.send(request, **kwargs) + + async def open(self): + pass + + async def close(self): + pass + + async def __aenter__(self): + pass + + async def __aexit__(self, *args): # pylint: disable=arguments-differ + pass diff --git a/sdk/storage/azure-storage-file/tests/test_share.py b/sdk/storage/azure-storage-file/tests/test_share.py index 53f6063e2e13..7d1b636ac016 100644 --- a/sdk/storage/azure-storage-file/tests/test_share.py +++ b/sdk/storage/azure-storage-file/tests/test_share.py @@ -10,6 +10,7 @@ import pytest import requests +from azure.core.pipeline.transport import RequestsTransport from azure.core.exceptions import ( HttpResponseError, ResourceNotFoundError, @@ -762,6 +763,23 @@ def test_create_permission_for_share(self): # server returned permission self.assertEquals(permission_key, permission_key2) + @record + def test_transport_closed_only_once(self): + if TestMode.need_recording_file(self.test_mode): + return + transport = RequestsTransport() + url = self.get_file_url() + credential = self.get_shared_key_credential() + prefix = TEST_SHARE_PREFIX + share_name = self.get_resource_name(prefix) + with FileServiceClient(url, credential=credential, transport=transport) as fsc: + fsc.get_service_properties() + assert transport.session is not None + with fsc.get_share_client(share_name) as fc: + assert transport.session is not None + fsc.get_service_properties() + assert transport.session is not None + # ------------------------------------------------------------------------------ if __name__ == '__main__': unittest.main() diff --git a/sdk/storage/azure-storage-file/tests/test_share_async.py b/sdk/storage/azure-storage-file/tests/test_share_async.py index fe4d3687d017..a6974b181de9 100644 --- a/sdk/storage/azure-storage-file/tests/test_share_async.py +++ b/sdk/storage/azure-storage-file/tests/test_share_async.py @@ -11,6 +11,7 @@ import pytest import requests from azure.core.pipeline.transport import AioHttpTransport +from azure.core.pipeline.transport import AsyncioRequestsTransport from multidict import CIMultiDict, CIMultiDictProxy from azure.core.exceptions import ( HttpResponseError, @@ -918,6 +919,27 @@ def test_create_permission_for_share_async(self): loop = asyncio.get_event_loop() loop.run_until_complete(self._test_create_permission_for_share()) + async def _test_transport_closed_only_once_async(self): + if TestMode.need_recording_file(self.test_mode): + return + transport = AioHttpTransport() + url = self.get_file_url() + credential = self.get_shared_key_credential() + prefix = TEST_SHARE_PREFIX + share_name = self.get_resource_name(prefix) + async with FileServiceClient(url, credential=credential, transport=transport) as fsc: + await fsc.get_service_properties() + assert transport.session is not None + async with fsc.get_share_client(share_name) as fc: + assert transport.session is not None + await fsc.get_service_properties() + assert transport.session is not None + + @record + def test_transport_closed_only_once_async(self): + loop = asyncio.get_event_loop() + loop.run_until_complete(self._test_transport_closed_only_once_async()) + # ------------------------------------------------------------------------------ if __name__ == '__main__': unittest.main() diff --git a/sdk/storage/azure-storage-queue/azure/storage/queue/_shared/base_client.py b/sdk/storage/azure-storage-queue/azure/storage/queue/_shared/base_client.py index 85111a60c90a..3ef7c427560b 100644 --- a/sdk/storage/azure-storage-queue/azure/storage/queue/_shared/base_client.py +++ b/sdk/storage/azure-storage-queue/azure/storage/queue/_shared/base_client.py @@ -28,7 +28,7 @@ from azure.core.configuration import Configuration from azure.core.exceptions import HttpResponseError from azure.core.pipeline import Pipeline -from azure.core.pipeline.transport import RequestsTransport +from azure.core.pipeline.transport import RequestsTransport, HttpTransport from azure.core.pipeline.policies import ( RedirectPolicy, ContentDecodePolicy, @@ -221,6 +221,30 @@ def _batch_send( process_storage_error(error) +class TransportWrapper(HttpTransport): + """Wrapper class that ensures that an inner client created + by a `get_client` method does not close the outer transport for the parent + when used in a context manager. + """ + def __init__(self, transport): + self._transport = transport + + def send(self, request, **kwargs): + return self._transport.send(request, **kwargs) + + def open(self): + pass + + def close(self): + pass + + def __enter__(self): + pass + + def __exit__(self, *args): # pylint: disable=arguments-differ + pass + + def format_shared_key_credential(account, credential): if isinstance(credential, six.string_types): if len(account) < 2: diff --git a/sdk/storage/azure-storage-queue/azure/storage/queue/_shared/base_client_async.py b/sdk/storage/azure-storage-queue/azure/storage/queue/_shared/base_client_async.py index 32a03fcf4f7a..77f6e48ef5ab 100644 --- a/sdk/storage/azure-storage-queue/azure/storage/queue/_shared/base_client_async.py +++ b/sdk/storage/azure-storage-queue/azure/storage/queue/_shared/base_client_async.py @@ -19,6 +19,7 @@ DistributedTracingPolicy ) +from azure.core.pipeline.transport import AsyncHttpTransport from .constants import STORAGE_OAUTH_SCOPE, DEFAULT_SOCKET_TIMEOUT from .authentication import SharedKeyCredentialPolicy from .base_client import create_configuration @@ -123,3 +124,27 @@ async def _batch_send( return response.parts() # Return an AsyncIterator except StorageErrorException as error: process_storage_error(error) + + +class AsyncTransportWrapper(AsyncHttpTransport): + """Wrapper class that ensures that an inner client created + by a `get_client` method does not close the outer transport for the parent + when used in a context manager. + """ + def __init__(self, async_transport): + self._transport = async_transport + + async def send(self, request, **kwargs): + return await self._transport.send(request, **kwargs) + + async def open(self): + pass + + async def close(self): + pass + + async def __aenter__(self): + pass + + async def __aexit__(self, *args): # pylint: disable=arguments-differ + pass diff --git a/sdk/storage/azure-storage-queue/tests/test_queue.py b/sdk/storage/azure-storage-queue/tests/test_queue.py index 0525c6f644a2..c36feaf7082e 100644 --- a/sdk/storage/azure-storage-queue/tests/test_queue.py +++ b/sdk/storage/azure-storage-queue/tests/test_queue.py @@ -18,6 +18,7 @@ ) from devtools_testutils import ResourceGroupPreparer, StorageAccountPreparer from azure.mgmt.storage.models import Endpoints +from azure.core.pipeline.transport import RequestsTransport from azure.core.exceptions import ( HttpResponseError, ResourceNotFoundError, @@ -963,6 +964,21 @@ def test_unicode_update_message_unicode_data(self, resource_group, location, sto self.assertIsInstance(message.expires_on, datetime) self.assertIsInstance(message.next_visible_on, datetime) + @ResourceGroupPreparer() + @StorageAccountPreparer(name_prefix='pyacrstorage') + def test_transport_closed_only_once(self, resource_group, location, storage_account, storage_account_key): + if not self.is_live: + return + transport = RequestsTransport() + prefix = TEST_QUEUE_PREFIX + queue_name = self.get_resource_name(prefix) + with QueueServiceClient(self._account_url(storage_account.name), credential=storage_account_key, transport=transport) as qsc: + qsc.get_service_properties() + assert transport.session is not None + with qsc.get_queue_client(queue_name) as qc: + assert transport.session is not None + qsc.get_service_properties() + assert transport.session is not None # ------------------------------------------------------------------------------ if __name__ == '__main__': diff --git a/sdk/storage/azure-storage-queue/tests/test_queue_async.py b/sdk/storage/azure-storage-queue/tests/test_queue_async.py index 3de50da06d05..8f528532f4a7 100644 --- a/sdk/storage/azure-storage-queue/tests/test_queue_async.py +++ b/sdk/storage/azure-storage-queue/tests/test_queue_async.py @@ -16,6 +16,7 @@ ) from devtools_testutils import ResourceGroupPreparer, StorageAccountPreparer from multidict import CIMultiDict, CIMultiDictProxy +from azure.core.pipeline.transport import AsyncioRequestsTransport from azure.core.exceptions import ( HttpResponseError, ResourceNotFoundError, @@ -1025,6 +1026,22 @@ async def test_unicode_update_message_unicode_data(self, resource_group, locatio self.assertIsInstance(message.expires_on, datetime) self.assertIsInstance(message.next_visible_on, datetime) + @ResourceGroupPreparer() + @StorageAccountPreparer(name_prefix='pyacrstorage') + @AsyncQueueTestCase.await_prepared_test + async def test_transport_closed_only_once_async(self, resource_group, location, storage_account, storage_account_key): + if not self.is_live: + return + transport = AioHttpTransport() + prefix = TEST_QUEUE_PREFIX + queue_name = self.get_resource_name(prefix) + async with QueueServiceClient(self._account_url(storage_account.name), credential=storage_account_key, transport=transport) as qsc: + await qsc.get_service_properties() + assert transport.session is not None + async with qsc.get_queue_client(queue_name) as qc: + assert transport.session is not None + await qsc.get_service_properties() + assert transport.session is not None # ------------------------------------------------------------------------------ if __name__ == '__main__':