fix bug in test hacking

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

@@ -1471,6 +1471,7 @@ async def _hack_amqp_sender_run_async(cls):
         async with ServiceBusClient.from_connection_string(
                 servicebus_namespace_connection_string, logging_enable=False) as sb_client:
             async with sb_client.get_queue_sender(servicebus_queue.name) as sender:
+                # this one doesn't need to reset the method, as it's hacking the method on the instance
                 sender._handler._client_run_async = types.MethodType(_hack_amqp_sender_run_async, sender._handler)
                 with pytest.raises(OperationTimeoutError):
                     await sender.send_messages(Message("body"), timeout=5)
@@ -1498,6 +1499,8 @@ async def hack_mgmt_execute_async(self, operation, op_type, message, timeout=0):
             response = self._responses.pop(operation_id)
             return response
 
+        original_execute_method = uamqp.async_ops.mgmt_operation_async.MgmtOperationAsync.execute_async
+        # hack the mgmt method on the class, not on an instance, so it needs reset
         uamqp.async_ops.mgmt_operation_async.MgmtOperationAsync.execute_async = hack_mgmt_execute_async
         async with ServiceBusClient.from_connection_string(
                 servicebus_namespace_connection_string, logging_enable=False) as sb_client:
@@ -1506,3 +1509,5 @@ async def hack_mgmt_execute_async(self, operation, op_type, message, timeout=0):
                     scheduled_time_utc = utc_now() + timedelta(seconds=30)
                     await sender.schedule_messages(Message("Message to be scheduled"), scheduled_time_utc, timeout=5)
 
+        # must reset the mgmt execute method, otherwise other test cases would use the hacked execute method, leading to timeout error
+        uamqp.async_ops.mgmt_operation_async.MgmtOperationAsync.execute_async = original_execute_method

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

@@ -1937,6 +1937,7 @@ def _hack_amqp_sender_run(cls):
         with ServiceBusClient.from_connection_string(
                 servicebus_namespace_connection_string, logging_enable=False) as sb_client:
             with sb_client.get_queue_sender(servicebus_queue.name) as sender:
+                # this one doesn't need to reset the method, as it's hacking the method on the instance
                 sender._handler._client_run = types.MethodType(_hack_amqp_sender_run, sender._handler)
                 with pytest.raises(OperationTimeoutError):
                     sender.send_messages(Message("body"), timeout=5)
@@ -1964,10 +1965,15 @@ def hack_mgmt_execute(self, operation, op_type, message, timeout=0):
             response = self._responses.pop(operation_id)
             return response
 
+        original_execute_method = uamqp.mgmt_operation.MgmtOperation.execute
+        # hack the mgmt method on the class, not on an instance, so it needs reset
         uamqp.mgmt_operation.MgmtOperation.execute = hack_mgmt_execute
         with ServiceBusClient.from_connection_string(
                 servicebus_namespace_connection_string, logging_enable=False) as sb_client:
             with sb_client.get_queue_sender(servicebus_queue.name) as sender:
                 with pytest.raises(OperationTimeoutError):
                     scheduled_time_utc = utc_now() + timedelta(seconds=30)
-                    sender.schedule_messages(Message("Message to be scheduled"), scheduled_time_utc, timeout=5)
+                    sender.schedule_messages(Message("Message to be scheduled"), scheduled_time_utc, timeout=5)
+
+        # must reset the mgmt execute method, otherwise other test cases would use the hacked execute method, leading to timeout error
+        uamqp.mgmt_operation.MgmtOperation.execute = original_execute_method