[ServiceBus] fix test import failures

sdk/servicebus/azure-servicebus/tests/__init__.py

@@ -0,0 +1,4 @@
+# ------------------------------------
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+# ----------------------------------

sdk/servicebus/azure-servicebus/tests/async_tests/__init__.py

@@ -0,0 +1,4 @@
+# ------------------------------------
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+# ----------------------------------

sdk/servicebus/azure-servicebus/tests/async_tests/test_queues_async.py

@@ -59,15 +59,15 @@
     OperationTimeoutError,
 )
 from devtools_testutils import AzureMgmtRecordedTestCase, AzureRecordedTestCase, get_credential
-from tests.servicebus_preparer import (
+from ..servicebus_preparer import (
     SERVICEBUS_ENDPOINT_SUFFIX,
     CachedServiceBusNamespacePreparer,
     CachedServiceBusQueuePreparer,
     ServiceBusQueuePreparer,
     CachedServiceBusResourceGroupPreparer,
 )
-from mocks_async import MockReceivedMessage, MockReceiver
-from tests.utilities import (
+from .mocks_async import MockReceivedMessage, MockReceiver
+from ..utilities import (
     get_logger,
     print_message,
     sleep_until_expired,

sdk/servicebus/azure-servicebus/tests/async_tests/test_sb_client_async.py

@@ -28,7 +28,7 @@
     ServiceBusConnectionError,
 )
 from devtools_testutils import AzureMgmtRecordedTestCase, get_credential
-from tests.servicebus_preparer import (
+from ..servicebus_preparer import (
     CachedServiceBusNamespacePreparer,
     ServiceBusTopicPreparer,
     ServiceBusQueuePreparer,
@@ -40,7 +40,7 @@
     CachedServiceBusResourceGroupPreparer,
     SERVICEBUS_ENDPOINT_SUFFIX,
 )
-from tests.utilities import (
+from ..utilities import (
     get_logger,
     uamqp_transport as get_uamqp_transport,
     ArgPasserAsync,

sdk/servicebus/azure-servicebus/tests/async_tests/test_sessions_async.py

@@ -34,7 +34,7 @@
     AutoLockRenewTimeout,
 )
 from devtools_testutils import AzureMgmtRecordedTestCase, get_credential
-from tests.servicebus_preparer import (
+from ..servicebus_preparer import (
     SERVICEBUS_ENDPOINT_SUFFIX,
     CachedServiceBusNamespacePreparer,
     CachedServiceBusQueuePreparer,
@@ -43,7 +43,7 @@
     ServiceBusSubscriptionPreparer,
     CachedServiceBusResourceGroupPreparer,
 )
-from tests.utilities import get_logger, print_message, uamqp_transport as get_uamqp_transport, ArgPasserAsync
+from ..utilities import get_logger, print_message, uamqp_transport as get_uamqp_transport, ArgPasserAsync
 
 uamqp_transport_params, uamqp_transport_ids = get_uamqp_transport()
 

sdk/servicebus/azure-servicebus/tests/async_tests/test_subscriptions_async.py

@@ -19,7 +19,7 @@
 from azure.servicebus._common.constants import ServiceBusSubQueue
 
 from devtools_testutils import AzureMgmtRecordedTestCase, RandomNameResourceGroupPreparer, get_credential
-from tests.servicebus_preparer import (
+from ..servicebus_preparer import (
     CachedServiceBusNamespacePreparer,
     CachedServiceBusTopicPreparer,
     CachedServiceBusSubscriptionPreparer,
@@ -28,7 +28,7 @@
     CachedServiceBusResourceGroupPreparer,
     SERVICEBUS_ENDPOINT_SUFFIX,
 )
-from tests.utilities import get_logger, print_message, uamqp_transport as get_uamqp_transport, ArgPasserAsync
+from ..utilities import get_logger, print_message, uamqp_transport as get_uamqp_transport, ArgPasserAsync
 
 uamqp_transport_params, uamqp_transport_ids = get_uamqp_transport()
 

sdk/servicebus/azure-servicebus/tests/async_tests/test_topic_async.py

@@ -17,15 +17,15 @@
 from azure.servicebus.aio import ServiceBusClient
 from azure.servicebus.aio._base_handler_async import ServiceBusSharedKeyCredential
 from azure.servicebus._common.message import ServiceBusMessage
-from tests.servicebus_preparer import (
+from ..servicebus_preparer import (
     ServiceBusNamespacePreparer,
     ServiceBusTopicPreparer,
     CachedServiceBusNamespacePreparer,
     CachedServiceBusTopicPreparer,
     CachedServiceBusResourceGroupPreparer,
     SERVICEBUS_ENDPOINT_SUFFIX,
 )
-from tests.utilities import get_logger, print_message, uamqp_transport as get_uamqp_transport, ArgPasserAsync
+from ..utilities import get_logger, print_message, uamqp_transport as get_uamqp_transport, ArgPasserAsync
 
 uamqp_transport_params, uamqp_transport_ids = get_uamqp_transport()
 

sdk/servicebus/azure-servicebus/tests/async_tests/utilities.py

@@ -0,0 +1,116 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License. See License.txt in the project root for
+# license information.
+# --------------------------------------------------------------------------
+
+import logging
+import sys
+import time
+import os
+
+try:
+    import uamqp
+
+    uamqp_available = True
+except (ModuleNotFoundError, ImportError):
+    uamqp_available = False
+from azure.servicebus import TransportType
+from azure.servicebus._common.utils import utc_now
+
+# temporary - disable uamqp if China b/c of 8+ hr runtime
+uamqp_available = uamqp_available and os.environ.get("SERVICEBUS_ENDPOINT_SUFFIX") != ".servicebus.chinacloudapi.cn"
+test_pyamqp = os.environ.get("TEST_PYAMQP", "true") == "true"
+
+
+def _get_default_handler():
+    handler = logging.StreamHandler(stream=sys.stdout)
+    handler.setFormatter(logging.Formatter("%(asctime)s %(name)-12s %(levelname)-8s %(message)s"))
+    return handler
+
+
+def _build_logger(name, level):
+    logger = logging.getLogger(name)
+    if not logger.handlers:
+        logger.setLevel(level)
+        handler = _get_default_handler()
+        logger.addHandler(handler)
+    return logger
+
+
+# Note: This was the initial generic logger entry point, kept to allow us to
+# move to more fine-grained logging controls incrementally.
+def get_logger(level, amqp_level=logging.INFO):
+    _build_logger("azure.servicebus._pyamqp", amqp_level)
+    _build_logger("uamqp", amqp_level)
+    return _build_logger("azure", level)
+
+
+def print_message(_logger, message):
+    _logger.info(f"Receiving: {message}")
+    _logger.debug(f"Time to live: {message.time_to_live}")
+    _logger.debug(f"Sequence number: {message.sequence_number}")
+    _logger.debug(f"Enqueue Sequence numger: {message.enqueued_sequence_number}")
+    _logger.debug(f"Partition Key: {message.partition_key}")
+    _logger.debug(f"Application Properties: {message.application_properties}")
+    _logger.debug(f"Delivery count: {message.delivery_count}")
+    try:
+        _logger.debug(f"Locked until: {message.locked_until_utc}")
+        _logger.debug(f"Lock Token: {message.lock_token}")
+    except (TypeError, AttributeError):
+        pass
+    _logger.debug(f"Enqueued time: {message.enqueued_time_utc}")
+
+
+def sleep_until_expired(entity):
+    time.sleep(max(0, (entity.locked_until_utc - utc_now()).total_seconds() + 1))
+
+
+def uamqp_transport(use_uamqp=uamqp_available, use_pyamqp=test_pyamqp):
+    uamqp_transport_params = []
+    uamqp_transport_ids = []
+    if use_uamqp:
+        uamqp_transport_params.append(True)
+        uamqp_transport_ids.append("uamqp")
+    if use_pyamqp:
+        uamqp_transport_params.append(False)
+        uamqp_transport_ids.append("pyamqp")
+    return uamqp_transport_params, uamqp_transport_ids
+
+
+def socket_transport():
+    socket_transport_params = [TransportType.Amqp, TransportType.AmqpOverWebsocket]
+    socket_transport_ids = ["amqp", "ws"]
+    return socket_transport_params, socket_transport_ids
+
+
+class ArgPasser:
+    def __call__(self, fn):
+        def _preparer(test_class, uamqp_transport, **kwargs):
+            fn(test_class, uamqp_transport=uamqp_transport, **kwargs)
+
+        return _preparer
+
+
+class ArgPasserAsync:
+    def __call__(self, fn):
+        async def _preparer(test_class, uamqp_transport, **kwargs):
+            await fn(test_class, uamqp_transport=uamqp_transport, **kwargs)
+
+        return _preparer
+
+
+class SocketArgPasser:
+    def __call__(self, fn):
+        def _preparer(test_class, uamqp_transport, socket_transport, **kwargs):
+            fn(test_class, uamqp_transport=uamqp_transport, socket_transport=socket_transport, **kwargs)
+
+        return _preparer
+
+
+class SocketArgPasserAsync:
+    def __call__(self, fn):
+        async def _preparer(test_class, uamqp_transport, socket_transport, **kwargs):
+            await fn(test_class, uamqp_transport=uamqp_transport, socket_transport=socket_transport, **kwargs)
+
+        return _preparer

sdk/servicebus/azure-servicebus/tests/mgmt_tests/__init__.py

@@ -0,0 +1,4 @@
+# ------------------------------------
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+# ----------------------------------

sdk/servicebus/azure-servicebus/tests/mgmt_tests/async/__init__.py

@@ -0,0 +1,4 @@
+# ------------------------------------
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+# ----------------------------------

sdk/servicebus/azure-servicebus/tests/mgmt_tests/async/test_mgmt_namespaces_async.py

@@ -9,7 +9,7 @@
 
 from devtools_testutils import AzureMgmtRecordedTestCase, get_credential
 from devtools_testutils.aio import recorded_by_proxy_async
-from tests.sb_env_loader import ServiceBusPreparer
+from ...sb_env_loader import ServiceBusPreparer
 
 
 class TestServiceBusManagementClientNamespaceAsync(AzureMgmtRecordedTestCase):

sdk/servicebus/azure-servicebus/tests/mgmt_tests/async/test_mgmt_queues_async.py

@@ -15,10 +15,10 @@
 
 from devtools_testutils import AzureMgmtRecordedTestCase, get_credential
 from devtools_testutils.aio import recorded_by_proxy_async
-from tests.sb_env_loader import ServiceBusPreparer
-from tests.servicebus_preparer import SERVICEBUS_ENDPOINT_SUFFIX
+from ...sb_env_loader import ServiceBusPreparer
+from ...servicebus_preparer import SERVICEBUS_ENDPOINT_SUFFIX
 
-from mgmt_test_utilities_async import (
+from .mgmt_test_utilities_async import (
     AsyncMgmtQueueListTestHelper,
     AsyncMgmtQueueListRuntimeInfoTestHelper,
     run_test_async_mgmt_list_with_parameters,

sdk/servicebus/azure-servicebus/tests/mgmt_tests/async/test_mgmt_rules_async.py

@@ -17,7 +17,7 @@
     SqlRuleAction,
 )
 from azure.servicebus.management._constants import INT32_MAX_VALUE
-from tests.utilities import get_logger
+from ...utilities import get_logger
 from azure.core.exceptions import HttpResponseError, ResourceExistsError
 
 from devtools_testutils import (
@@ -27,9 +27,9 @@
     get_credential,
 )
 from devtools_testutils.aio import recorded_by_proxy_async
-from tests.sb_env_loader import ServiceBusPreparer
+from ...sb_env_loader import ServiceBusPreparer
 
-from mgmt_test_utilities_async import async_pageable_to_list, clear_topics
+from .mgmt_test_utilities_async import async_pageable_to_list, clear_topics
 
 _logger = get_logger(logging.DEBUG)
 

sdk/servicebus/azure-servicebus/tests/mgmt_tests/async/test_mgmt_subscriptions_async.py

@@ -9,15 +9,15 @@
 
 from azure.servicebus.aio.management import ServiceBusAdministrationClient
 from azure.servicebus.management import SubscriptionProperties
-from tests.utilities import get_logger
+from ...utilities import get_logger
 from azure.core.exceptions import HttpResponseError, ResourceExistsError
 
 from devtools_testutils import AzureMgmtRecordedTestCase, CachedResourceGroupPreparer, get_credential
 from devtools_testutils.aio import recorded_by_proxy_async
-from tests.sb_env_loader import ServiceBusPreparer
-from tests.servicebus_preparer import SERVICEBUS_ENDPOINT_SUFFIX
+from ...sb_env_loader import ServiceBusPreparer
+from ...servicebus_preparer import SERVICEBUS_ENDPOINT_SUFFIX
 
-from mgmt_test_utilities_async import async_pageable_to_list, clear_topics
+from .mgmt_test_utilities_async import async_pageable_to_list, clear_topics
 
 _logger = get_logger(logging.DEBUG)
 

sdk/servicebus/azure-servicebus/tests/mgmt_tests/async/test_mgmt_topics_async.py

@@ -13,14 +13,14 @@
 from azure.servicebus.management import TopicProperties, AccessRights, AuthorizationRule
 from azure.servicebus.aio._base_handler_async import ServiceBusSharedKeyCredential
 from azure.servicebus.management import ApiVersion
-from tests.utilities import get_logger
+from ...utilities import get_logger
 from azure.core.exceptions import HttpResponseError, ResourceExistsError
 
 from devtools_testutils import AzureMgmtRecordedTestCase, CachedResourceGroupPreparer, get_credential
 from devtools_testutils.aio import recorded_by_proxy_async
-from tests.sb_env_loader import ServiceBusPreparer
+from ...sb_env_loader import ServiceBusPreparer
 
-from mgmt_test_utilities_async import async_pageable_to_list, clear_topics
+from .mgmt_test_utilities_async import async_pageable_to_list, clear_topics
 
 _logger = get_logger(logging.DEBUG)
 

sdk/servicebus/azure-servicebus/tests/mgmt_tests/test_mgmt_namespaces.py

@@ -8,7 +8,7 @@
 from azure.servicebus.management import ServiceBusAdministrationClient
 
 from devtools_testutils import AzureMgmtRecordedTestCase, CachedResourceGroupPreparer, recorded_by_proxy, get_credential
-from tests.sb_env_loader import ServiceBusPreparer
+from ..sb_env_loader import ServiceBusPreparer
 
 
 class TestServiceBusManagementClientNamespace(AzureMgmtRecordedTestCase):

sdk/servicebus/azure-servicebus/tests/mgmt_tests/test_mgmt_queues.py

@@ -14,17 +14,17 @@
 
 from azure.servicebus.management import ServiceBusAdministrationClient, QueueProperties, ApiVersion, AuthorizationRule, AccessRights
 from azure.servicebus._common.utils import utc_now
-from tests.utilities import get_logger
+from ..utilities import get_logger
 from azure.core.exceptions import HttpResponseError, ResourceNotFoundError, ResourceExistsError
 from azure.servicebus._base_handler import ServiceBusSharedKeyCredential
 
 from devtools_testutils import AzureMgmtRecordedTestCase, CachedResourceGroupPreparer, recorded_by_proxy, get_credential
-from tests.sb_env_loader import (
+from ..sb_env_loader import (
     ServiceBusPreparer,
 )
-from tests.servicebus_preparer import SERVICEBUS_ENDPOINT_SUFFIX
+from ..servicebus_preparer import SERVICEBUS_ENDPOINT_SUFFIX
 
-from mgmt_test_utilities import (
+from .mgmt_test_utilities import (
     MgmtQueueListTestHelper,
     MgmtQueueListRuntimeInfoTestHelper,
     run_test_mgmt_list_with_parameters,

sdk/servicebus/azure-servicebus/tests/mgmt_tests/test_mgmt_rules.py

@@ -17,7 +17,7 @@
     SqlRuleAction,
 )
 from azure.servicebus.management._constants import INT32_MAX_VALUE
-from tests.utilities import get_logger
+from ..utilities import get_logger
 from azure.core.exceptions import HttpResponseError, ResourceExistsError
 
 from devtools_testutils import (
@@ -27,9 +27,9 @@
     set_bodiless_matcher,
     get_credential,
 )
-from tests.sb_env_loader import ServiceBusPreparer
+from ..sb_env_loader import ServiceBusPreparer
 
-from mgmt_test_utilities import clear_topics
+from .mgmt_test_utilities import clear_topics
 
 _logger = get_logger(logging.DEBUG)
 

sdk/servicebus/azure-servicebus/tests/mgmt_tests/test_mgmt_subscriptions.py

@@ -8,14 +8,14 @@
 import datetime
 
 from azure.servicebus.management import ServiceBusAdministrationClient, SubscriptionProperties
-from tests.utilities import get_logger
+from ..utilities import get_logger
 from azure.core.exceptions import HttpResponseError, ResourceExistsError
 
 from devtools_testutils import AzureMgmtRecordedTestCase, CachedResourceGroupPreparer, recorded_by_proxy, get_credential
-from tests.sb_env_loader import ServiceBusPreparer
-from tests.servicebus_preparer import SERVICEBUS_ENDPOINT_SUFFIX
+from ..sb_env_loader import ServiceBusPreparer
+from ..servicebus_preparer import SERVICEBUS_ENDPOINT_SUFFIX
 
-from mgmt_test_utilities import clear_topics
+from .mgmt_test_utilities import clear_topics
 
 _logger = get_logger(logging.DEBUG)