From eeb9c4483ea5593af181e0e8f9a00f9c8aabf73e Mon Sep 17 00:00:00 2001 From: Erik Shestopal Date: Thu, 4 Jun 2026 14:55:09 -0400 Subject: [PATCH 1/2] Add regressions for returned message content handling --- tests/channel/test_publish.py | 480 +++++++++++++++++++++++++++++++++- 1 file changed, 479 insertions(+), 1 deletion(-) diff --git a/tests/channel/test_publish.py b/tests/channel/test_publish.py index f4927fa..23a7ec6 100644 --- a/tests/channel/test_publish.py +++ b/tests/channel/test_publish.py @@ -1,14 +1,139 @@ +import typing + +import anyio import pytest +from serena.channel import Channel +from serena.connection import AMQPConnection from serena.enums import ReplyCode from serena.exc import MessageReturnedError +from serena.frame import BodyFrame, Frame, HeartbeatFrame from serena.message import AMQPMessage -from serena.payloads.header import BasicHeader +from serena.payloads.header import BasicHeader, ContentHeaderFrame, ContentHeaderPayload +from serena.payloads.method import ( + BasicDeliverPayload, + BasicGetEmptyPayload, + BasicGetOkPayload, + BasicReturnPayload, + MethodFrame, +) from tests import _open_connection pytestmark = pytest.mark.anyio +def _return_frames( + routing_key: str, + body: bytes = b"x", + *, + channel_id: int = 1, + header: BasicHeader | None = None, +) -> tuple[MethodFrame[BasicReturnPayload], ContentHeaderFrame, BodyFrame]: + return_frame: MethodFrame[BasicReturnPayload] = MethodFrame( + channel_id=channel_id, + payload=BasicReturnPayload( + reply_code=ReplyCode.no_route, + reply_text="NO_ROUTE", + exchange="", + routing_key=routing_key, + ), + ) + header_frame = ContentHeaderFrame( + channel_id=channel_id, + payload=ContentHeaderPayload( + class_id=return_frame.payload.klass, + full_size=len(body), + flags=0, + payload=header or BasicHeader(), + ), + ) + body_frame = BodyFrame(channel_id=channel_id, data=body) + return return_frame, header_frame, body_frame + + +def _delivery_frames( + routing_key: str, + body: bytes = b"x", + *, + channel_id: int = 2, + header: BasicHeader | None = None, +) -> tuple[MethodFrame[BasicDeliverPayload], ContentHeaderFrame, BodyFrame]: + deliver_frame: MethodFrame[BasicDeliverPayload] = MethodFrame( + channel_id=channel_id, + payload=BasicDeliverPayload( + consumer_tag="consumer", + delivery_tag=1, + redelivered=False, + exchange_name="", + routing_key=routing_key, + ), + ) + header_frame = ContentHeaderFrame( + channel_id=channel_id, + payload=ContentHeaderPayload( + class_id=deliver_frame.payload.klass, + full_size=len(body), + flags=0, + payload=header or BasicHeader(), + ), + ) + body_frame = BodyFrame(channel_id=channel_id, data=body) + return deliver_frame, header_frame, body_frame + + +def _get_ok_frames( + routing_key: str, + body: bytes = b"x", + *, + channel_id: int = 1, + header: BasicHeader | None = None, +) -> tuple[MethodFrame[BasicGetOkPayload], ContentHeaderFrame, BodyFrame]: + get_ok_frame: MethodFrame[BasicGetOkPayload] = MethodFrame( + channel_id=channel_id, + payload=BasicGetOkPayload( + delivery_tag=1, + redelivered=False, + exchange_name="", + routing_key=routing_key, + message_count=0, + ), + ) + header_frame = ContentHeaderFrame( + channel_id=channel_id, + payload=ContentHeaderPayload( + class_id=get_ok_frame.payload.klass, + full_size=len(body), + flags=0, + payload=header or BasicHeader(), + ), + ) + body_frame = BodyFrame(channel_id=channel_id, data=body) + return get_ok_frame, header_frame, body_frame + + +async def _read_delivery_message(frames: list[Frame]) -> AMQPMessage | None: + channel = Channel(1, typing.cast(typing.Any, object()), stream_buffer_size=8) + received = None + message_received = anyio.Event() + + async with anyio.create_task_group() as task_group: + async def receive_message() -> None: + nonlocal received + received = await channel._receive_delivery_message() + message_received.set() + + task_group.start_soon(receive_message) + for frame in frames: + await channel._enqueue_delivery(frame) + + with anyio.fail_after(1): + await message_received.wait() + + task_group.cancel_scope.cancel() + + return received + + async def test_basic_publish(): """ Tests publishing a message to a queue, and getting the message back. @@ -110,3 +235,356 @@ async def test_return(): assert channel.open assert e.value.reply_code == ReplyCode.no_route + + +async def test_return_preserves_returned_body_and_header(): + async with _open_connection() as conn, conn.open_channel() as channel: + header = BasicHeader(message_id="returned-message") + + with pytest.raises(MessageReturnedError) as e: + await channel.basic_publish( + "", + routing_key="non-existent-queue", + body=b"returned body", + header=header, + ) + + assert e.value.reply_code == ReplyCode.no_route + assert e.value.body == b"returned body" + assert e.value.header == header + + +async def test_basic_get_empty_frame_returns_none(): + message = await _read_delivery_message( + [MethodFrame(channel_id=1, payload=BasicGetEmptyPayload(reserved_1=""))] + ) + + assert message is None + + +async def test_basic_get_ok_reassembles_empty_body(): + method_frame, header_frame, _body_frame = _get_ok_frames("queue", b"") + + message = await _read_delivery_message([method_frame, header_frame]) + + assert message is not None + assert message.body == b"" + assert message.envelope.routing_key == "queue" + assert message.envelope.delivery_tag == 1 + + +async def test_basic_get_ok_reassembles_body_across_multiple_frames(): + header = BasicHeader(message_id="get-message") + method_frame, header_frame, _body_frame = _get_ok_frames("queue", b"hello world", header=header) + + message = await _read_delivery_message( + [ + method_frame, + header_frame, + BodyFrame(channel_id=1, data=b"hello "), + BodyFrame(channel_id=1, data=b"world"), + ] + ) + + assert message is not None + assert message.body == b"hello world" + assert message.header == header + assert message.envelope.routing_key == "queue" + + +async def test_basic_deliver_reassembles_empty_body_and_envelope(): + method_frame, header_frame, _body_frame = _delivery_frames("queue", b"", channel_id=1) + + message = await _read_delivery_message([method_frame, header_frame]) + + assert message is not None + assert message.body == b"" + assert message.envelope.consumer_tag == "consumer" + assert message.envelope.routing_key == "queue" + assert message.envelope.delivery_tag == 1 + + +async def test_basic_deliver_reassembles_body_and_headers_across_multiple_frames(): + header = BasicHeader(message_id="delivered-message") + method_frame, header_frame, _body_frame = _delivery_frames( + "queue", b"hello world", channel_id=1, header=header + ) + + message = await _read_delivery_message( + [ + method_frame, + header_frame, + BodyFrame(channel_id=1, data=b"hello "), + BodyFrame(channel_id=1, data=b"world"), + ] + ) + + assert message is not None + assert message.body == b"hello world" + assert message.header == header + assert message.envelope.consumer_tag == "consumer" + assert message.envelope.routing_key == "queue" + + +async def test_empty_return_content_frames_do_not_buffer_as_deliveries(): + connection = AMQPConnection(typing.cast(typing.Any, object()), channel_buffer_size=8) + connection._actual_heartbeat_interval = 10 + channel = Channel(1, connection, stream_buffer_size=8) + connection._channel_channels[1] = channel + + send_frames, receive_frames = anyio.create_memory_object_stream[Frame](10) + heartbeat_read = anyio.Event() + + async def read_single_frame() -> Frame: + frame = await receive_frames.receive() + if isinstance(frame, HeartbeatFrame): + heartbeat_read.set() + return frame + + connection._read_single_frame = read_single_frame # type: ignore[method-assign] + + received_return: MethodFrame[BasicReturnPayload] | None = None + receiver_ready = anyio.Event() + return_received = anyio.Event() + + async def receive_return() -> None: + nonlocal received_return + receiver_ready.set() + frame = await channel._receive_frame() + assert isinstance(frame.payload, BasicReturnPayload) + received_return = typing.cast(MethodFrame[BasicReturnPayload], frame) + return_received.set() + + return_frame, header_frame, _body_frame = _return_frames("missing-queue", b"") + + async with anyio.create_task_group() as task_group: + task_group.start_soon(connection._listen_for_messages) + task_group.start_soon(receive_return) + + await receiver_ready.wait() + await send_frames.send(return_frame) + await send_frames.send(header_frame) + await send_frames.send(HeartbeatFrame(channel_id=0)) + + with anyio.fail_after(1): + await return_received.wait() + await heartbeat_read.wait() + + assert received_return == return_frame + assert channel.current_buffer_size == 0 + + task_group.cancel_scope.cancel() + + +async def test_multi_frame_return_content_frames_do_not_buffer_as_deliveries(): + connection = AMQPConnection(typing.cast(typing.Any, object()), channel_buffer_size=8) + connection._actual_heartbeat_interval = 10 + channel = Channel(1, connection, stream_buffer_size=8) + connection._channel_channels[1] = channel + + send_frames, receive_frames = anyio.create_memory_object_stream[Frame](10) + heartbeat_read = anyio.Event() + + async def read_single_frame() -> Frame: + frame = await receive_frames.receive() + if isinstance(frame, HeartbeatFrame): + heartbeat_read.set() + return frame + + connection._read_single_frame = read_single_frame # type: ignore[method-assign] + + received_return: MethodFrame[BasicReturnPayload] | None = None + receiver_ready = anyio.Event() + return_received = anyio.Event() + + async def receive_return() -> None: + nonlocal received_return + receiver_ready.set() + frame = await channel._receive_frame() + assert isinstance(frame.payload, BasicReturnPayload) + received_return = typing.cast(MethodFrame[BasicReturnPayload], frame) + return_received.set() + + return_frame, header_frame, _body_frame = _return_frames("missing-queue", b"hello world") + + async with anyio.create_task_group() as task_group: + task_group.start_soon(connection._listen_for_messages) + task_group.start_soon(receive_return) + + await receiver_ready.wait() + await send_frames.send(return_frame) + await send_frames.send(header_frame) + await send_frames.send(BodyFrame(channel_id=1, data=b"hello ")) + await send_frames.send(BodyFrame(channel_id=1, data=b"world")) + await send_frames.send(HeartbeatFrame(channel_id=0)) + + with anyio.fail_after(1): + await return_received.wait() + await heartbeat_read.wait() + + assert received_return == return_frame + assert channel.current_buffer_size == 0 + + task_group.cancel_scope.cancel() + + +async def test_return_content_frames_do_not_buffer_as_deliveries(): + connection = AMQPConnection(typing.cast(typing.Any, object()), channel_buffer_size=8) + connection._actual_heartbeat_interval = 10 + channel = Channel(1, connection, stream_buffer_size=8) + connection._channel_channels[1] = channel + + send_frames, receive_frames = anyio.create_memory_object_stream[Frame](10) + heartbeat_read = anyio.Event() + second_heartbeat_read = anyio.Event() + heartbeat_count = 0 + + async def read_single_frame() -> Frame: + nonlocal heartbeat_count + frame = await receive_frames.receive() + if isinstance(frame, HeartbeatFrame): + heartbeat_count += 1 + heartbeat_read.set() + if heartbeat_count == 2: + second_heartbeat_read.set() + return frame + + connection._read_single_frame = read_single_frame # type: ignore[method-assign] + + received_return: MethodFrame[BasicReturnPayload] | None = None + receiver_ready = anyio.Event() + return_received = anyio.Event() + + async def receive_return() -> None: + nonlocal received_return + receiver_ready.set() + frame = await channel._receive_frame() + assert isinstance(frame.payload, BasicReturnPayload) + received_return = typing.cast(MethodFrame[BasicReturnPayload], frame) + return_received.set() + + return_frame, header_frame, body_frame = _return_frames("missing-queue") + + async with anyio.create_task_group() as task_group: + task_group.start_soon(connection._listen_for_messages) + task_group.start_soon(receive_return) + + await receiver_ready.wait() + await send_frames.send(return_frame) + await send_frames.send(HeartbeatFrame(channel_id=0)) + await send_frames.send(header_frame) + await send_frames.send(body_frame) + await send_frames.send(HeartbeatFrame(channel_id=0)) + + with anyio.fail_after(1): + await return_received.wait() + await heartbeat_read.wait() + await second_heartbeat_read.wait() + + assert received_return == return_frame + assert channel.current_buffer_size == 0 + + task_group.cancel_scope.cancel() + + +async def test_many_returned_publishes_do_not_block_connection_reader(): + connection = AMQPConnection(typing.cast(typing.Any, object()), channel_buffer_size=8) + connection._actual_heartbeat_interval = 10 + channel = Channel(1, connection, stream_buffer_size=8) + connection._channel_channels[1] = channel + + send_frames, receive_frames = anyio.create_memory_object_stream[Frame](100) + heartbeat_read = anyio.Event() + + async def read_single_frame() -> Frame: + frame = await receive_frames.receive() + if isinstance(frame, HeartbeatFrame): + heartbeat_read.set() + return frame + + connection._read_single_frame = read_single_frame # type: ignore[method-assign] + + async with anyio.create_task_group() as task_group: + task_group.start_soon(connection._listen_for_messages) + + for index in range(channel.max_buffer_size): + for frame in _return_frames(f"missing-queue-{index}"): + await send_frames.send(frame) + await send_frames.send(HeartbeatFrame(channel_id=0)) + + with anyio.fail_after(1): + await heartbeat_read.wait() + + assert channel.current_buffer_size == 0 + + task_group.cancel_scope.cancel() + + +async def test_returned_message_does_not_interfere_with_delivery_on_another_channel(): + connection = AMQPConnection(typing.cast(typing.Any, object()), channel_buffer_size=8) + connection._actual_heartbeat_interval = 10 + return_channel = Channel(1, connection, stream_buffer_size=8) + delivery_channel = Channel(2, connection, stream_buffer_size=8) + connection._channel_channels[1] = return_channel + connection._channel_channels[2] = delivery_channel + + send_frames, receive_frames = anyio.create_memory_object_stream[Frame](20) + heartbeat_read = anyio.Event() + + async def read_single_frame() -> Frame: + frame = await receive_frames.receive() + if isinstance(frame, HeartbeatFrame): + heartbeat_read.set() + return frame + + connection._read_single_frame = read_single_frame # type: ignore[method-assign] + + received_return: MethodFrame[BasicReturnPayload] | None = None + return_receiver_ready = anyio.Event() + return_received = anyio.Event() + + async def receive_return() -> None: + nonlocal received_return + return_receiver_ready.set() + frame = await return_channel._receive_frame() + assert isinstance(frame.payload, BasicReturnPayload) + received_return = typing.cast(MethodFrame[BasicReturnPayload], frame) + return_received.set() + + received_delivery: AMQPMessage | None = None + delivery_received = anyio.Event() + + async def receive_delivery() -> None: + nonlocal received_delivery + received_delivery = await delivery_channel._receive_delivery_message() + delivery_received.set() + + return_frame, return_header, return_body = _return_frames("missing-queue", b"returned") + deliver_frame, deliver_header, deliver_body = _delivery_frames("queue", b"delivered") + + async with anyio.create_task_group() as task_group: + task_group.start_soon(connection._listen_for_messages) + task_group.start_soon(receive_return) + task_group.start_soon(receive_delivery) + + await return_receiver_ready.wait() + await send_frames.send(return_frame) + await send_frames.send(deliver_frame) + await send_frames.send(deliver_header) + await send_frames.send(return_header) + await send_frames.send(deliver_body) + await send_frames.send(return_body) + await send_frames.send(HeartbeatFrame(channel_id=0)) + + with anyio.fail_after(1): + await return_received.wait() + await delivery_received.wait() + await heartbeat_read.wait() + + assert received_return == return_frame + assert received_delivery is not None + assert received_delivery.body == b"delivered" + assert return_channel.current_buffer_size == 0 + assert delivery_channel.current_buffer_size == 0 + + task_group.cancel_scope.cancel() From 1cac8f8749e83ffcfe4cf72e37c667c2e25a9e84 Mon Sep 17 00:00:00 2001 From: Erik Shestopal Date: Thu, 4 Jun 2026 15:04:50 -0400 Subject: [PATCH 2/2] Assemble returned message content --- src/serena/channel.py | 115 ++++++++++++++++++++++++++------------- src/serena/connection.py | 29 ++++++++-- src/serena/exc.py | 18 +++++- 3 files changed, 120 insertions(+), 42 deletions(-) diff --git a/src/serena/channel.py b/src/serena/channel.py index 6c0180f..05a7a18 100644 --- a/src/serena/channel.py +++ b/src/serena/channel.py @@ -12,8 +12,9 @@ ) import anyio +import attr import outcome -from anyio import CancelScope, ClosedResourceError, EndOfStream, Event, Lock +from anyio import CancelScope, ClosedResourceError, EndOfStream, Event, Lock, WouldBlock from anyio.abc import TaskStatus from anyio.lowlevel import checkpoint from outcome import Error, Value @@ -29,7 +30,7 @@ from serena.frame import BodyFrame, Frame from serena.message import AMQPEnvelope, AMQPMessage from serena.mixin import ChannelLike -from serena.payloads.header import BasicHeader, ContentHeaderFrame, ContentHeaderPayload +from serena.payloads.header import BasicHeader, ContentHeaderFrame from serena.payloads.method import ( BasicAckPayload, BasicCancelPayload, @@ -76,6 +77,47 @@ PayloadType = TypeVar("PayloadType", bound=MethodPayload) +@attr.s(slots=True, frozen=True) +class _ContentMessage: + method: MethodFrame[MethodPayload] = attr.ib() + header: BasicHeader = attr.ib() + body: bytes = attr.ib() + + +class _ContentAssembly: + def __init__(self, method: MethodFrame[MethodPayload]) -> None: + self.method = method + self._header: ContentHeaderFrame | None = None + self._body = bytearray() + + def feed(self, frame: ContentHeaderFrame | BodyFrame) -> _ContentMessage | None: + if self._header is None: + if not isinstance(frame, ContentHeaderFrame): + raise AMQPStateError(f"Expected a header frame, got {frame} instead") + + if frame.payload.class_id != self.method.payload.klass: + raise AMQPStateError( + f"Class mismatch ({frame.payload.class_id} != {self.method.payload.klass})" + ) + + self._header = frame + else: + if not isinstance(frame, BodyFrame): + raise AMQPStateError(f"Expected a body frame, got {frame} instead") + + self._body += frame.data + + header = self._header + if len(self._body) >= header.payload.full_size: + return _ContentMessage( + method=self.method, + header=header.payload.payload, + body=bytes(self._body), + ) + + return None + + class Channel(ChannelLike): """ A wrapper around an AMQP channel. @@ -106,6 +148,7 @@ def __init__(self, channel_id: int, connection: AMQPConnection, stream_buffer_si self._delivery_send, self._delivery_receive = anyio.create_memory_object_stream[Frame]( max_buffer_size=stream_buffer_size ) + self._returned_content: dict[int, _ContentMessage] = {} self._close_info: ChannelClosePayload | None = None self._lock = Lock() @@ -194,30 +237,30 @@ async def _enqueue_delivery(self, frame: Frame) -> None: await self._delivery_send.send(frame) + def _enqueue_return(self, message: _ContentMessage) -> None: + """ + Enqueues a complete Basic.Return message. + """ + + self._returned_content[id(message.method)] = message + try: + self._enqueue_regular(message.method) + except WouldBlock: + self._returned_content.pop(id(message.method), None) + raise + + def _pop_return_content(self, frame: MethodFrame[MethodPayload]) -> _ContentMessage | None: + return self._returned_content.pop(id(frame), None) + async def _receive_delivery_message(self) -> AMQPMessage | None: """ Receives a single delivery message. This will reassemble a full message into its constituent frames. """ - method: MethodFrame[MethodPayload] | None = None - headers = None - body = b"" + assembly: _ContentAssembly | None = None while True: - # check for successful reassembly first - if headers is not None and len(body) >= headers.payload.full_size: - assert method is not None, "reached reassembly with invalid method" - - await checkpoint() - # hehe payload.payload - return AMQPMessage( - channel=self, - envelope=AMQPEnvelope.of(method.payload), - header=headers.payload.payload, - body=body, - ) - try: next_frame = await self._delivery_receive.receive() except EndOfStream: @@ -226,14 +269,14 @@ async def _receive_delivery_message(self) -> AMQPMessage | None: raise UnexpectedCloseError.of(self._close_info) from None - if method is None: + if assembly is None: if not isinstance(next_frame, MethodFrame): raise AMQPStateError(f"Expected a method frame, got {next_frame} instead") method_frame: MethodFrame[MethodPayload] = next_frame if isinstance(method_frame.payload, BasicGetOkPayload | BasicDeliverPayload): - method = method_frame + assembly = _ContentAssembly(method_frame) elif isinstance(method_frame.payload, BasicGetEmptyPayload): return None @@ -244,24 +287,19 @@ async def _receive_delivery_message(self) -> AMQPMessage | None: f"got {method_frame.payload} instead" ) - elif headers is None: - if not isinstance(next_frame, ContentHeaderFrame): - raise AMQPStateError(f"Expected a header frame, got {next_frame} instead") - - # explicit type hint as pycharm incorrectly infers based on the previous if check - payload: ContentHeaderPayload = next_frame.payload - if payload.class_id != method.payload.klass: - raise AMQPStateError( - f"Class mismatch ({payload.class_id} != {method.payload.klass})" - ) - - headers = next_frame - else: - if not isinstance(next_frame, BodyFrame): - raise AMQPStateError(f"Expected a body frame, got {next_frame} instead") - - body += next_frame.data + if not isinstance(next_frame, ContentHeaderFrame | BodyFrame): + raise AMQPStateError(f"Expected a content frame, got {next_frame} instead") + + message = assembly.feed(next_frame) + if message is not None: + await checkpoint() + return AMQPMessage( + channel=self, + envelope=AMQPEnvelope.of(message.method.payload), + header=message.header, + body=message.body, + ) async def _receive_frame(self) -> MethodFrame[MethodPayload]: """ @@ -814,11 +852,14 @@ async def basic_publish( logger.trace(f"C#{self.id}: Server ACKed published message") elif isinstance(payload, BasicReturnPayload): + returned_content = self._pop_return_content(response) raise MessageReturnedError( exchange=payload.exchange, routing_key=payload.routing_key, reply_code=payload.reply_code, reply_text=payload.reply_text, + header=returned_content.header if returned_content is not None else None, + body=returned_content.body if returned_content is not None else b"", ) elif isinstance(payload, BasicNackPayload): diff --git a/src/serena/connection.py b/src/serena/connection.py index 7cfdfe6..5db7f42 100644 --- a/src/serena/connection.py +++ b/src/serena/connection.py @@ -20,7 +20,7 @@ from anyio.abc import ByteStream, TaskGroup from anyio.lowlevel import checkpoint -from serena.channel import Channel +from serena.channel import Channel, _ContentAssembly from serena.enums import ClassID, ReplyCode from serena.exc import ( AMQPError, @@ -29,15 +29,16 @@ InvalidProtocolError, UnexpectedCloseError, ) -from serena.frame import Frame, FrameType +from serena.frame import BodyFrame, Frame, FrameType from serena.frameparser import NEED_DATA, FrameParser -from serena.payloads.header import BasicHeader +from serena.payloads.header import BasicHeader, ContentHeaderFrame from serena.payloads.method import ( BasicDeliverPayload, BasicGetEmptyPayload, BasicGetOkPayload, BasicQOSOkPayload, BasicQOSPayload, + BasicReturnPayload, ChannelCloseOkPayload, ChannelClosePayload, ChannelOpenPayload, @@ -159,6 +160,7 @@ def __init__( # mapping of channel id -> Channel self._channel_channels: dict[int, Channel] = {} + self._return_assemblies: dict[int, _ContentAssembly] = {} self._write_lock = Lock() @@ -492,6 +494,7 @@ def _remove_channel(self, channel_id: int, payload: ChannelClosePayload | None) """ self._channels[channel_id] = False + self._return_assemblies.pop(channel_id, None) chan = self._channel_channels.pop(channel_id) chan._close(payload) @@ -603,7 +606,10 @@ async def _listen_for_messages(self) -> None: continue - if isinstance( + if isinstance(payload, BasicReturnPayload): + self._return_assemblies[channel] = _ContentAssembly(method_frame) + + elif isinstance( payload, BasicDeliverPayload | BasicGetOkPayload | BasicGetEmptyPayload ): # requires special logic @@ -623,6 +629,21 @@ async def _listen_for_messages(self) -> None: channel = frame.channel_id assert channel != 0, "header frame cannot happen on control channel" channel_object = self._channel_channels[channel] + + if assembly := self._return_assemblies.get(channel): + assert isinstance(frame, ContentHeaderFrame | BodyFrame) + content = assembly.feed(frame) + if content is not None: + self._return_assemblies.pop(channel, None) + try: + channel_object._enqueue_return(content) + except WouldBlock: + logger.warning( + f"Channel #{channel} was not listening for returned message, " + "dropping" + ) + continue + await self._enqueue_frame(channel_object, frame) async def _listen_wrapper(self) -> None: diff --git a/src/serena/exc.py b/src/serena/exc.py index 350a76a..06db852 100644 --- a/src/serena/exc.py +++ b/src/serena/exc.py @@ -6,6 +6,7 @@ from typing_extensions import override from serena.enums import ReplyCode +from serena.payloads.header import BasicHeader from serena.payloads.method import ( ChannelClosePayload, ConnectionClosePayload, @@ -52,7 +53,16 @@ class MessageReturnedError(AMQPError): Thrown when a message is returned unexpectedly. """ - def __init__(self, reply_code: ReplyCode, reply_text: str, exchange: str, routing_key: str): + def __init__( + self, + reply_code: ReplyCode, + reply_text: str, + exchange: str, + routing_key: str, + *, + header: BasicHeader | None = None, + body: bytes = b"", + ): #: The server-provided error code. self.reply_code = reply_code #: The server-provided error text. @@ -64,6 +74,12 @@ def __init__(self, reply_code: ReplyCode, reply_text: str, exchange: str, routin #: The routing key the message was using. self.routing_key = routing_key + #: The returned message header. + self.header = header or BasicHeader() + + #: The returned message body. + self.body = body + @override def __str__(self) -> str: return (