[SB] Fix Peek Last Sequence ID Capture (Azure#39513)

* set last sequence number for next call

* live test

* fix test

* fix increment

* pylint fixes

* make change in mgmt common

* fix test

* fix test with async cred

* additional testing

* fix sync tests

Author: kashifkhan

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

@@ -52,12 +52,14 @@ def peek_op(  # pylint: disable=inconsistent-return-statements
 ):
     condition = message.application_properties.get(MGMT_RESPONSE_MESSAGE_ERROR_CONDITION)
     if status_code == 200:
-        return amqp_transport.parse_received_message(
+        parsed_messages = amqp_transport.parse_received_message(
             message, message_type=ServiceBusReceivedMessage, receiver=receiver, is_peeked_message=True
         )
+        if parsed_messages:
+            receiver._last_received_sequenced_number = parsed_messages[-1].sequence_number # pylint: disable=protected-access
+        return parsed_messages
     if status_code in [202, 204]:
         return []
-
     amqp_transport.handle_amqp_mgmt_error(
         _LOGGER, "Message peek failed.", condition, description, status_code
     )

sdk/servicebus/azure-servicebus/azure/servicebus/_servicebus_receiver.py

@@ -72,7 +72,7 @@
 _LOGGER = logging.getLogger(__name__)
 
 
-class ServiceBusReceiver(BaseHandler, ReceiverMixin):
+class ServiceBusReceiver(BaseHandler, ReceiverMixin): # pylint: disable=too-many-instance-attributes
     """The ServiceBusReceiver class defines a high level interface for
     receiving messages from the Azure Service Bus Queue or Topic Subscription.
 
@@ -769,7 +769,11 @@ def peek_messages(
         if timeout is not None and timeout <= 0:
             raise ValueError("The timeout must be greater than 0.")
         if not sequence_number:
-            sequence_number = self._last_received_sequenced_number or 1
+            sequence_number = (
+                self._last_received_sequenced_number + 1
+                if self._last_received_sequenced_number
+                else 1
+            )
         if int(max_message_count) < 0:
             raise ValueError("max_message_count must be 1 or greater.")
 

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

@@ -748,7 +748,11 @@ async def peek_messages(
         if timeout is not None and timeout <= 0:
             raise ValueError("The timeout must be greater than 0.")
         if not sequence_number:
-            sequence_number = self._last_received_sequenced_number or 1
+            sequence_number = (
+                self._last_received_sequenced_number + 1
+                if self._last_received_sequenced_number
+                else 1
+            )
         if int(max_message_count) < 0:
             raise ValueError("max_message_count must be 1 or greater.")
 

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

@@ -3682,3 +3682,117 @@ async def test_queue_complete_message_on_different_receiver_async(
                 messages_in_queue = await receiver1.peek_messages()
 
                 assert len(messages_in_queue) == 0
+    
+    @pytest.mark.asyncio
+    @pytest.mark.liveTest
+    @pytest.mark.live_test_only
+    @CachedServiceBusResourceGroupPreparer(name_prefix="servicebustest")
+    @CachedServiceBusNamespacePreparer(name_prefix="servicebustest")
+    @ServiceBusQueuePreparer(name_prefix="servicebustest", dead_lettering_on_message_expiration=True)
+    @pytest.mark.parametrize("uamqp_transport", uamqp_transport_params, ids=uamqp_transport_ids)
+    @ArgPasserAsync()
+    async def test_queue_async_by_queue_client_peek_auto_increment(
+        self, uamqp_transport, *, servicebus_namespace=None, servicebus_queue=None, **kwargs
+    ):
+        fully_qualified_namespace = f"{servicebus_namespace.name}{SERVICEBUS_ENDPOINT_SUFFIX}"
+        credential = get_credential(is_async=True)
+        async with ServiceBusClient(
+            fully_qualified_namespace, credential, logging_enable=False, uamqp_transport=uamqp_transport
+        ) as sb_client:
+
+            sender = sb_client.get_queue_sender(servicebus_queue.name)
+            async with sender:
+                messages = []
+                for i in range(3):
+                    message = ServiceBusMessage("Handler message no. {}".format(i), application_properties={"index": i})
+                    messages.append(message)
+                await sender.send_messages(messages)
+
+            receiver = sb_client.get_queue_receiver(servicebus_queue.name, max_wait_time=5)
+            async with receiver:
+                peek_message = await receiver.peek_messages()
+                assert peek_message[0].application_properties[b"index"] == 0
+                assert peek_message[0].sequence_number == 1
+                peek_message = await receiver.peek_messages()
+                assert peek_message[0].application_properties[b"index"] == 1
+                assert peek_message[0].sequence_number == 2
+                peek_message = await receiver.peek_messages()
+                assert peek_message[0].application_properties[b"index"] == 2
+                assert peek_message[0].sequence_number == 3
+    
+    @pytest.mark.asyncio
+    @pytest.mark.liveTest
+    @pytest.mark.live_test_only
+    @CachedServiceBusResourceGroupPreparer(name_prefix="servicebustest")
+    @CachedServiceBusNamespacePreparer(name_prefix="servicebustest")
+    @ServiceBusQueuePreparer(name_prefix="servicebustest", dead_lettering_on_message_expiration=True)
+    @pytest.mark.parametrize("uamqp_transport", uamqp_transport_params, ids=uamqp_transport_ids)
+    @ArgPasserAsync()
+    async def test_queue_async_by_queue_client_peek_auto_increment_multiple(
+        self, uamqp_transport, *, servicebus_namespace=None, servicebus_queue=None, **kwargs
+    ):
+        fully_qualified_namespace = f"{servicebus_namespace.name}{SERVICEBUS_ENDPOINT_SUFFIX}"
+        credential = get_credential(is_async=True)
+        async with ServiceBusClient(
+            fully_qualified_namespace, credential, logging_enable=False, uamqp_transport=uamqp_transport
+        ) as sb_client:
+
+            sender = sb_client.get_queue_sender(servicebus_queue.name)
+            async with sender:
+                messages = []
+                for i in range(4):
+                    message = ServiceBusMessage("Handler message no. {}".format(i), application_properties={"index": i})
+                    messages.append(message)
+                await sender.send_messages(messages)
+
+            receiver = sb_client.get_queue_receiver(servicebus_queue.name, max_wait_time=5)
+            async with receiver:
+                peek_message = await receiver.peek_messages(max_message_count=2)
+                assert len(peek_message) == 2
+                assert peek_message[0].application_properties[b"index"] == 0
+                assert peek_message[0].sequence_number == 1
+                assert peek_message[1].application_properties[b"index"] == 1
+                assert peek_message[1].sequence_number == 2
+                peek_message = await receiver.peek_messages(max_message_count=2)
+                assert len(peek_message) == 2
+                assert peek_message[0].application_properties[b"index"] == 2
+                assert peek_message[0].sequence_number == 3
+                assert peek_message[1].application_properties[b"index"] == 3
+                assert peek_message[1].sequence_number == 4
+    
+    @pytest.mark.asyncio
+    @pytest.mark.liveTest
+    @pytest.mark.live_test_only
+    @CachedServiceBusResourceGroupPreparer(name_prefix="servicebustest")
+    @CachedServiceBusNamespacePreparer(name_prefix="servicebustest")
+    @ServiceBusQueuePreparer(name_prefix="servicebustest", dead_lettering_on_message_expiration=True)
+    @pytest.mark.parametrize("uamqp_transport", uamqp_transport_params, ids=uamqp_transport_ids)
+    @ArgPasserAsync()
+    async def test_queue_async_by_queue_client_peek_and_receive(
+        self, uamqp_transport, *, servicebus_namespace=None, servicebus_queue=None, **kwargs
+    ):
+        fully_qualified_namespace = f"{servicebus_namespace.name}{SERVICEBUS_ENDPOINT_SUFFIX}"
+        credential = get_credential(is_async=True)
+        async with ServiceBusClient(
+            fully_qualified_namespace, credential, logging_enable=False, uamqp_transport=uamqp_transport
+        ) as sb_client:
+
+            sender = sb_client.get_queue_sender(servicebus_queue.name)
+            async with sender:
+                messages = []
+                for i in range(4):
+                    message = ServiceBusMessage("Handler message no. {}".format(i), application_properties={"index": i})
+                    messages.append(message)
+                await sender.send_messages(messages)
+
+            receiver = sb_client.get_queue_receiver(servicebus_queue.name, max_wait_time=5)
+            async with receiver:
+                peeked_messages = await receiver.peek_messages(max_message_count=2)
+
+                messages = await receiver.receive_messages(max_message_count=3)
+                last_received_sequnece_number = messages[-1].sequence_number
+                for message in messages:
+                    await receiver.complete_message(message)
+
+                peeked_messages = await receiver.peek_messages(max_message_count=2)
+                assert peeked_messages[0].sequence_number == last_received_sequnece_number + 1

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

@@ -3720,3 +3720,117 @@ def test_queue_complete_message_on_different_receiver(
                 messages_in_queue = receiver1.peek_messages()
 
                 assert len(messages_in_queue) == 0
+    
+    @pytest.mark.liveTest
+    @pytest.mark.live_test_only
+    @CachedServiceBusResourceGroupPreparer(name_prefix="servicebustest")
+    @CachedServiceBusNamespacePreparer(name_prefix="servicebustest")
+    @ServiceBusQueuePreparer(name_prefix="servicebustest", dead_lettering_on_message_expiration=True)
+    @pytest.mark.parametrize("uamqp_transport", uamqp_transport_params, ids=uamqp_transport_ids)
+    @ArgPasser()
+    def test_queue_by_queue_client_peek_auto_increment(
+        self, uamqp_transport, *, servicebus_namespace=None, servicebus_queue=None, **kwargs
+    ):
+        fully_qualified_namespace = f"{servicebus_namespace.name}{SERVICEBUS_ENDPOINT_SUFFIX}"
+        credential = get_credential()
+        with ServiceBusClient(
+            fully_qualified_namespace, credential, logging_enable=False, uamqp_transport=uamqp_transport
+        ) as sb_client:
+
+            sender = sb_client.get_queue_sender(servicebus_queue.name)
+            with sender:
+                messages = []
+                for i in range(3):
+                    message = ServiceBusMessage("Handler message no. {}".format(i), application_properties={"index": i})
+                    messages.append(message)
+                sender.send_messages(messages)
+
+            receiver = sb_client.get_queue_receiver(servicebus_queue.name, max_wait_time=5)
+            with receiver:
+                peek_message = receiver.peek_messages()
+                assert peek_message[0].application_properties[b"index"] == 0
+                assert peek_message[0].sequence_number == 1
+                peek_message = receiver.peek_messages()
+                assert peek_message[0].application_properties[b"index"] == 1
+                assert peek_message[0].sequence_number == 2
+                peek_message = receiver.peek_messages()
+                assert peek_message[0].application_properties[b"index"] == 2
+                assert peek_message[0].sequence_number == 3
+    
+    @pytest.mark.liveTest
+    @pytest.mark.live_test_only
+    @CachedServiceBusResourceGroupPreparer(name_prefix="servicebustest")
+    @CachedServiceBusNamespacePreparer(name_prefix="servicebustest")
+    @ServiceBusQueuePreparer(name_prefix="servicebustest", dead_lettering_on_message_expiration=True)
+    @pytest.mark.parametrize("uamqp_transport", uamqp_transport_params, ids=uamqp_transport_ids)
+    @ArgPasser()
+    def test_queue_by_queue_client_peek_auto_increment_multiple(
+        self, uamqp_transport, *, servicebus_namespace=None, servicebus_queue=None, **kwargs
+    ):
+        fully_qualified_namespace = f"{servicebus_namespace.name}{SERVICEBUS_ENDPOINT_SUFFIX}"
+        credential = get_credential()
+        with ServiceBusClient(
+            fully_qualified_namespace, credential, logging_enable=False, uamqp_transport=uamqp_transport
+        ) as sb_client:
+
+            sender = sb_client.get_queue_sender(servicebus_queue.name)
+            with sender:
+                messages = []
+                for i in range(4):
+                    message = ServiceBusMessage("Handler message no. {}".format(i), application_properties={"index": i})
+                    messages.append(message)
+                sender.send_messages(messages)
+
+            receiver = sb_client.get_queue_receiver(servicebus_queue.name, max_wait_time=5)
+            with receiver:
+                peek_message = receiver.peek_messages(max_message_count=2)
+                assert len(peek_message) == 2
+                assert peek_message[0].application_properties[b"index"] == 0
+                assert peek_message[0].sequence_number == 1
+                assert peek_message[1].application_properties[b"index"] == 1
+                assert peek_message[1].sequence_number == 2
+                peek_message = receiver.peek_messages(max_message_count=2)
+                assert len(peek_message) == 2
+                assert peek_message[0].application_properties[b"index"] == 2
+                assert peek_message[0].sequence_number == 3
+                assert peek_message[1].application_properties[b"index"] == 3
+                assert peek_message[1].sequence_number == 4
+    
+    @pytest.mark.liveTest
+    @pytest.mark.live_test_only
+    @CachedServiceBusResourceGroupPreparer(name_prefix="servicebustest")
+    @CachedServiceBusNamespacePreparer(name_prefix="servicebustest")
+    @ServiceBusQueuePreparer(name_prefix="servicebustest", dead_lettering_on_message_expiration=True)
+    @pytest.mark.parametrize("uamqp_transport", uamqp_transport_params, ids=uamqp_transport_ids)
+    @ArgPasser()
+    def test_queue_by_queue_client_peek_and_receive(
+        self, uamqp_transport, *, servicebus_namespace=None, servicebus_queue=None, **kwargs
+    ):
+        fully_qualified_namespace = f"{servicebus_namespace.name}{SERVICEBUS_ENDPOINT_SUFFIX}"
+        credential = get_credential()
+        with ServiceBusClient(
+            fully_qualified_namespace, credential, logging_enable=False, uamqp_transport=uamqp_transport
+        ) as sb_client:
+
+            sender = sb_client.get_queue_sender(servicebus_queue.name)
+            with sender:
+                messages = []
+                for i in range(4):
+                    message = ServiceBusMessage("Handler message no. {}".format(i), application_properties={"index": i})
+                    messages.append(message)
+                sender.send_messages(messages)
+
+            receiver = sb_client.get_queue_receiver(servicebus_queue.name, max_wait_time=5)
+            with receiver:
+                receiver = sb_client.get_queue_receiver(servicebus_queue.name, max_wait_time=5)
+                peeked_messages = receiver.peek_messages(max_message_count=2)
+
+                messages = receiver.receive_messages(max_message_count=3)
+                last_received_sequnece_number = messages[-1].sequence_number
+                for message in messages:
+                    receiver.complete_message(message)
+
+                peeked_messages = receiver.peek_messages(max_message_count=2)
+                assert peeked_messages[0].sequence_number == last_received_sequnece_number + 1
+
+