tests.py

import utime as time
from ubinascii import hexlify
from umqtt.simple2 import MQTTClient as _MQTTClient, MQTTException, pid_gen


def debug_print(data):
    print('HEX: %s STR: /*' % hexlify(data).decode('ascii'), end='')
    for i, d in enumerate(data):
        if type(d) == str:
            d = ord(d)
        if d > 31 and d < 127:
            print(chr(d), end='')
        else:
            print('.', end='')
    print('*/')


def debug_func_name(f_name):
    print('FUNC: %s' % f_name)


class MQTTClient(_MQTTClient):
    MAX_DBG_LEN = 80

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.newpid = pid_gen(65535 - 1)

    def _read(self, n):
        out = super()._read(n)
        if type(out) == bytes:
            print(' R(%3d) - ' % len(out), end='')
            debug_print(out[:self.MAX_DBG_LEN])
        else:
            print(' R(---) - %s' % out)
        return out

    def _write(self, bytes_wr, *args, **kwargs):
        print(' W(%3d,%3s) - ' % (len(bytes_wr), args[0] if args else '---'), end='')
        debug_print(bytes_wr[:self.MAX_DBG_LEN])
        return super()._write(bytes_wr, *args, **kwargs)

    def connect(self, *a, **k):
        debug_func_name('connect')
        return super().connect(*a, **k)

    def disconnect(self, *a, **k):
        debug_func_name('disconnect')
        return super().disconnect(*a, **k)

    def ping(self, *a, **k):
        debug_func_name('ping')
        return super().ping(*a, **k)

    def publish(self, *a, **k):
        debug_func_name('publish')
        return super().publish(*a, **k)

    def subscribe(self, *a, **k):
        debug_func_name('subscribe')
        return super().subscribe(*a, **k)

    def wait_msg(self, *a, **k):
        debug_func_name('wait_msg')
        return super().wait_msg(*a, **k)


