diff --git a/sdk/servicebus/azure-servicebus/CHANGELOG.md b/sdk/servicebus/azure-servicebus/CHANGELOG.md index f16eabf8f771..5ce496f03c21 100644 --- a/sdk/servicebus/azure-servicebus/CHANGELOG.md +++ b/sdk/servicebus/azure-servicebus/CHANGELOG.md @@ -12,6 +12,7 @@ - Read `com.microsoft:max-message-batch-size` vendor property from the AMQP sender link to correctly limit batch size on Premium large-message entities, where `max-message-size` can be up to 100 MB but the batch limit is 1 MB. - Fixed a bug where sending a batched or multi-message payload with `uamqp_transport=True` raised `TypeError: 'BatchMessage' object is not subscriptable` (and a masked `AttributeError` on the list path) when the first message carried a `message_id`, `session_id`, or `partition_key`. The batch envelope properties are now set through the transport-appropriate code path. (regression from [#42598](https://github.com/Azure/azure-sdk-for-python/pull/42598)) - Fixed a bug where the async receiver factory methods on `azure.servicebus.aio.ServiceBusClient` (`get_queue_receiver`, `get_subscription_receiver`) and the async `ServiceBusReceiver` annotated the `auto_lock_renewer` keyword with the synchronous `AutoLockRenewer`, causing static type checkers to reject the documented `azure.servicebus.aio.AutoLockRenewer` usage. The annotation now references the async `AutoLockRenewer`, matching the docstrings and runtime behavior. ([#47948](https://github.com/Azure/azure-sdk-for-python/issues/47948)) +- Fixed a bug where closing a `PEEK_LOCK` receiver did not release messages that had been prefetched into the client buffer or were still in flight, so they remained locked at the broker until lock expiry — delaying their redelivery and inflating their delivery count. On close, a non-session `PEEK_LOCK` receiver now drains the link (stopping the broker and flushing in-flight transfers) and releases the buffered messages (`released` disposition), so the broker can redeliver them immediately without incrementing the delivery count. ([#42917](https://github.com/Azure/azure-sdk-for-python/issues/42917)) ## 7.14.3 (2025-11-11) diff --git a/sdk/servicebus/azure-servicebus/azure/servicebus/_servicebus_receiver.py b/sdk/servicebus/azure-servicebus/azure/servicebus/_servicebus_receiver.py index 3c8c1dd0c87f..074032215d58 100644 --- a/sdk/servicebus/azure-servicebus/azure/servicebus/_servicebus_receiver.py +++ b/sdk/servicebus/azure-servicebus/azure/servicebus/_servicebus_receiver.py @@ -538,6 +538,15 @@ def _renew_locks(self, *lock_tokens: str, **kwargs: Any) -> Any: def _close_handler(self): self._message_iter = None + if ( + self._handler + and self._receive_mode == ServiceBusReceiveMode.PEEK_LOCK + and not self._session + ): + # Drain the link and release buffered/in-flight messages so they are not + # left locked at the broker until lock expiry (delaying redelivery, + # inflating delivery count). + self._amqp_transport.drain_receive_link_and_release_messages(self._handler) super(ServiceBusReceiver, self)._close_handler() @property diff --git a/sdk/servicebus/azure-servicebus/azure/servicebus/_transport/_base.py b/sdk/servicebus/azure-servicebus/azure/servicebus/_transport/_base.py index 5ee0409c4f34..d95a6cd89ab6 100644 --- a/sdk/servicebus/azure-servicebus/azure/servicebus/_transport/_base.py +++ b/sdk/servicebus/azure-servicebus/azure/servicebus/_transport/_base.py @@ -302,6 +302,14 @@ def reset_link_credit(handler, link_credit): :param int link_credit: The link credit. """ + @staticmethod + @abstractmethod + def drain_receive_link_and_release_messages(handler): + """ + Drain the receive link and release buffered/in-flight messages on close. + :param ~uamqp.ReceiveClient or ~pyamqp.ReceiveClient handler: The handler. + """ + @staticmethod @abstractmethod def settle_message_via_receiver_link( diff --git a/sdk/servicebus/azure-servicebus/azure/servicebus/_transport/_pyamqp_transport.py b/sdk/servicebus/azure-servicebus/azure/servicebus/_transport/_pyamqp_transport.py index ac5a3e1add7e..bbf084d6f3b7 100644 --- a/sdk/servicebus/azure-servicebus/azure/servicebus/_transport/_pyamqp_transport.py +++ b/sdk/servicebus/azure-servicebus/azure/servicebus/_transport/_pyamqp_transport.py @@ -157,6 +157,10 @@ def is_retryable(self, error): ERROR_CODE_TIMEOUT: OperationTimeoutError, } +# Safety cap (seconds) for draining a receive link on close; the drain loop exits +# early on quiescence, so a normal close returns well before this. +RECEIVE_LINK_DRAIN_TIMEOUT = 5 + class PyamqpTransport(AmqpTransport): # pylint: disable=too-many-public-methods """ @@ -799,6 +803,53 @@ def reset_link_credit(handler: "ReceiveClient", link_credit: int) -> None: """ handler._link.flow(link_credit=link_credit) # pylint: disable=protected-access + @staticmethod + def drain_receive_link_and_release_messages(handler: "ReceiveClient") -> None: + """ + Drain the receive link and release buffered/in-flight messages on close, so + they are not left locked at the broker until lock expiry. Intended for + non-session PEEK_LOCK receivers only (gated by the caller). + :param ReceiveClient handler: Client whose receive link to drain. + :rtype: None + """ + # pylint: disable=protected-access + link = handler._link + if link is None or link._is_closed: + return # cannot drain or settle through a closed/absent link + try: + if link.current_link_credit > 0: + # Outstanding credit: stop the broker and pump in-flight transfers + # into the buffer before releasing them below. + outstanding = link.current_link_credit + link.flow(link_credit=0, drain=True) + deadline = time.time() + RECEIVE_LINK_DRAIN_TIMEOUT + idle_cycles = 0 + while True: + remaining = deadline - time.time() + if remaining <= 0: + break + before = handler._received_messages.qsize() + # listen() directly, not do_work() (which re-issues credit at 0 and + # undoes the drain). batch=outstanding assembles multi-frame transfers + # in one cycle; cap each read by the remaining budget so close stays + # bounded. pyamqp drops the drain echo, so stop after 2 idle cycles. + wait = remaining if handler._socket_timeout is None else min(handler._socket_timeout, remaining) + handler._connection.listen(wait=wait, batch=outstanding) + if handler._received_messages.qsize() == before: + idle_cycles += 1 + if idle_cycles >= 2: + break + else: + idle_cycles = 0 + # Always release buffered deliveries (incl. the credit==0 prefetch case) so + # they are not left locked until lock expiry. + while not handler._received_messages.empty(): + frame, _ = handler._received_messages.get_nowait() + handler.settle_messages(frame[1], frame[2], "released") + handler._received_messages.task_done() + except Exception: # pylint: disable=broad-except + pass # a faulted/detached link must not block close + @staticmethod def settle_message_via_receiver_link( handler: ReceiveClient, diff --git a/sdk/servicebus/azure-servicebus/azure/servicebus/_transport/_uamqp_transport.py b/sdk/servicebus/azure-servicebus/azure/servicebus/_transport/_uamqp_transport.py index 1df6cbae992f..bc0c5a1cf90d 100644 --- a/sdk/servicebus/azure-servicebus/azure/servicebus/_transport/_uamqp_transport.py +++ b/sdk/servicebus/azure-servicebus/azure/servicebus/_transport/_uamqp_transport.py @@ -848,6 +848,15 @@ def reset_link_credit(handler: "ReceiveClient", link_credit: int) -> None: """ handler.message_handler.reset_link_credit(link_credit) + @staticmethod + def drain_receive_link_and_release_messages(handler: "ReceiveClient") -> None: + """ + No-op for uamqp: drain-on-close is only implemented for the pyamqp + transport (the default). uamqp is deprecated. + :param ~uamqp.ReceiveClient handler: The handler. + :rtype: None + """ + # Executes message settlement, implementation is in settle_message_via_receiver_link_impl # May be able to remove and just call methods in private method. @staticmethod diff --git a/sdk/servicebus/azure-servicebus/azure/servicebus/aio/_servicebus_receiver_async.py b/sdk/servicebus/azure-servicebus/azure/servicebus/aio/_servicebus_receiver_async.py index 795e6a5b1d04..f247fb7c154b 100644 --- a/sdk/servicebus/azure-servicebus/azure/servicebus/aio/_servicebus_receiver_async.py +++ b/sdk/servicebus/azure-servicebus/azure/servicebus/aio/_servicebus_receiver_async.py @@ -527,6 +527,15 @@ async def _renew_locks(self, *lock_tokens: str, timeout: Optional[float] = None) async def _close_handler(self): self._message_iter = None + if ( + self._handler + and self._receive_mode == ServiceBusReceiveMode.PEEK_LOCK + and not self._session + ): + # Drain the link and release buffered/in-flight messages so they are not + # left locked at the broker until lock expiry (delaying redelivery, + # inflating delivery count). + await self._amqp_transport.drain_receive_link_and_release_messages_async(self._handler) await super(ServiceBusReceiver, self)._close_handler() @property diff --git a/sdk/servicebus/azure-servicebus/azure/servicebus/aio/_transport/_base_async.py b/sdk/servicebus/azure-servicebus/azure/servicebus/aio/_transport/_base_async.py index 00697accbd42..3e8c8d049593 100644 --- a/sdk/servicebus/azure-servicebus/azure/servicebus/aio/_transport/_base_async.py +++ b/sdk/servicebus/azure-servicebus/azure/servicebus/aio/_transport/_base_async.py @@ -244,6 +244,16 @@ async def reset_link_credit_async(handler, link_credit): :rtype: None """ + @staticmethod + @abstractmethod + async def drain_receive_link_and_release_messages_async(handler): + """ + Drain the receive link and release buffered/in-flight messages on close. + :param ~uamqp.ReceiveClientAsync + or ~pyamqp.aio.ReceiveClientAsync handler: Client whose receive link to drain. + :rtype: None + """ + @staticmethod @abstractmethod async def settle_message_via_receiver_link_async( diff --git a/sdk/servicebus/azure-servicebus/azure/servicebus/aio/_transport/_pyamqp_transport_async.py b/sdk/servicebus/azure-servicebus/azure/servicebus/aio/_transport/_pyamqp_transport_async.py index 105eda4572b2..4296e3ed41fe 100644 --- a/sdk/servicebus/azure-servicebus/azure/servicebus/aio/_transport/_pyamqp_transport_async.py +++ b/sdk/servicebus/azure-servicebus/azure/servicebus/aio/_transport/_pyamqp_transport_async.py @@ -42,7 +42,7 @@ OPERATION_TIMEOUT, NEXT_AVAILABLE_SESSION, ) -from ..._transport._pyamqp_transport import PyamqpTransport +from ..._transport._pyamqp_transport import PyamqpTransport, RECEIVE_LINK_DRAIN_TIMEOUT from ...exceptions import ( OperationTimeoutError, ServiceBusConnectionError, @@ -301,6 +301,54 @@ async def reset_link_credit_async(handler: "ReceiveClientAsync", link_credit: in """ await handler._link.flow(link_credit=link_credit) # pylint: disable=protected-access + @staticmethod + async def drain_receive_link_and_release_messages_async(handler: "ReceiveClientAsync") -> None: + """ + Drain the receive link and release buffered/in-flight messages on close, so + they are not left locked at the broker until lock expiry. Intended for + non-session PEEK_LOCK receivers only (gated by the caller). + :param ReceiveClientAsync handler: Client whose receive link to drain. + :rtype: None + """ + # pylint: disable=protected-access + link = handler._link + if link is None or link._is_closed: + return # cannot drain or settle through a closed/absent link + try: + if link.current_link_credit > 0: + # Outstanding credit: stop the broker and pump in-flight transfers + # into the buffer before releasing them below. + outstanding = link.current_link_credit + await link.flow(link_credit=0, drain=True) + deadline = time.time() + RECEIVE_LINK_DRAIN_TIMEOUT + idle_cycles = 0 + while True: + remaining = deadline - time.time() + if remaining <= 0: + break + before = handler._received_messages.qsize() + # listen() directly, not do_work_async() (which re-issues credit at + # 0 and undoes the drain). batch=outstanding assembles multi-frame + # transfers in one cycle; cap each read by the remaining budget so + # close stays bounded. pyamqp drops the drain echo, so stop after 2 + # idle cycles. + wait = remaining if handler._socket_timeout is None else min(handler._socket_timeout, remaining) + await handler._connection.listen(wait=wait, batch=outstanding) + if handler._received_messages.qsize() == before: + idle_cycles += 1 + if idle_cycles >= 2: + break + else: + idle_cycles = 0 + # Always release buffered deliveries (incl. the credit==0 prefetch case) so + # they are not left locked until lock expiry. + while not handler._received_messages.empty(): + frame, _ = handler._received_messages.get_nowait() + await handler.settle_messages_async(frame[1], frame[2], "released") + handler._received_messages.task_done() + except Exception: # pylint: disable=broad-except + pass # a faulted/detached link must not block close + @staticmethod async def settle_message_via_receiver_link_async( handler: "ReceiveClientAsync", diff --git a/sdk/servicebus/azure-servicebus/azure/servicebus/aio/_transport/_uamqp_transport_async.py b/sdk/servicebus/azure-servicebus/azure/servicebus/aio/_transport/_uamqp_transport_async.py index f59a33f0f511..c7ac0038eec1 100644 --- a/sdk/servicebus/azure-servicebus/azure/servicebus/aio/_transport/_uamqp_transport_async.py +++ b/sdk/servicebus/azure-servicebus/azure/servicebus/aio/_transport/_uamqp_transport_async.py @@ -248,6 +248,15 @@ async def reset_link_credit_async(handler: "ReceiveClientAsync", link_credit: in """ await handler.message_handler.reset_link_credit_async(link_credit) + @staticmethod + async def drain_receive_link_and_release_messages_async(handler: "ReceiveClientAsync") -> None: + """ + No-op for uamqp: drain-on-close is only implemented for the pyamqp + transport (the default). uamqp is deprecated. + :param ReceiveClientAsync handler: The handler. + :rtype: None + """ + @staticmethod async def settle_message_via_receiver_link_async( handler: "ReceiveClientAsync", diff --git a/sdk/servicebus/azure-servicebus/tests/unittests/test_receiver_drain_on_close.py b/sdk/servicebus/azure-servicebus/tests/unittests/test_receiver_drain_on_close.py new file mode 100644 index 000000000000..650646502f37 --- /dev/null +++ b/sdk/servicebus/azure-servicebus/tests/unittests/test_receiver_drain_on_close.py @@ -0,0 +1,303 @@ +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. + +# pylint: disable=protected-access,unused-argument + +""" +Unit tests for draining and releasing buffered/in-flight messages when a Service +Bus receiver is closed (regression guard for azure-sdk-for-python#42917). + +When ``receive_messages`` grants AMQP link credit, the broker may transfer more +messages than the application consumes before the call returns. Those messages +end up either in the client-side buffer ``_received_messages`` or still in flight +on the wire (leftover credit). Prior to the fix, closing the receiver cleared the +buffer without settling it and never drained the link, leaving those messages +locked at the broker until lock expiry (delaying redelivery and inflating the +delivery count). + +The fix lives in the Service Bus transport layer (keeping the vendored ``_pyamqp`` +package untouched): + + * ``PyamqpTransport.drain_receive_link_and_release_messages`` (and the async + twin) drains the link -- ``flow(link_credit=0, drain=True)`` + pumping the + connection until quiescent, bounded by a timeout -- then releases every + buffered message with the ``released`` disposition. + * ``ServiceBusReceiver._close_handler`` calls it, gated to non-session + PEEK_LOCK receivers (parity with the .NET SDK). + +These tests exercise the transport method directly and the receiver gating, with +no network. +""" + +import queue +import time + +import pytest +from unittest.mock import AsyncMock, MagicMock, patch + +from azure.servicebus import ServiceBusReceiveMode +from azure.servicebus._servicebus_receiver import ServiceBusReceiver +from azure.servicebus._base_handler import BaseHandler +from azure.servicebus._transport._pyamqp_transport import PyamqpTransport +from azure.servicebus.aio._servicebus_receiver_async import ServiceBusReceiver as ServiceBusReceiverAsync +from azure.servicebus.aio._base_handler_async import BaseHandler as BaseHandlerAsync +from azure.servicebus.aio._transport._pyamqp_transport_async import PyamqpTransportAsync +from azure.servicebus._pyamqp.performatives import TransferFrame + + +def _buffer_with(*deliveries): + buf = queue.Queue() + for delivery_id, delivery_tag in deliveries: + frame = TransferFrame(handle=0, delivery_id=delivery_id, delivery_tag=delivery_tag) + buf.put((frame, MagicMock(name=f"message-{delivery_id}"))) + return buf + + +def _handler(*, is_closed=False, credit=5, buffer=None, connection=None, is_async=False): + h = MagicMock(name="handler") + h._link = MagicMock(name="link") + h._link._is_closed = is_closed + h._link.current_link_credit = credit + h._link.flow = AsyncMock(name="flow") if is_async else MagicMock(name="flow") + if connection is None: + connection = MagicMock(name="connection") + connection.listen = AsyncMock(name="listen") if is_async else MagicMock(name="listen") + h._connection = connection + h._received_messages = buffer if buffer is not None else queue.Queue() + h._socket_timeout = 0.2 + if is_async: + h.settle_messages_async = AsyncMock(name="settle_messages_async") + else: + h.settle_messages = MagicMock(name="settle_messages") + return h + + +# --------------------------------------------------------------------------- # +# Transport method: drain + release # +# --------------------------------------------------------------------------- # + +class TestPyamqpTransportDrain: + def test_drain_sent_and_buffer_released(self): + h = _handler(credit=5, buffer=_buffer_with((10, b"tag-10"), (11, b"tag-11"))) + PyamqpTransport.drain_receive_link_and_release_messages(h) + + h._link.flow.assert_called_once_with(link_credit=0, drain=True) + assert h._connection.listen.call_count >= 1 + released = [(a[0], a[1], a[2]) for a, _ in h.settle_messages.call_args_list] + assert released == [(10, b"tag-10", "released"), (11, b"tag-11", "released")] + assert h._received_messages.empty() + + def test_pulls_inflight_then_releases(self): + buffer = queue.Queue() + connection = MagicMock(name="connection") + calls = {"n": 0} + + def _listen(*args, **kwargs): + calls["n"] += 1 + if calls["n"] == 1: + frame = TransferFrame(handle=0, delivery_id=99, delivery_tag=b"tag-99") + buffer.put((frame, MagicMock())) + + connection.listen = MagicMock(side_effect=_listen) + h = _handler(credit=3, buffer=buffer, connection=connection) + + PyamqpTransport.drain_receive_link_and_release_messages(h) + + h._link.flow.assert_called_once_with(link_credit=0, drain=True) + h.settle_messages.assert_called_once_with(99, b"tag-99", "released") + + def test_no_credit_skips_drain_but_releases_buffer(self): + # A prefetch receiver can hit zero credit with deliveries still buffered: + # skip the drain flow, but the buffered messages must still be released. + h = _handler(credit=0, buffer=_buffer_with((10, b"tag-10"))) + PyamqpTransport.drain_receive_link_and_release_messages(h) + h._link.flow.assert_not_called() + h.settle_messages.assert_called_once_with(10, b"tag-10", "released") + assert h._received_messages.empty() + + def test_no_credit_empty_buffer_is_noop(self): + h = _handler(credit=0) + PyamqpTransport.drain_receive_link_and_release_messages(h) + h._link.flow.assert_not_called() + h.settle_messages.assert_not_called() + + def test_skipped_when_link_closed(self): + h = _handler(credit=5, is_closed=True, buffer=_buffer_with((10, b"tag-10"))) + PyamqpTransport.drain_receive_link_and_release_messages(h) + h._link.flow.assert_not_called() + h.settle_messages.assert_not_called() # cannot settle through a closed link + + def test_skipped_when_link_none(self): + h = _handler(credit=5) + h._link = None + PyamqpTransport.drain_receive_link_and_release_messages(h) # must not raise + h.settle_messages.assert_not_called() + + def test_failure_does_not_raise(self): + h = _handler(credit=5, buffer=_buffer_with((10, b"tag-10"))) + h._link.flow.side_effect = ValueError("link faulted") + PyamqpTransport.drain_receive_link_and_release_messages(h) # must not raise + # drain aborted before releasing; no exception escapes + h.settle_messages.assert_not_called() + + def test_bounded_by_deadline_even_with_blocking_listen(self): + # A large socket_timeout must not let a single blocking read overrun the cap: + # each listen() should wait at most the remaining drain budget. + buffer = queue.Queue() + connection = MagicMock(name="connection") + + def _listen(*args, **kwargs): + time.sleep(kwargs.get("wait", 0)) # block for the requested wait + frame = TransferFrame(handle=0, delivery_id=1, delivery_tag=b"t") + buffer.put((frame, MagicMock())) # never quiescent + + connection.listen = MagicMock(side_effect=_listen) + h = _handler(credit=3, buffer=buffer, connection=connection) + h._socket_timeout = 5 # much larger than the drain cap + + start = time.time() + with patch("azure.servicebus._transport._pyamqp_transport.RECEIVE_LINK_DRAIN_TIMEOUT", 0.1): + PyamqpTransport.drain_receive_link_and_release_messages(h) + elapsed = time.time() - start + + assert h._connection.listen.call_count >= 1 + assert elapsed < 2.0 # bounded by the cap, not the 5s socket timeout + + +class TestPyamqpTransportDrainAsync: + @pytest.mark.asyncio + async def test_drain_sent_and_buffer_released(self): + h = _handler(credit=5, buffer=_buffer_with((20, b"tag-20"), (21, b"tag-21")), is_async=True) + await PyamqpTransportAsync.drain_receive_link_and_release_messages_async(h) + + h._link.flow.assert_awaited_once_with(link_credit=0, drain=True) + assert h._connection.listen.await_count >= 1 + released = [(a[0], a[1], a[2]) for a, _ in h.settle_messages_async.call_args_list] + assert released == [(20, b"tag-20", "released"), (21, b"tag-21", "released")] + assert h._received_messages.empty() + + @pytest.mark.asyncio + async def test_pulls_inflight_then_releases(self): + buffer = queue.Queue() + connection = MagicMock(name="connection") + calls = {"n": 0} + + async def _listen(*args, **kwargs): + calls["n"] += 1 + if calls["n"] == 1: + frame = TransferFrame(handle=0, delivery_id=88, delivery_tag=b"tag-88") + buffer.put((frame, MagicMock())) + + connection.listen = AsyncMock(side_effect=_listen) + h = _handler(credit=3, buffer=buffer, connection=connection, is_async=True) + + await PyamqpTransportAsync.drain_receive_link_and_release_messages_async(h) + + h._link.flow.assert_awaited_once_with(link_credit=0, drain=True) + h.settle_messages_async.assert_awaited_once_with(88, b"tag-88", "released") + + @pytest.mark.asyncio + async def test_no_credit_skips_drain_but_releases_buffer(self): + h = _handler(credit=0, buffer=_buffer_with((20, b"tag-20")), is_async=True) + await PyamqpTransportAsync.drain_receive_link_and_release_messages_async(h) + h._link.flow.assert_not_called() + h.settle_messages_async.assert_awaited_once_with(20, b"tag-20", "released") + assert h._received_messages.empty() + + @pytest.mark.asyncio + async def test_failure_does_not_raise(self): + h = _handler(credit=5, buffer=_buffer_with((20, b"tag-20")), is_async=True) + h._link.flow.side_effect = ValueError("link faulted") + await PyamqpTransportAsync.drain_receive_link_and_release_messages_async(h) # must not raise + h.settle_messages_async.assert_not_called() + + @pytest.mark.asyncio + async def test_bounded_by_deadline_even_with_blocking_listen(self): + import asyncio # pylint: disable=do-not-import-asyncio + + buffer = queue.Queue() + connection = MagicMock(name="connection") + + async def _listen(*args, **kwargs): + await asyncio.sleep(kwargs.get("wait", 0)) # block for the requested wait + frame = TransferFrame(handle=0, delivery_id=1, delivery_tag=b"t") + buffer.put((frame, MagicMock())) # never quiescent + + connection.listen = AsyncMock(side_effect=_listen) + h = _handler(credit=3, buffer=buffer, connection=connection, is_async=True) + h._socket_timeout = 5 # much larger than the drain cap + + start = time.time() + with patch("azure.servicebus.aio._transport._pyamqp_transport_async.RECEIVE_LINK_DRAIN_TIMEOUT", 0.1): + await PyamqpTransportAsync.drain_receive_link_and_release_messages_async(h) + elapsed = time.time() - start + + assert h._connection.listen.await_count >= 1 + assert elapsed < 2.0 # bounded by the cap, not the 5s socket timeout + + +# --------------------------------------------------------------------------- # +# Receiver gating: only non-session PEEK_LOCK receivers drain on close # +# --------------------------------------------------------------------------- # + +def _sync_receiver(mode, session): + r = ServiceBusReceiver.__new__(ServiceBusReceiver) + r._handler = MagicMock(name="handler") + r._message_iter = None + r._receive_mode = mode + r._session = session + r._amqp_transport = MagicMock(name="transport") + r._amqp_transport.drain_receive_link_and_release_messages = MagicMock(name="drain") + return r + + +def _async_receiver(mode, session): + r = ServiceBusReceiverAsync.__new__(ServiceBusReceiverAsync) + r._handler = MagicMock(name="handler") + r._message_iter = None + r._receive_mode = mode + r._session = session + r._amqp_transport = MagicMock(name="transport") + r._amqp_transport.drain_receive_link_and_release_messages_async = AsyncMock(name="drain") + return r + + +class TestReceiverCloseGating: + def test_peek_lock_non_session_drains(self): + r = _sync_receiver(ServiceBusReceiveMode.PEEK_LOCK, session=None) + with patch.object(BaseHandler, "_close_handler"): + r._close_handler() + r._amqp_transport.drain_receive_link_and_release_messages.assert_called_once_with(r._handler) + + def test_receive_and_delete_does_not_drain(self): + r = _sync_receiver(ServiceBusReceiveMode.RECEIVE_AND_DELETE, session=None) + with patch.object(BaseHandler, "_close_handler"): + r._close_handler() + r._amqp_transport.drain_receive_link_and_release_messages.assert_not_called() + + def test_session_receiver_does_not_drain(self): + r = _sync_receiver(ServiceBusReceiveMode.PEEK_LOCK, session=MagicMock(name="session")) + with patch.object(BaseHandler, "_close_handler"): + r._close_handler() + r._amqp_transport.drain_receive_link_and_release_messages.assert_not_called() + + @pytest.mark.asyncio + async def test_peek_lock_non_session_drains_async(self): + r = _async_receiver(ServiceBusReceiveMode.PEEK_LOCK, session=None) + with patch.object(BaseHandlerAsync, "_close_handler", new=AsyncMock()): + await r._close_handler() + r._amqp_transport.drain_receive_link_and_release_messages_async.assert_awaited_once_with(r._handler) + + @pytest.mark.asyncio + async def test_receive_and_delete_does_not_drain_async(self): + r = _async_receiver(ServiceBusReceiveMode.RECEIVE_AND_DELETE, session=None) + with patch.object(BaseHandlerAsync, "_close_handler", new=AsyncMock()): + await r._close_handler() + r._amqp_transport.drain_receive_link_and_release_messages_async.assert_not_called() + + @pytest.mark.asyncio + async def test_session_receiver_does_not_drain_async(self): + r = _async_receiver(ServiceBusReceiveMode.PEEK_LOCK, session=MagicMock(name="session")) + with patch.object(BaseHandlerAsync, "_close_handler", new=AsyncMock()): + await r._close_handler() + r._amqp_transport.drain_receive_link_and_release_messages_async.assert_not_called()