Fixed base handler async/added tests for SBSender + SBReceiver (Azure#16051)

* fixed base handler async to use async credential + added test for sync/async SBSender + SBReceiver

* changed sb base handler methods for reusability

* changed base handler to take class type of credentials

* small update to kwargs call

* remove unnecessary import

Author: swathipil

sdk/servicebus/azure-servicebus/azure/servicebus/_base_handler.py

@@ -226,19 +226,17 @@ def _convert_connection_string_to_kwargs(cls, conn_str, **kwargs):
 
         kwargs["fully_qualified_namespace"] = host
         kwargs["entity_name"] = entity_in_conn_str or entity_in_kwargs
-        # This has to be defined seperately to support sync vs async credentials.
-        kwargs["credential"] = cls._create_credential_from_connection_string_parameters(
-            token, token_expiry, policy, key
-        )
-        return kwargs
 
-    @classmethod
-    def _create_credential_from_connection_string_parameters(
-        cls, token, token_expiry, policy, key
-    ):
+        # Set the type to sync credentials, unless async credentials are passed in.
+        token_cred_type = kwargs.pop("token_cred_type", ServiceBusSASTokenCredential)
+        key_cred_type = kwargs.pop("key_cred_type", ServiceBusSharedKeyCredential)
+
         if token and token_expiry:
-            return ServiceBusSASTokenCredential(token, token_expiry)
-        return ServiceBusSharedKeyCredential(policy, key)
+            kwargs["credential"] = token_cred_type(token, token_expiry)
+        else:
+            kwargs["credential"] = key_cred_type(policy, key)
+
+        return kwargs
 
     def __enter__(self):
         if self._shutdown.is_set():

sdk/servicebus/azure-servicebus/azure/servicebus/aio/_base_handler_async.py

@@ -112,15 +112,12 @@ def __init__(
     @classmethod
     def _convert_connection_string_to_kwargs(cls, conn_str, **kwargs):
         # pylint:disable=protected-access
-        return BaseHandlerSync._convert_connection_string_to_kwargs(conn_str, **kwargs)
-
-    @classmethod
-    def _create_credential_from_connection_string_parameters(
-        cls, token, token_expiry, policy, key
-    ):
-        if token and token_expiry:
-            return ServiceBusSASTokenCredential(token, token_expiry)
-        return ServiceBusSharedKeyCredential(policy, key)
+        return BaseHandlerSync._convert_connection_string_to_kwargs(
+            conn_str,
+            token_cred_type=ServiceBusSASTokenCredential,
+            key_cred_type=ServiceBusSharedKeyCredential,
+            **kwargs
+        )
 
     async def __aenter__(self):
         if self._shutdown.is_set():

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

@@ -9,7 +9,7 @@
 import pytest
 
 from azure.mgmt.servicebus.models import AccessRights
-from azure.servicebus.aio import ServiceBusClient, ServiceBusSender
+from azure.servicebus.aio import ServiceBusClient, ServiceBusSender, ServiceBusReceiver
 from azure.servicebus import ServiceBusMessage
 from azure.servicebus.aio._base_handler_async import ServiceBusSharedKeyCredential
 from azure.servicebus.exceptions import (
@@ -173,6 +173,60 @@ async def test_async_sb_client_close_spawned_handlers(self, servicebus_namespace
         assert not receiver._handler and not receiver._running
         assert len(client._handlers) == 0
 
+    @pytest.mark.liveTest
+    @pytest.mark.live_test_only
+    @CachedResourceGroupPreparer(name_prefix='servicebustest')
+    @CachedServiceBusNamespacePreparer(name_prefix='servicebustest')
+    @ServiceBusNamespaceAuthorizationRulePreparer(name_prefix='servicebustest')
+    @ServiceBusQueuePreparer(name_prefix='servicebustest_qone', parameter_name='wrong_queue', dead_lettering_on_message_expiration=True)
+    @ServiceBusQueuePreparer(name_prefix='servicebustest_qtwo', dead_lettering_on_message_expiration=True)
+    @ServiceBusQueueAuthorizationRulePreparer(name_prefix='servicebustest_qtwo')
+    async def test_sb_client_incorrect_queue_conn_str_async(self, servicebus_queue_authorization_rule_connection_string, servicebus_queue, wrong_queue, **kwargs):
+        
+        client = ServiceBusClient.from_connection_string(servicebus_queue_authorization_rule_connection_string)
+        async with client:
+            # Validate that the wrong sender/receiver queues with the right credentials fail.
+            with pytest.raises(ValueError):
+                async with client.get_queue_sender(wrong_queue.name) as sender:
+                    await sender.send_messages(ServiceBusMessage("test"))
+
+            with pytest.raises(ValueError):
+                async with client.get_queue_receiver(wrong_queue.name) as receiver:
+                    messages = await receiver.receive_messages(max_message_count=1, max_wait_time=1)
+
+            # But that the correct ones work.
+            async with client.get_queue_sender(servicebus_queue.name) as sender:
+                await sender.send_messages(ServiceBusMessage("test")) 
+
+            async with client.get_queue_receiver(servicebus_queue.name) as receiver:
+                messages = await receiver.receive_messages(max_message_count=1, max_wait_time=1)
+
+            # Now do the same but with direct connstr initialization.
+            with pytest.raises(ValueError):
+                async with ServiceBusSender._from_connection_string(
+                    servicebus_queue_authorization_rule_connection_string,
+                    queue_name=wrong_queue.name,
+                ) as sender:
+                        await sender.send_messages(ServiceBusMessage("test"))
+
+            with pytest.raises(ValueError):
+                async with ServiceBusReceiver._from_connection_string(
+                    servicebus_queue_authorization_rule_connection_string,
+                    queue_name=wrong_queue.name,
+                ) as receiver:
+                    messages = await receiver.receive_messages(max_message_count=1, max_wait_time=1)
+
+            async with ServiceBusSender._from_connection_string(
+                servicebus_queue_authorization_rule_connection_string,
+                queue_name=servicebus_queue.name,
+            ) as sender:
+                await sender.send_messages(ServiceBusMessage("test"))
+
+            async with ServiceBusReceiver._from_connection_string(
+                servicebus_queue_authorization_rule_connection_string,
+                queue_name=servicebus_queue.name,
+            ) as receiver:
+                messages = await receiver.receive_messages(max_message_count=1, max_wait_time=1)
 
     @pytest.mark.liveTest
     @pytest.mark.live_test_only

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

@@ -13,7 +13,7 @@
 
 from azure.common import AzureHttpError, AzureConflictHttpError
 from azure.mgmt.servicebus.models import AccessRights
-from azure.servicebus import ServiceBusClient, ServiceBusSender
+from azure.servicebus import ServiceBusClient, ServiceBusSender, ServiceBusReceiver
 from azure.servicebus._base_handler import ServiceBusSharedKeyCredential
 from azure.servicebus._common.message import ServiceBusMessage, ServiceBusReceivedMessage
 from azure.servicebus.exceptions import (
@@ -152,29 +152,49 @@ def test_sb_client_incorrect_queue_conn_str(self, servicebus_queue_authorization
 
         client = ServiceBusClient.from_connection_string(servicebus_queue_authorization_rule_connection_string)
         with client:
-            # Validate that the wrong queue with the right credentials fails.
+            # Validate that the wrong sender/receiver queues with the right credentials fail.
             with pytest.raises(ValueError):
                 with client.get_queue_sender(wrong_queue.name) as sender:
                     sender.send_messages(ServiceBusMessage("test"))
 
-            # But that the correct one works.
+            with pytest.raises(ValueError):
+                with client.get_queue_receiver(wrong_queue.name) as receiver:
+                    messages =  receiver.receive_messages(max_message_count=1, max_wait_time=1)
+
+            # But that the correct ones work.
             with client.get_queue_sender(servicebus_queue.name) as sender:
                 sender.send_messages(ServiceBusMessage("test")) 
 
+            with client.get_queue_receiver(servicebus_queue.name) as receiver:
+               messages =  receiver.receive_messages(max_message_count=1, max_wait_time=1)
+
             # Now do the same but with direct connstr initialization.
             with pytest.raises(ValueError):
                 with ServiceBusSender._from_connection_string(
                     servicebus_queue_authorization_rule_connection_string,
                     queue_name=wrong_queue.name,
                 ) as sender:
-                    sender.send_messages(ServiceBusMessage("test"))
+                        sender.send_messages(ServiceBusMessage("test"))
+
+            with pytest.raises(ValueError):
+                with ServiceBusReceiver._from_connection_string(
+                    servicebus_queue_authorization_rule_connection_string,
+                    queue_name=wrong_queue.name,
+                ) as receiver:
+                    messages =  receiver.receive_messages(max_message_count=1, max_wait_time=1)
 
             with ServiceBusSender._from_connection_string(
                 servicebus_queue_authorization_rule_connection_string,
                 queue_name=servicebus_queue.name,
             ) as sender:
                 sender.send_messages(ServiceBusMessage("test"))
 
+            with ServiceBusReceiver._from_connection_string(
+                servicebus_queue_authorization_rule_connection_string,
+                queue_name=servicebus_queue.name,
+            ) as receiver:
+                messages =  receiver.receive_messages(max_message_count=1, max_wait_time=1)
+
     @pytest.mark.liveTest
     @pytest.mark.live_test_only
     @CachedResourceGroupPreparer()