amqtt/tests/mqtt/protocol/test_handler.py

# Copyright (c) 2015 Nicolas JOUANIN
#
# See the file license.txt for copying permission.
import unittest
import asyncio
import logging
import random
from hbmqtt.plugins.manager import PluginManager
from hbmqtt.session import Session, OutgoingApplicationMessage, IncomingApplicationMessage
from hbmqtt.mqtt.protocol.handler import ProtocolHandler
from hbmqtt.adapters import StreamWriterAdapter, StreamReaderAdapter
from hbmqtt.mqtt.constants import QOS_0, QOS_1, QOS_2
from hbmqtt.mqtt.publish import PublishPacket
from hbmqtt.mqtt.puback import PubackPacket
from hbmqtt.mqtt.pubrec import PubrecPacket
from hbmqtt.mqtt.pubrel import PubrelPacket
from hbmqtt.mqtt.pubcomp import PubcompPacket

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("hbmqtt.test.plugins", context=None, loop=self.loop)

    def tearDown(self):
        self.loop.close()

    def test_init_handler(self):
        Session()
        handler = ProtocolHandler(self.plugin_manager, loop=self.loop)
        self.assertIsNone(handler.session)
        self.assertIs(handler._loop, self.loop)
        self.check_empty_waiters(handler)

    def test_start_stop(self):
        async def server_mock(reader, writer):
            pass

        async def test_coro():
            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(loop=self.loop)
        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):
            try:
                packet = await PublishPacket.from_stream(reader)
                self.assertEqual(packet.variable_header.topic_name, "/topic")
                self.assertEqual(packet.qos, QOS_0)
                self.assertIsNone(packet.packet_id)
            except Exception as ae:
                future.set_exception(ae)

        async def test_coro():
            try:
                s = Session()
                reader, writer = await asyncio.open_connection(
                    "127.0.0.1", 8888, loop=self.loop
                )
                reader_adapted, writer_adapted = adapt(reader, writer)
                handler = ProtocolHandler(self.plugin_manager, loop=self.loop)
                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)
                self.assertIsInstance(message, OutgoingApplicationMessage)
                self.assertIsNotNone(message.publish_packet)
                self.assertIsNone(message.puback_packet)
                self.assertIsNone(message.pubrec_packet)
                self.assertIsNone(message.pubrel_packet)
                self.assertIsNone(message.pubcomp_packet)
                await self.stop_handler(handler, s)
                future.set_result(True)
            except Exception as ae:
                future.set_exception(ae)

        future = asyncio.Future(loop=self.loop)
        coro = asyncio.start_server(server_mock, "127.0.0.1", 8888, loop=self.loop)
        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):
            packet = await PublishPacket.from_stream(reader)
            try:
                self.assertEqual(packet.variable_header.topic_name, "/topic")
                self.assertEqual(packet.qos, QOS_1)
                self.assertIsNotNone(packet.packet_id)
                self.assertIn(packet.packet_id, self.session.inflight_out)
                self.assertIn(packet.packet_id, 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():
            try:
                reader, writer = await asyncio.open_connection(
                    "127.0.0.1", 8888, loop=self.loop
                )
                reader_adapted, writer_adapted = adapt(reader, writer)
                self.handler = ProtocolHandler(self.plugin_manager, loop=self.loop)
                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)
                self.assertIsInstance(message, OutgoingApplicationMessage)
                self.assertIsNotNone(message.publish_packet)
                self.assertIsNotNone(message.puback_packet)
                self.assertIsNone(message.pubrec_packet)
                self.assertIsNone(message.pubrel_packet)
                self.assertIsNone(message.pubcomp_packet)
                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(loop=self.loop)

        coro = asyncio.start_server(server_mock, "127.0.0.1", 8888, loop=self.loop)
        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):
            try:
                packet = await PublishPacket.from_stream(reader)
                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(writer)

                await PubrelPacket.from_stream(reader)
                self.assertIn(packet.packet_id, 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():
            try:
                reader, writer = await asyncio.open_connection(
                    "127.0.0.1", 8888, loop=self.loop
                )
                reader_adapted, writer_adapted = adapt(reader, writer)
                self.handler = ProtocolHandler(self.plugin_manager, loop=self.loop)
                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)
                self.assertIsInstance(message, OutgoingApplicationMessage)
                self.assertIsNotNone(message.publish_packet)
                self.assertIsNone(message.puback_packet)
                self.assertIsNotNone(message.pubrec_packet)
                self.assertIsNotNone(message.pubrel_packet)
                self.assertIsNotNone(message.pubcomp_packet)
                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(loop=self.loop)

        coro = asyncio.start_server(server_mock, "127.0.0.1", 8888, loop=self.loop)
        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):
            packet = PublishPacket.build(
                "/topic", b"test_data", rand_packet_id(), False, QOS_0, False
            )
            await packet.to_stream(writer)

        async def test_coro():
            try:
                reader, writer = await asyncio.open_connection(
                    "127.0.0.1", 8888, loop=self.loop
                )
                reader_adapted, writer_adapted = adapt(reader, writer)
                self.handler = ProtocolHandler(self.plugin_manager, loop=self.loop)
                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()
                self.assertIsInstance(message, IncomingApplicationMessage)
                self.assertIsNotNone(message.publish_packet)
                self.assertIsNone(message.puback_packet)
                self.assertIsNone(message.pubrec_packet)
                self.assertIsNone(message.pubrel_packet)
                self.assertIsNone(message.pubcomp_packet)
                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(loop=self.loop)
        coro = asyncio.start_server(server_mock, "127.0.0.1", 8888, loop=self.loop)
        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):
            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)
                self.assertIsNotNone(puback)
                self.assertEqual(packet.packet_id, puback.packet_id)
            except Exception as ae:
                print(ae)
                future.set_exception(ae)

        async def test_coro():
            try:
                reader, writer = await asyncio.open_connection(
                    "127.0.0.1", 8888, loop=self.loop
                )
                reader_adapted, writer_adapted = adapt(reader, writer)
                self.handler = ProtocolHandler(self.plugin_manager, loop=self.loop)
                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()
                self.assertIsInstance(message, IncomingApplicationMessage)
                self.assertIsNotNone(message.publish_packet)
                self.assertIsNotNone(message.puback_packet)
                self.assertIsNone(message.pubrec_packet)
                self.assertIsNone(message.pubrel_packet)
                self.assertIsNone(message.pubcomp_packet)
                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(loop=self.loop)
        self.event = asyncio.Event(loop=self.loop)
        coro = asyncio.start_server(server_mock, "127.0.0.1", 8888, loop=self.loop)
        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):
            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)
                self.assertIsNotNone(pubrec)
                self.assertEqual(packet.packet_id, pubrec.packet_id)
                self.assertIn(packet.packet_id, self.handler._pubrel_waiters)
                pubrel = PubrelPacket.build(packet.packet_id)
                await pubrel.to_stream(writer)
                pubcomp = await PubcompPacket.from_stream(reader)
                self.assertIsNotNone(pubcomp)
                self.assertEqual(packet.packet_id, pubcomp.packet_id)
            except Exception as ae:
                future.set_exception(ae)

        async def test_coro():
            try:
                reader, writer = await asyncio.open_connection(
                    "127.0.0.1", 8888, loop=self.loop
                )
                reader_adapted, writer_adapted = adapt(reader, writer)
                self.handler = ProtocolHandler(self.plugin_manager, loop=self.loop)
                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()
                self.assertIsInstance(message, IncomingApplicationMessage)
                self.assertIsNotNone(message.publish_packet)
                self.assertIsNone(message.puback_packet)
                self.assertIsNotNone(message.pubrec_packet)
                self.assertIsNotNone(message.pubrel_packet)
                self.assertIsNotNone(message.pubcomp_packet)
                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(loop=self.loop)
        coro = asyncio.start_server(server_mock, "127.0.0.1", 8888, loop=self.loop)
        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):
        self.assertFalse(handler._puback_waiters)
        self.assertFalse(handler._pubrec_waiters)
        self.assertFalse(handler._pubrel_waiters)
        self.assertFalse(handler._pubcomp_waiters)

    def check_no_message(self, session):
        self.assertFalse(session.inflight_out)
        self.assertFalse(session.inflight_in)

    def test_publish_qos1_retry(self):
        async def server_mock(reader, writer):
            packet = await PublishPacket.from_stream(reader)
            try:
                self.assertEqual(packet.topic_name, "/topic")
                self.assertEqual(packet.qos, QOS_1)
                self.assertIsNotNone(packet.packet_id)
                self.assertIn(packet.packet_id, self.session.inflight_out)
                self.assertIn(packet.packet_id, 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():
            try:
                reader, writer = await asyncio.open_connection(
                    "127.0.0.1", 8888, loop=self.loop
                )
                reader_adapted, writer_adapted = adapt(reader, writer)
                self.handler = ProtocolHandler(self.plugin_manager, loop=self.loop)
                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(loop=self.loop)

        coro = asyncio.start_server(server_mock, "127.0.0.1", 8888, loop=self.loop)
        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):
            try:
                packet = await PublishPacket.from_stream(reader)
                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(writer)

                await PubrelPacket.from_stream(reader)
                self.assertIn(packet.packet_id, 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():
            try:
                reader, writer = await asyncio.open_connection(
                    "127.0.0.1", 8888, loop=self.loop
                )
                reader_adapted, writer_adapted = adapt(reader, writer)
                self.handler = ProtocolHandler(self.plugin_manager, loop=self.loop)
                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(loop=self.loop)

        coro = asyncio.start_server(server_mock, "127.0.0.1", 8888, loop=self.loop)
        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()