class TestMQTT:
    HIDE_SENSITIVE = True

    TESTS = [
        'test_publish_qos_0',
        'test_subscribe_qos_0',
        # 'test_subscribe_qos_1', # not implenemted
        # 'test_subscribe_qos_2', # not implenemted
        'test_subscribe_long_topic',
        'test_publish_qos_1',
        'test_publish_qos_1_no_puback',
        'test_publish_qos_2',
        'test_publish_retain',
        'test_publish_lastwill',
    ]

    def __init__(self, *args, **kwargs):
        self.mqtt_client_args = (args, kwargs)
        self.msg_id = args[0]
        self.subsctiption_out = {}
        self.status_out = {}
        self.clients = {}
        self.client = None

    def init_mqtt_client(self, clientid_postfix='_1'):
        args = list(self.mqtt_client_args[0][:])
        kwargs = self.mqtt_client_args[1].copy()
        if len(args) > 0:
            args[0] += clientid_postfix
        if 'client_id' in kwargs:
            kwargs['client_id'] += clientid_postfix
        print_args = args.copy()
        print_kwargs = kwargs.copy()
        if 'ssl_params' in kwargs:
            print_kwargs['ssl_params'] = kwargs['ssl_params'].copy()
        hidden_val = '../hidden/..'
        if self.HIDE_SENSITIVE:
            if 'ssl_params' in print_kwargs:
                if 'key' in print_kwargs['ssl_params']:
                    print_kwargs['ssl_params']['key'] = hidden_val
                if 'cert' in print_kwargs['ssl_params']:
                    print_kwargs['ssl_params']['cert'] = hidden_val
            if 'password' in print_kwargs:
                print_kwargs['password'] = hidden_val
            if len(print_args) >= 5:
                print_args[4] = hidden_val  # password
            if len(print_args) >= 8:
                print_args[7] = hidden_val  # ssl_params
                if 'key' in print_args[7]:
                    print_args[7]['key'] = hidden_val
                if 'cert' in print_args[7]:
                    print_args[7]['cert'] = hidden_val

        print('MQTT connection args:', args, print_kwargs)
        client = MQTTClient(*args, **kwargs)
        client.set_callback(self.sub_cb_gen(clientid_postfix))
        client.set_callback_status(self.stat_cb_gen(clientid_postfix))
        self.subsctiption_out[clientid_postfix] = None
        self.status_out[clientid_postfix] = None
        self.clients[clientid_postfix] = client
        return client

    def sub_cb_gen(self, clientid_postfix='_1'):
        def sub_cb(topic, msg, retained, dup):
            print('TOPIC: %s MSG: %s R: %s D: %s' % (topic, msg, retained, dup))
            self.subsctiption_out[clientid_postfix] = (topic, msg, retained)

        return sub_cb

    def stat_cb_gen(self, clientid_postfix='_1'):
        def stat_cb(pid, status):
            print('PID: %s STATUS: %d' % (pid, status))
            self.status_out[clientid_postfix] = (pid, status)

        return stat_cb

    def get_subscription_out(self, timeout=5, clientid_postfix='_1'):
        print('WAIT SUB: timeout=%d' % (timeout,))
        client = self.clients[clientid_postfix]
        for i in range(timeout):
            client.check_msg()
            if clientid_postfix in self.subsctiption_out and self.subsctiption_out[clientid_postfix] is not None:
                o = self.subsctiption_out[clientid_postfix]
                self.subsctiption_out[clientid_postfix] = None
                return o
            time.sleep(1)
        raise Exception('timeout')

    def get_status_out(self, timeout=5, pid=None, clientid_postfix='_1'):
        print('WAIT STAT: timeout=%d pid=%s' % (timeout, pid))
        client = self.clients[clientid_postfix]
        for i in range(timeout + 1):
            time.sleep(1)
            client.check_msg()
            if clientid_postfix in self.status_out and self.status_out[clientid_postfix] is not None:
                o = self.status_out[clientid_postfix]
                self.status_out[clientid_postfix] = None
                if pid:
                    if pid != o[0]:
                        continue
                return o
        raise Exception('timeout')

    def disable_net(self):
        raise RuntimeError('Not implemented method')

    def enable_net(self):
        raise RuntimeError('Not implemented method')

    def network_status(self):
        # Display networki information
        raise RuntimeError('Not implemented method')

    def device_status(self):
        from micropython import mem_info, qstr_info
        print('DEVICE STATUS:')
        mem_info()
        qstr_info()

    def get_topic(self, test_name):
        return '%s/umqtt.simple2/%s/' % (self.msg_id, test_name)

    def verbose_tests(self, test_values: dict):
        tests_fail=[k for k, v in test_values.items() if not v[0]]
        if tests_fail:
            print('\nTests fails: %s\n' % ', '.join(tests_fail))
        else:
            print('All the tests were finished successfully!')

    def run(self, verbose=True, tests=None, hide_sensitive=HIDE_SENSITIVE):
        self.HIDE_SENSITIVE = hide_sensitive
        if not tests:
            tests = self.TESTS
        tests_values = {}
        for test_name in tests:
            a = time.ticks_ms()
            s = self.run_test(test_name)
            b = time.ticks_ms()
            tests_values[test_name] = (s, time.ticks_diff(b, a))
        if verbose:
            self.verbose_tests(tests_values)
        return tests_values

    def run_test(self, test_name):
        try:
            if self.client:
                self.client.disconnect()
        except:
            pass
        self.client = self.init_mqtt_client()
        self.subsctiption_out = {}
        self.status_out = {}
        test = getattr(self, test_name)
        print('RUN [%s]' % test_name)
        test_pass = True
        self.enable_net()
        try:
            test(self.get_topic(test_name))
        except Exception as e:
            from sys import print_exception
            self.network_status()
            self.device_status()
            print_exception(e)
            test_pass = False
        print('END [%s] %s\n' % (test_name, 'succes' if test_pass else 'FAIL'))
        return test_pass

    def test_publish_qos_0(self, topic):
        self.client.connect()
        self.client.publish(topic, 'test QoS 0')
        self.client.disconnect()

    def test_publish_qos_1(self, topic):
        self.client.connect()
        pid = self.client.publish(topic, 'test QoS 1', qos=1)
        out_pid, status = self.get_status_out(pid=pid)
        assert status == 1
        self.client.disconnect()

    def test_publish_qos_1_no_puback(self, topic):
        self.client.connect()
        pid = self.client.publish(topic, 'test QoS 1', qos=1)
        pid = next(self.client.newpid)
        self.client.rcv_pids[pid] = time.ticks_add(time.ticks_ms(), self.client.message_timeout * 1000)
        out_pid, status = self.get_status_out(10, pid=pid)
        assert status == 0
        self.client.disconnect()

    def test_publish_qos_2(self, topic):
        self.client.connect()
        try:
            self.client.publish(topic, 'test QoS 2', qos=2)
        except AssertionError as e:
            print(e)
        self.client.disconnect()

    def test_publish_retain(self, topic):
        self.client.connect()
        pid = self.client.publish(topic, 'test retain', qos=1, retain=True)
        self.client.disconnect()
        self.client.connect()
        self.client.subscribe(topic + '#')
        t, m, r = self.get_subscription_out()
        assert t.decode('ascii') == topic
        assert m.decode('ascii') == 'test retain'
        assert r == True
        self.client.disconnect()

    def test_subscribe_qos_0(self, topic):
        self.client.connect()
        pid = self.client.subscribe(topic + '#')
        out_pid, status = self.get_status_out(pid=pid)
        assert status == 1
        assert pid == 65535
        msg_in = 'abc123'
        self.client.publish(topic, msg_in, qos=1)
        msg_out = self.get_subscription_out()[1]
        assert msg_in == msg_out.decode('ascii')
        msg_in = ''
        self.client.publish(topic, msg_in, qos=1)
        msg_out = self.get_subscription_out()[1]
        assert msg_in == msg_out.decode('ascii')
        self.client.disconnect()

    def test_subscribe_qos_1(self, topic):
        pass

    def test_subscribe_qos_2(self, topic):
        pass

    def test_subscribe_long_topic(self, topic):
        self.client.connect()
        topic = topic + '3' * (500 - len(topic))
        pid = self.client.subscribe(topic + '/#')
        out_pid, status = self.get_status_out(pid=pid)
        assert status == 1
        msg_in = 'abc123'
        self.client.publish(topic, msg_in, qos=1)
        msg_out = self.get_subscription_out()[1]
        assert msg_in == msg_out.decode('ascii')
        self.client.disconnect()

    def test_publish_lastwill(self, topic):
        self.client.set_last_will(topic, 'offline', True, 1)

        self.client.connect()
        self.client.publish(topic, 'online', retain=True, qos=1)
        self.client.subscribe(topic)
        msg_out = self.get_subscription_out()[1]
        assert b'online' == msg_out
        self.client.sock.close()
        time.sleep(1)

        client_2 = self.init_mqtt_client('_2')
        client_2.connect()
        client_2.subscribe(topic)
        msg_out = self.get_subscription_out(clientid_postfix='_2')[1]
        assert 'offline' == msg_out.decode('ascii')
        client_2.disconnect()