fix: Save lock-token from receive_deferred_messages AMQP payload, into ServiceBusReceivedMessage

sdk/servicebus/azure-servicebus/azure/servicebus/_common/message.py

@@ -775,6 +775,7 @@ def __init__(
         self._settled = receive_mode == ServiceBusReceiveMode.RECEIVE_AND_DELETE
         self._delivery_tag = self._amqp_transport.get_message_delivery_tag(message, frame)
         self._delivery_id = self._amqp_transport.get_message_delivery_id(message, frame)  # only used by pyamqp
+        self._received_lock_token = kwargs.get("lock_token", None)
         self._received_timestamp_utc = utc_now()
         self._is_deferred_message = kwargs.get("is_deferred_message", False)
         self._is_peeked_message = kwargs.get("is_peeked_message", False)
@@ -1078,6 +1079,10 @@ def lock_token(self) -> Optional[Union[uuid.UUID, str]]:
         if self._settled:
             return None
 
+        if self._received_lock_token:
+            # this will already be a uuid.UUID
+            return self._received_lock_token
+
         if self._delivery_tag:
             return uuid.UUID(bytes_le=self._delivery_tag)
 

sdk/servicebus/azure-servicebus/azure/servicebus/_transport/_pyamqp_transport.py

@@ -818,15 +818,21 @@ def parse_received_message( # pylint:disable=docstring-keyword-should-match-keyw
         :keyword ~azure.servicebus.ServiceBusReceiver receiver: Required.
         :keyword bool is_peeked_message: Optional. For peeked messages.
         :keyword bool is_deferred_message: Optional. For deferred messages.
+        :keyword uuid.UUID lock_token: Optional. Lock token, if it is given by the message receiver.
         :keyword ~azure.servicebus.ServiceBusReceiveMode receive_mode: Optional.
         :return: List of service bus received messages.
         :rtype: list[~azure.servicebus.ServiceBusReceivedMessage]
         """
+        is_deferred_message = kwargs.get("is_deferred_message", False)
         parsed = []
         if message.value:
             for m in message.value[b"messages"]:
                 wrapped = decode_payload(memoryview(m[b"message"]))
-                parsed.append(message_type(wrapped, **kwargs))
+                if is_deferred_message and b"lock-token" in m:
+                    lock_token = m[b"lock-token"]
+                else:
+                    lock_token = kwargs.pop("lock_token", None)
+                parsed.append(message_type(wrapped, lock_token=lock_token, **kwargs))
         return parsed
 
     @staticmethod

sdk/servicebus/azure-servicebus/azure/servicebus/_transport/_uamqp_transport.py

@@ -863,14 +863,20 @@ def parse_received_message( # pylint:disable=docstring-keyword-should-match-keyw
             :keyword ~azure.servicebus.ServiceBusReceiver receiver: Required.
             :keyword bool is_peeked_message: Optional. For peeked messages.
             :keyword bool is_deferred_message: Optional. For deferred messages.
+            :keyword uuid.UUID lock_token: Optional. Lock token, if it is given by the message receiver.
             :keyword ~azure.servicebus.ServiceBusReceiveMode receive_mode: Optional.
             :return: List of ServiceBusReceivedMessage.
             :rtype: list[~azure.servicebus.ServiceBusReceivedMessage]
             """
+            is_deferred_message = kwargs.get("is_deferred_message", False)
             parsed = []
             for m in message.get_data()[b"messages"]:
                 wrapped = Message.decode_from_bytes(bytearray(m[b"message"]))
-                parsed.append(message_type(wrapped, **kwargs))
+                if is_deferred_message and b"lock-token" in m:
+                    lock_token = m[b"lock-token"]
+                else:
+                    lock_token = kwargs.pop("lock_token", None)
+                parsed.append(message_type(wrapped, lock_token=lock_token, **kwargs))
             return parsed
 
         @staticmethod