async def _on_get(
self, frame: Union[spec.Basic.GetOk, spec.Basic.GetEmpty]) -> None:
message = None
if isinstance(frame, spec.Basic.GetOk):
header: ContentHeader = await self.__get_content_header()
message = await self._read_content(frame, header)
if isinstance(frame, spec.Basic.GetEmpty):
message = DeliveredMessage(
delivery=frame,
header=ContentHeader(),
body=b"",
channel=self,
)
getter = getattr(self, "getter", None)
if getter is None:
raise RuntimeError("Getter is None")
if getter.done():
log.error("Got message but no active getter")
return
getter.set_result((frame, message))
return
def invoke_mock_command(self, command_str, command_id=0):
if issubclass(actor.__class__, (clu.LegacyActor, clu.JSONActor)):
if isinstance(command_str, str):
command_str = command_str.encode("utf-8")
full_command = f" {command_id} ".encode("utf-8") + command_str
return self.new_command(actor.transports[user_id], full_command)
elif issubclass(actor.__class__, clu.AMQPActor):
command_id = str(command_id)
headers = {
"command_id": command_id,
"commander_id": "mock_test_client"
}
header = ContentHeader(
properties=pamqp.specification.Basic.Properties(
content_type="text/json",
headers=headers,
))
message_body = {"command_string": command_str}
message = aio_pika.IncomingMessage(
DeliveredMessage(
pamqp.specification.Basic.Deliver(),
header,
json.dumps(message_body).encode(),
None,
), )
return self.new_command(message, ack=False)
def test_channel_build_message_headers(self):
channel = Channel(0, None, 360)
deliver = specification.Basic.Deliver()
header = ContentHeader(body_size=10)
channel._inbound = [deliver, header]
result = channel._build_message_headers()
self.assertIsInstance(result[0], specification.Basic.Deliver)
self.assertIsInstance(result[1], ContentHeader)
self.assertEqual(result[1].body_size, 10)
def test_channel_build_out_of_order_message(self):
channel = Channel(0, None, 360)
message = b'Hello World!'
message_len = len(message)
deliver = specification.Basic.Deliver()
header = ContentHeader(body_size=message_len)
body = ContentBody(value=message)
channel._inbound = [deliver, deliver, header, body]
result = channel._build_message()
self.assertEqual(result, None)
def test_channel_build_inbound_messages(self):
channel = Channel(0, FakeConnection(), 360)
channel.set_state(Channel.OPEN)
message = b'Hello World!'
message_len = len(message)
deliver = specification.Basic.Deliver()
header = ContentHeader(body_size=message_len)
body = ContentBody(value=message)
channel._inbound = [deliver, header, body]
for message in channel.build_inbound_messages(break_on_empty=True):
self.assertIsInstance(message, Message)
def test_channel_build_out_of_order_message_deliver(self):
channel = Channel(0, None, 360)
message = b'Hello World!'
message_len = len(message)
deliver = specification.Basic.Deliver()
header = ContentHeader(body_size=message_len)
channel._inbound = [deliver, deliver, header]
result = channel._build_message()
self.assertEqual(result, None)
self.assertIn("Received an out-of-order frame:",
self.logging_handler.messages['warning'][0])
def test_channel_build_message_headers_out_of_order(self):
channel = Channel(0, None, 360)
deliver = specification.Basic.Deliver()
header = ContentHeader(body_size=10)
channel._inbound = [header, deliver]
result = channel._build_message_headers()
self.assertEqual(result, None)
self.assertIn("Received an out-of-order frame:",
self.logging_handler.messages['warning'][0])
self.logging_handler.messages['warning'] = []
channel._inbound = [deliver, deliver]
result = channel._build_message_headers()
self.assertEqual(result, None)
self.assertIn("Received an out-of-order frame:",
self.logging_handler.messages['warning'][0])
async def _consumer_task(self, queue_name: str, consumer_tag: str,
channel_id: int) -> None:
_logger.debug(f"* New consumer {consumer_tag}")
async for message in self._on_consume(queue_name):
_logger.debug(f"--> Message {message}")
delivery_tag = self._get_delivery_tag()
self._delivered_messages[delivery_tag] = message.id
frame_out = spec.Basic.Deliver(
consumer_tag=consumer_tag,
delivery_tag=delivery_tag,
exchange=message.exchange,
routing_key=message.routing_key,
)
await self._send_frame(channel_id, frame_out)
encoded = json.dumps(message.value).encode()
header, body = ContentHeader(
body_size=len(encoded)), ContentBody(encoded)
await self._send_frame(channel_id, header)
await self._send_frame(channel_id, body)
def _marshal_content_header_frame(value: header.ContentHeader,
channel_id: int) -> bytes:
"""Marshal a content header frame"""
return _marshal(constants.FRAME_HEADER, channel_id, value.marshal())
def on_get_frame(*_):
channel.rpc.on_frame(specification.Basic.GetOk())
channel.rpc.on_frame(ContentHeader(body_size=message_len))
channel.rpc.on_frame(ContentBody(value=message))
async def basic_publish(
self,
body: bytes,
*,
exchange: str = "",
routing_key: str = "",
properties: spec.Basic.Properties = None,
mandatory: bool = False,
immediate: bool = False,
timeout: Union[int,
float] = None) -> Optional[ConfirmationFrameType]:
frame = spec.Basic.Publish(
exchange=exchange,
routing_key=routing_key,
mandatory=mandatory,
immediate=immediate,
)
content_header = ContentHeader(
properties=properties or spec.Basic.Properties(delivery_mode=1),
body_size=len(body),
)
if not content_header.properties.message_id:
# UUID compatible random bytes
rnd_id = os.urandom(16)
content_header.properties.message_id = hexlify(rnd_id).decode()
confirmation = None
async with self.lock:
self.delivery_tag += 1
if self.publisher_confirms:
message_id = content_header.properties.message_id
if self.delivery_tag not in self.confirmations:
self.confirmations[
self.delivery_tag] = self.create_future()
confirmation = self.confirmations[self.delivery_tag]
self.message_id_delivery_tag[message_id] = self.delivery_tag
confirmation.add_done_callback(
lambda _: self.message_id_delivery_tag.pop(
message_id,
None,
), )
self.writer.write(pamqp.frame.marshal(frame, self.number))
# noinspection PyTypeChecker
self.writer.write(pamqp.frame.marshal(content_header, self.number))
with BytesIO(body) as buf:
read_chunk = partial(buf.read, self.max_content_size)
reader = iter(read_chunk, b"")
for chunk in reader:
# noinspection PyTypeChecker
self.writer.write(
pamqp.frame.marshal(ContentBody(chunk), self.number), )
if not self.publisher_confirms:
return
return await asyncio.wait_for(confirmation, timeout=timeout)
async def basic_publish(
self,
body: bytes,
*,
exchange: str = "",
routing_key: str = "",
properties: spec.Basic.Properties = None,
mandatory: bool = False,
immediate: bool = False,
timeout: TimeoutType = None,
wait: bool = True,
) -> Optional[ConfirmationFrameType]:
body_io = BytesIO(body)
drain_future = self.create_future() if wait else None
countdown = Countdown(timeout=timeout)
publish_frame = spec.Basic.Publish(
exchange=exchange,
routing_key=routing_key,
mandatory=mandatory,
immediate=immediate,
)
content_header = ContentHeader(
properties=properties or spec.Basic.Properties(delivery_mode=1),
body_size=len(body),
)
if not content_header.properties.message_id:
# UUID compatible random bytes
rnd_uuid = UUID(int=getrandbits(128), version=4)
content_header.properties.message_id = rnd_uuid.hex
confirmation: Optional[ConfirmationType] = None
async with countdown.enter_context(self.lock):
self.delivery_tag += 1
if self.publisher_confirms:
message_id = content_header.properties.message_id
if self.delivery_tag not in self.confirmations:
self.confirmations[
self.delivery_tag] = self.create_future()
confirmation = self.confirmations[self.delivery_tag]
self.message_id_delivery_tag[message_id] = self.delivery_tag
if confirmation is None:
return
confirmation.add_done_callback(
lambda _: self.message_id_delivery_tag.pop(
message_id,
None,
), )
def frame_generator() -> Generator[FrameType, Any, None]:
yield publish_frame
yield content_header
with body_io as fp:
read_chunk = partial(fp.read, self.max_content_size)
for chunk in iter(read_chunk, b""):
yield ContentBody(chunk)
await countdown(
self.write_queue.put(
ChannelFrame(
frames=frame_generator(),
channel_number=self.number,
drain_future=drain_future,
), ), )
if drain_future:
await drain_future
if not self.publisher_confirms:
return None
if confirmation is None:
return None
return await countdown(confirmation)
def on_get_frame(*_):
channel.rpc.on_frame(commands.Basic.GetOk())
channel.rpc.on_frame(ContentHeader(body_size=message_len))
channel.rpc.on_frame(ContentBody(value=message))