azure service bus: fix TypeError when using Managed Identities (#1825)

* azure servicebus: fix TypeError: argument of type 'DefaultAzureCredential' is not iterable * Update t/unit/transport/test_azureservicebus.py * Update t/unit/transport/test_azureservicebus.py * Update t/unit/transport/test_azureservicebus.py --------- Co-authored-by: Asif Saif Uddin <[email protected]>
celery · Dec 29, 2023 · 0dec03d · 0dec03d
1 parent 8aff4e7
commit 0dec03d
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 2 deletions.
diff --git a/kombu/transport/azureservicebus.py b/kombu/transport/azureservicebus.py
@@ -481,7 +481,7 @@ def parse_uri(uri: str) -> tuple[str, str | DefaultAzureCredential |
     @classmethod
     def as_uri(cls, uri: str, include_password=False, mask='**') -> str:
         namespace, credential = cls.parse_uri(uri)
-        if ":" in credential:
+        if isinstance(credential, str) and ":" in credential:
             policy, sas_key = credential.split(':', 1)
             return 'azureservicebus://{}:{}@{}'.format(
                 policy,
@@ -490,6 +490,6 @@ def as_uri(cls, uri: str, include_password=False, mask='**') -> str:
             )
 
         return 'azureservicebus://{}@{}'.format(
-            credential,
+            credential.__class__.__name__,
             namespace
         )
diff --git a/t/unit/transport/test_azureservicebus.py b/t/unit/transport/test_azureservicebus.py
@@ -446,3 +446,18 @@ def test_basic_ack_reject_message_when_raises_exception(
         assert mock_queue.channel._get_asb_receiver.call_count == 1
         assert queue_object_mock.receiver.complete_message.call_count == 1
         assert super_basic_reject.call_count == 1
+
+
+def test_returning_sas():
+    conn = Connection(URL_CREDS_SAS, transport=azureservicebus.Transport)
+    assert conn.as_uri(True) == URL_CREDS_SAS
+
+
+def test_returning_da():
+    conn = Connection(URL_CREDS_DA, transport=azureservicebus.Transport)
+    assert conn.as_uri(True) == URL_CREDS_DA
+
+
+def test_returning_mi():
+    conn = Connection(URL_CREDS_MI, transport=azureservicebus.Transport)
+    assert conn.as_uri(True) == URL_CREDS_MI