kopia lustrzana https://github.com/Yakifo/amqtt
499 wiersze
20 KiB
Python
499 wiersze
20 KiB
Python
# Copyright (c) 2015 Nicolas JOUANIN
|
|
#
|
|
# See the file license.txt for copying permission.
|
|
import asyncio
|
|
import logging
|
|
import random
|
|
import unittest
|
|
|
|
from amqtt.adapters import StreamReaderAdapter, StreamWriterAdapter
|
|
from amqtt.mqtt.constants import QOS_0, QOS_1, QOS_2
|
|
from amqtt.mqtt.protocol.handler import ProtocolHandler
|
|
from amqtt.mqtt.puback import PubackPacket
|
|
from amqtt.mqtt.pubcomp import PubcompPacket
|
|
from amqtt.mqtt.publish import PublishPacket
|
|
from amqtt.mqtt.pubrec import PubrecPacket
|
|
from amqtt.mqtt.pubrel import PubrelPacket
|
|
from amqtt.plugins.manager import PluginManager
|
|
from amqtt.session import (
|
|
IncomingApplicationMessage,
|
|
OutgoingApplicationMessage,
|
|
Session,
|
|
)
|
|
|
|
formatter = "[%(asctime)s] %(name)s {%(filename)s:%(lineno)d} %(levelname)s - %(message)s"
|
|
logging.basicConfig(level=logging.DEBUG, format=formatter)
|
|
log = logging.getLogger(__name__)
|
|
|
|
|
|
def rand_packet_id():
|
|
return random.randint(0, 65535)
|
|
|
|
|
|
def adapt(reader, writer):
|
|
return StreamReaderAdapter(reader), StreamWriterAdapter(writer)
|
|
|
|
|
|
class ProtocolHandlerTest(unittest.TestCase):
|
|
def setUp(self):
|
|
self.loop = asyncio.new_event_loop()
|
|
asyncio.set_event_loop(self.loop)
|
|
self.plugin_manager = PluginManager("amqtt.test.plugins", context=None)
|
|
|
|
def tearDown(self):
|
|
self.loop.close()
|
|
|
|
def test_init_handler(self):
|
|
Session()
|
|
handler = ProtocolHandler(self.plugin_manager)
|
|
assert handler.session is None
|
|
assert handler._loop is self.loop
|
|
self.check_empty_waiters(handler)
|
|
|
|
def test_start_stop(self):
|
|
async def server_mock(reader, writer) -> None:
|
|
pass
|
|
|
|
async def test_coro() -> None:
|
|
try:
|
|
s = Session()
|
|
reader, writer = await asyncio.open_connection("127.0.0.1", 8888)
|
|
reader_adapted, writer_adapted = adapt(reader, writer)
|
|
handler = ProtocolHandler(self.plugin_manager)
|
|
handler.attach(s, reader_adapted, writer_adapted)
|
|
await self.start_handler(handler, s)
|
|
await self.stop_handler(handler, s)
|
|
future.set_result(True)
|
|
except Exception as ae:
|
|
future.set_exception(ae)
|
|
|
|
future = asyncio.Future()
|
|
coro = asyncio.start_server(server_mock, "127.0.0.1", 8888)
|
|
server = self.loop.run_until_complete(coro)
|
|
self.loop.run_until_complete(test_coro())
|
|
server.close()
|
|
self.loop.run_until_complete(server.wait_closed())
|
|
if future.exception():
|
|
raise future.exception()
|
|
|
|
def test_publish_qos0(self):
|
|
async def server_mock(reader, writer) -> None:
|
|
try:
|
|
packet = await PublishPacket.from_stream(reader)
|
|
assert packet.variable_header.topic_name == "/topic"
|
|
assert packet.qos == QOS_0
|
|
assert packet.packet_id is None
|
|
except Exception as ae:
|
|
future.set_exception(ae)
|
|
|
|
async def test_coro() -> None:
|
|
try:
|
|
s = Session()
|
|
reader, writer = await asyncio.open_connection("127.0.0.1", 8888)
|
|
reader_adapted, writer_adapted = adapt(reader, writer)
|
|
handler = ProtocolHandler(self.plugin_manager)
|
|
handler.attach(s, reader_adapted, writer_adapted)
|
|
await self.start_handler(handler, s)
|
|
message = await handler.mqtt_publish(
|
|
"/topic",
|
|
b"test_data",
|
|
QOS_0,
|
|
False,
|
|
)
|
|
assert isinstance(message, OutgoingApplicationMessage)
|
|
assert message.publish_packet is not None
|
|
assert message.puback_packet is None
|
|
assert message.pubrec_packet is None
|
|
assert message.pubrel_packet is None
|
|
assert message.pubcomp_packet is None
|
|
await self.stop_handler(handler, s)
|
|
future.set_result(True)
|
|
except Exception as ae:
|
|
future.set_exception(ae)
|
|
|
|
future = asyncio.Future()
|
|
coro = asyncio.start_server(server_mock, "127.0.0.1", 8888)
|
|
server = self.loop.run_until_complete(coro)
|
|
self.loop.run_until_complete(test_coro())
|
|
server.close()
|
|
self.loop.run_until_complete(server.wait_closed())
|
|
if future.exception():
|
|
raise future.exception()
|
|
|
|
def test_publish_qos1(self):
|
|
async def server_mock(reader, writer) -> None:
|
|
packet = await PublishPacket.from_stream(reader)
|
|
try:
|
|
assert packet.variable_header.topic_name == "/topic"
|
|
assert packet.qos == QOS_1
|
|
assert packet.packet_id is not None
|
|
assert packet.packet_id in self.session.inflight_out
|
|
assert packet.packet_id in self.handler._puback_waiters
|
|
puback = PubackPacket.build(packet.packet_id)
|
|
await puback.to_stream(writer)
|
|
except Exception as ae:
|
|
future.set_exception(ae)
|
|
|
|
async def test_coro() -> None:
|
|
try:
|
|
reader, writer = await asyncio.open_connection("127.0.0.1", 8888)
|
|
reader_adapted, writer_adapted = adapt(reader, writer)
|
|
self.handler = ProtocolHandler(self.plugin_manager)
|
|
self.handler.attach(self.session, reader_adapted, writer_adapted)
|
|
await self.start_handler(self.handler, self.session)
|
|
message = await self.handler.mqtt_publish(
|
|
"/topic",
|
|
b"test_data",
|
|
QOS_1,
|
|
False,
|
|
)
|
|
assert isinstance(message, OutgoingApplicationMessage)
|
|
assert message.publish_packet is not None
|
|
assert message.puback_packet is not None
|
|
assert message.pubrec_packet is None
|
|
assert message.pubrel_packet is None
|
|
assert message.pubcomp_packet is None
|
|
await self.stop_handler(self.handler, self.session)
|
|
if not future.done():
|
|
future.set_result(True)
|
|
except Exception as ae:
|
|
future.set_exception(ae)
|
|
|
|
self.handler = None
|
|
self.session = Session()
|
|
future = asyncio.Future()
|
|
|
|
coro = asyncio.start_server(server_mock, "127.0.0.1", 8888)
|
|
server = self.loop.run_until_complete(coro)
|
|
self.loop.run_until_complete(test_coro())
|
|
server.close()
|
|
self.loop.run_until_complete(server.wait_closed())
|
|
if future.exception():
|
|
raise future.exception()
|
|
|
|
def test_publish_qos2(self):
|
|
async def server_mock(reader, writer) -> None:
|
|
try:
|
|
packet = await PublishPacket.from_stream(reader)
|
|
assert packet.topic_name == "/topic"
|
|
assert packet.qos == QOS_2
|
|
assert packet.packet_id is not None
|
|
assert packet.packet_id in self.session.inflight_out
|
|
assert packet.packet_id in self.handler._pubrec_waiters
|
|
pubrec = PubrecPacket.build(packet.packet_id)
|
|
await pubrec.to_stream(writer)
|
|
|
|
await PubrelPacket.from_stream(reader)
|
|
assert packet.packet_id in self.handler._pubcomp_waiters
|
|
pubcomp = PubcompPacket.build(packet.packet_id)
|
|
await pubcomp.to_stream(writer)
|
|
except Exception as ae:
|
|
future.set_exception(ae)
|
|
|
|
async def test_coro() -> None:
|
|
try:
|
|
reader, writer = await asyncio.open_connection("127.0.0.1", 8888)
|
|
reader_adapted, writer_adapted = adapt(reader, writer)
|
|
self.handler = ProtocolHandler(self.plugin_manager)
|
|
self.handler.attach(self.session, reader_adapted, writer_adapted)
|
|
await self.start_handler(self.handler, self.session)
|
|
message = await self.handler.mqtt_publish(
|
|
"/topic",
|
|
b"test_data",
|
|
QOS_2,
|
|
False,
|
|
)
|
|
assert isinstance(message, OutgoingApplicationMessage)
|
|
assert message.publish_packet is not None
|
|
assert message.puback_packet is None
|
|
assert message.pubrec_packet is not None
|
|
assert message.pubrel_packet is not None
|
|
assert message.pubcomp_packet is not None
|
|
await self.stop_handler(self.handler, self.session)
|
|
if not future.done():
|
|
future.set_result(True)
|
|
except Exception as ae:
|
|
future.set_exception(ae)
|
|
|
|
self.handler = None
|
|
self.session = Session()
|
|
future = asyncio.Future()
|
|
|
|
coro = asyncio.start_server(server_mock, "127.0.0.1", 8888)
|
|
server = self.loop.run_until_complete(coro)
|
|
self.loop.run_until_complete(test_coro())
|
|
server.close()
|
|
self.loop.run_until_complete(server.wait_closed())
|
|
if future.exception():
|
|
raise future.exception()
|
|
|
|
def test_receive_qos0(self):
|
|
async def server_mock(reader, writer) -> None:
|
|
packet = PublishPacket.build(
|
|
"/topic",
|
|
b"test_data",
|
|
rand_packet_id(),
|
|
False,
|
|
QOS_0,
|
|
False,
|
|
)
|
|
await packet.to_stream(writer)
|
|
|
|
async def test_coro() -> None:
|
|
try:
|
|
reader, writer = await asyncio.open_connection("127.0.0.1", 8888)
|
|
reader_adapted, writer_adapted = adapt(reader, writer)
|
|
self.handler = ProtocolHandler(self.plugin_manager)
|
|
self.handler.attach(self.session, reader_adapted, writer_adapted)
|
|
await self.start_handler(self.handler, self.session)
|
|
message = await self.handler.mqtt_deliver_next_message()
|
|
assert isinstance(message, IncomingApplicationMessage)
|
|
assert message.publish_packet is not None
|
|
assert message.puback_packet is None
|
|
assert message.pubrec_packet is None
|
|
assert message.pubrel_packet is None
|
|
assert message.pubcomp_packet is None
|
|
await self.stop_handler(self.handler, self.session)
|
|
future.set_result(True)
|
|
except Exception as ae:
|
|
future.set_exception(ae)
|
|
|
|
self.handler = None
|
|
self.session = Session()
|
|
future = asyncio.Future()
|
|
coro = asyncio.start_server(server_mock, "127.0.0.1", 8888)
|
|
server = self.loop.run_until_complete(coro)
|
|
self.loop.run_until_complete(test_coro())
|
|
server.close()
|
|
self.loop.run_until_complete(server.wait_closed())
|
|
if future.exception():
|
|
raise future.exception()
|
|
|
|
def test_receive_qos1(self):
|
|
async def server_mock(reader, writer) -> None:
|
|
try:
|
|
packet = PublishPacket.build(
|
|
"/topic",
|
|
b"test_data",
|
|
rand_packet_id(),
|
|
False,
|
|
QOS_1,
|
|
False,
|
|
)
|
|
await packet.to_stream(writer)
|
|
puback = await PubackPacket.from_stream(reader)
|
|
assert puback is not None
|
|
assert packet.packet_id == puback.packet_id
|
|
except Exception as ae:
|
|
future.set_exception(ae)
|
|
|
|
async def test_coro() -> None:
|
|
try:
|
|
reader, writer = await asyncio.open_connection("127.0.0.1", 8888)
|
|
reader_adapted, writer_adapted = adapt(reader, writer)
|
|
self.handler = ProtocolHandler(self.plugin_manager)
|
|
self.handler.attach(self.session, reader_adapted, writer_adapted)
|
|
await self.start_handler(self.handler, self.session)
|
|
message = await self.handler.mqtt_deliver_next_message()
|
|
assert isinstance(message, IncomingApplicationMessage)
|
|
assert message.publish_packet is not None
|
|
assert message.puback_packet is not None
|
|
assert message.pubrec_packet is None
|
|
assert message.pubrel_packet is None
|
|
assert message.pubcomp_packet is None
|
|
await self.stop_handler(self.handler, self.session)
|
|
future.set_result(True)
|
|
except Exception as ae:
|
|
future.set_exception(ae)
|
|
|
|
self.handler = None
|
|
self.session = Session()
|
|
future = asyncio.Future()
|
|
self.event = asyncio.Event()
|
|
coro = asyncio.start_server(server_mock, "127.0.0.1", 8888)
|
|
server = self.loop.run_until_complete(coro)
|
|
self.loop.run_until_complete(test_coro())
|
|
server.close()
|
|
self.loop.run_until_complete(server.wait_closed())
|
|
if future.exception():
|
|
raise future.exception()
|
|
|
|
def test_receive_qos2(self):
|
|
async def server_mock(reader, writer) -> None:
|
|
try:
|
|
packet = PublishPacket.build(
|
|
"/topic",
|
|
b"test_data",
|
|
rand_packet_id(),
|
|
False,
|
|
QOS_2,
|
|
False,
|
|
)
|
|
await packet.to_stream(writer)
|
|
pubrec = await PubrecPacket.from_stream(reader)
|
|
assert pubrec is not None
|
|
assert packet.packet_id == pubrec.packet_id
|
|
assert packet.packet_id in self.handler._pubrel_waiters
|
|
pubrel = PubrelPacket.build(packet.packet_id)
|
|
await pubrel.to_stream(writer)
|
|
pubcomp = await PubcompPacket.from_stream(reader)
|
|
assert pubcomp is not None
|
|
assert packet.packet_id == pubcomp.packet_id
|
|
except Exception as ae:
|
|
future.set_exception(ae)
|
|
|
|
async def test_coro() -> None:
|
|
try:
|
|
reader, writer = await asyncio.open_connection("127.0.0.1", 8888)
|
|
reader_adapted, writer_adapted = adapt(reader, writer)
|
|
self.handler = ProtocolHandler(self.plugin_manager)
|
|
self.handler.attach(self.session, reader_adapted, writer_adapted)
|
|
await self.start_handler(self.handler, self.session)
|
|
message = await self.handler.mqtt_deliver_next_message()
|
|
assert isinstance(message, IncomingApplicationMessage)
|
|
assert message.publish_packet is not None
|
|
assert message.puback_packet is None
|
|
assert message.pubrec_packet is not None
|
|
assert message.pubrel_packet is not None
|
|
assert message.pubcomp_packet is not None
|
|
await self.stop_handler(self.handler, self.session)
|
|
future.set_result(True)
|
|
except Exception as ae:
|
|
future.set_exception(ae)
|
|
|
|
self.handler = None
|
|
self.session = Session()
|
|
future = asyncio.Future()
|
|
coro = asyncio.start_server(server_mock, "127.0.0.1", 8888)
|
|
server = self.loop.run_until_complete(coro)
|
|
self.loop.run_until_complete(test_coro())
|
|
server.close()
|
|
self.loop.run_until_complete(server.wait_closed())
|
|
if future.exception():
|
|
raise future.exception()
|
|
|
|
async def start_handler(self, handler, session):
|
|
self.check_empty_waiters(handler)
|
|
self.check_no_message(session)
|
|
await handler.start()
|
|
assert handler._reader_ready
|
|
|
|
async def stop_handler(self, handler, session):
|
|
await handler.stop()
|
|
assert handler._reader_stopped
|
|
self.check_empty_waiters(handler)
|
|
self.check_no_message(session)
|
|
|
|
def check_empty_waiters(self, handler):
|
|
assert not handler._puback_waiters
|
|
assert not handler._pubrec_waiters
|
|
assert not handler._pubrel_waiters
|
|
assert not handler._pubcomp_waiters
|
|
|
|
def check_no_message(self, session):
|
|
assert not session.inflight_out
|
|
assert not session.inflight_in
|
|
|
|
def test_publish_qos1_retry(self):
|
|
async def server_mock(reader, writer) -> None:
|
|
packet = await PublishPacket.from_stream(reader)
|
|
try:
|
|
assert packet.topic_name == "/topic"
|
|
assert packet.qos == QOS_1
|
|
assert packet.packet_id is not None
|
|
assert packet.packet_id in self.session.inflight_out
|
|
assert packet.packet_id in self.handler._puback_waiters
|
|
puback = PubackPacket.build(packet.packet_id)
|
|
await puback.to_stream(writer)
|
|
except Exception as ae:
|
|
future.set_exception(ae)
|
|
|
|
async def test_coro() -> None:
|
|
try:
|
|
reader, writer = await asyncio.open_connection("127.0.0.1", 8888)
|
|
reader_adapted, writer_adapted = adapt(reader, writer)
|
|
self.handler = ProtocolHandler(self.plugin_manager)
|
|
self.handler.attach(self.session, reader_adapted, writer_adapted)
|
|
await self.handler.start()
|
|
await self.stop_handler(self.handler, self.session)
|
|
if not future.done():
|
|
future.set_result(True)
|
|
except Exception as ae:
|
|
future.set_exception(ae)
|
|
|
|
self.handler = None
|
|
self.session = Session()
|
|
message = OutgoingApplicationMessage(1, "/topic", QOS_1, b"test_data", False)
|
|
message.publish_packet = PublishPacket.build(
|
|
"/topic",
|
|
b"test_data",
|
|
rand_packet_id(),
|
|
False,
|
|
QOS_1,
|
|
False,
|
|
)
|
|
self.session.inflight_out[1] = message
|
|
future = asyncio.Future()
|
|
|
|
coro = asyncio.start_server(server_mock, "127.0.0.1", 8888)
|
|
server = self.loop.run_until_complete(coro)
|
|
self.loop.run_until_complete(test_coro())
|
|
server.close()
|
|
self.loop.run_until_complete(server.wait_closed())
|
|
if future.exception():
|
|
raise future.exception()
|
|
|
|
def test_publish_qos2_retry(self):
|
|
async def server_mock(reader, writer) -> None:
|
|
try:
|
|
packet = await PublishPacket.from_stream(reader)
|
|
assert packet.topic_name == "/topic"
|
|
assert packet.qos == QOS_2
|
|
assert packet.packet_id is not None
|
|
assert packet.packet_id in self.session.inflight_out
|
|
assert packet.packet_id in self.handler._pubrec_waiters
|
|
pubrec = PubrecPacket.build(packet.packet_id)
|
|
await pubrec.to_stream(writer)
|
|
|
|
await PubrelPacket.from_stream(reader)
|
|
assert packet.packet_id in self.handler._pubcomp_waiters
|
|
pubcomp = PubcompPacket.build(packet.packet_id)
|
|
await pubcomp.to_stream(writer)
|
|
except Exception as ae:
|
|
future.set_exception(ae)
|
|
|
|
async def test_coro() -> None:
|
|
try:
|
|
reader, writer = await asyncio.open_connection("127.0.0.1", 8888)
|
|
reader_adapted, writer_adapted = adapt(reader, writer)
|
|
self.handler = ProtocolHandler(self.plugin_manager)
|
|
self.handler.attach(self.session, reader_adapted, writer_adapted)
|
|
await self.handler.start()
|
|
await self.stop_handler(self.handler, self.session)
|
|
if not future.done():
|
|
future.set_result(True)
|
|
except Exception as ae:
|
|
future.set_exception(ae)
|
|
|
|
self.handler = None
|
|
self.session = Session()
|
|
message = OutgoingApplicationMessage(1, "/topic", QOS_2, b"test_data", False)
|
|
message.publish_packet = PublishPacket.build(
|
|
"/topic",
|
|
b"test_data",
|
|
rand_packet_id(),
|
|
False,
|
|
QOS_2,
|
|
False,
|
|
)
|
|
self.session.inflight_out[1] = message
|
|
future = asyncio.Future()
|
|
|
|
coro = asyncio.start_server(server_mock, "127.0.0.1", 8888)
|
|
server = self.loop.run_until_complete(coro)
|
|
self.loop.run_until_complete(test_coro())
|
|
server.close()
|
|
self.loop.run_until_complete(server.wait_closed())
|
|
if future.exception():
|
|
raise future.exception()
|