async def server_mock(stream):
packet = await PublishPacket.from_stream(stream)
self.assertEqual(packet.topic_name, "/topic")
self.assertEqual(packet.qos, QOS_2)
self.assertIsNotNone(packet.packet_id)
self.assertIn(packet.packet_id, self.session.inflight_out)
self.assertIn(packet.packet_id, self.handler._pubrec_waiters)
pubrec = PubrecPacket.build(packet.packet_id)
await pubrec.to_stream(stream)
await PubrelPacket.from_stream(stream)
self.assertIn(packet.packet_id, self.handler._pubcomp_waiters)
pubcomp = PubcompPacket.build(packet.packet_id)
await pubcomp.to_stream(stream)
async def _handle_qos2_message_flow(self, app_message):
"""
Handle QOS_2 application message acknowledgment
For incoming messages, this method stores the message, sends PUBREC, waits for PUBREL, initiate delivery
and send PUBCOMP
For outgoing messages, this methods sends PUBLISH, waits for PUBREC, discards messages and wait for PUBCOMP
:param app_message:
:return:
"""
assert app_message.qos == QOS_2
if app_message.direction == OUTGOING:
if app_message.pubrel_packet and app_message.pubcomp_packet:
raise DistMQTTException(
"Message '%d' has already been acknowledged" %
app_message.packet_id)
if not app_message.pubrel_packet:
# Store message
if app_message.publish_packet is not None:
# This is a retry flow, no need to store just check the message exists in session
if app_message.packet_id not in self.session.inflight_out:
raise DistMQTTException(
"Unknown inflight message '%d' in session" %
app_message.packet_id)
publish_packet = app_message.build_publish_packet(dup=True)
else:
# Store message in session
self.session.inflight_out[
app_message.packet_id] = app_message
publish_packet = app_message.build_publish_packet()
# Wait PUBREC
if app_message.packet_id in self._pubrec_waiters:
# PUBREC waiter already exists for this packet ID
self.logger.warning(
"Can't add PUBREC waiter, a waiter already exists for message Id '%s'",
app_message.packet_id,
)
raise DistMQTTException(app_message)
waiter = Future()
self._pubrec_waiters[app_message.packet_id] = waiter
# Send PUBLISH packet
try:
await self._send_packet(publish_packet)
app_message.publish_packet = publish_packet
app_message.pubrec_packet = await waiter.get()
finally:
del self._pubrec_waiters[app_message.packet_id]
if not app_message.pubcomp_packet:
# Wait for PUBCOMP
waiter = Future()
self._pubcomp_waiters[app_message.packet_id] = waiter
# Send pubrel
try:
app_message.pubrel_packet = PubrelPacket.build(
app_message.packet_id)
await self._send_packet(app_message.pubrel_packet)
app_message.pubcomp_packet = await waiter.get()
finally:
del self._pubcomp_waiters[app_message.packet_id]
# Discard inflight message
del self.session.inflight_out[app_message.packet_id]
elif app_message.direction == INCOMING:
self.session.inflight_in[app_message.packet_id] = app_message
# Wait PUBREL
if (app_message.packet_id in self._pubrel_waiters and
not self._pubrel_waiters[app_message.packet_id].done()):
# PUBREL waiter already exists for this packet ID
self.logger.warning(
"A waiter already exists for message Id '%s', canceling it",
app_message.packet_id,
)
await self._pubrel_waiters[app_message.packet_id].cancel()
waiter = Future()
self._pubrel_waiters[app_message.packet_id] = waiter
# Send pubrec
try:
pubrec_packet = PubrecPacket.build(app_message.packet_id)
await self._send_packet(pubrec_packet)
app_message.pubrec_packet = pubrec_packet
app_message.pubrel_packet = await waiter.get()
finally:
if self._pubrel_waiters.get(app_message.packet_id) is waiter:
del self._pubrel_waiters[app_message.packet_id]
# Initiate delivery and discard message
await self.session.put_message(app_message)
del self.session.inflight_in[app_message.packet_id]
# Send pubcomp
pubcomp_packet = PubcompPacket.build(app_message.packet_id)
await self._send_packet(pubcomp_packet)
app_message.pubcomp_packet = pubcomp_packet
def test_to_bytes(self):
variable_header = PacketIdVariableHeader(10)
publish = PubrecPacket(variable_header=variable_header)
out = publish.to_bytes()
self.assertEqual(out, b'\x50\x02\x00\x0a')