Fix tests

apache · May 31, 2022 · 68f1ec5 · 68f1ec5
1 parent 1b96156
commit 68f1ec5
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 12 deletions.
diff --git a/tests/providers/microsoft/azure/hooks/test_asb_message.py b/tests/providers/microsoft/azure/hooks/test_asb_message.py
@@ -130,12 +130,14 @@ def test_receive_message(self, mock_sb_client):
         hook = ServiceBusMessageHook(azure_service_bus_conn_id=self.conn_id)
         hook.receive_message(self.queue_name)
         expected_calls = [
-            mock.call()
-            .__enter__()
-            .get_queue_receiver(self.queue_name)
-            .__enter__()
-            .receive_messages(max_message_count=10, max_wait_time=5)
-            .__iter__()
+            mock.call(),
+            mock.call().__enter__(),
+            mock.call().get_queue_receiver(queue_name='test_queue'),
+            mock.call().get_queue_receiver().__enter__(),
+            mock.call().get_queue_receiver().receive_messages(max_message_count=1, max_wait_time=None),
+            mock.call().get_queue_receiver().receive_messages().__iter__(),
+            mock.call().get_queue_receiver().__exit__(None, None, None),
+            mock.call().__exit__(None, None, None),
         ]
         mock_sb_client.assert_has_calls(expected_calls)
 

diff --git a/tests/providers/microsoft/azure/operators/test_azure_service_queue.py b/tests/providers/microsoft/azure/operators/test_azure_service_queue.py
@@ -221,11 +221,13 @@ def test_receive_message_queue(self, mock_get_conn):
         )
         asb_receive_queue_operator.execute(None)
         expected_calls = [
-            mock.call()
-            .__enter__()
-            .get_queue_receiver(QUEUE_NAME)
-            .__enter__()
-            .receive_messages(max_message_count=10, max_wait_time=5)
-            .__iter__()
+            mock.call(),
+            mock.call().__enter__(),
+            mock.call().get_queue_receiver(queue_name='test_queue'),
+            mock.call().get_queue_receiver().__enter__(),
+            mock.call().get_queue_receiver().receive_messages(max_message_count=10, max_wait_time=5),
+            mock.call().get_queue_receiver().receive_messages().__iter__(),
+            mock.call().get_queue_receiver().__exit__(None, None, None),
+            mock.call().__exit__(None, None, None),
         ]
         mock_get_conn.assert_has_calls(expected_calls)