diff --git a/micropython/umqtt.simple/test_umqtt_simple.py b/micropython/umqtt.simple/test_umqtt_simple.py new file mode 100644 index 000000000..c31efd17a --- /dev/null +++ b/micropython/umqtt.simple/test_umqtt_simple.py @@ -0,0 +1,76 @@ +import io +import sys + + +class Socket: + def __init__(self, read_data=b""): + self._write_buffer = io.BytesIO() + self._read_buffer = io.BytesIO(read_data) + + def write(self, buf, length=None): + if length is None: + length = len(buf) + self._write_buffer.write(buf[:length]) + + def read(self, n): + return self._read_buffer.read(n) + + def setblocking(self, blocking): + pass + + def close(self): + pass + + +sys.path.insert(0, "micropython/umqtt.simple") +# ruff: noqa: E402 +from umqtt.simple import MQTTClient + + +def make_client(read_data): + c = MQTTClient(b"cid", "127.0.0.1") + c.sock = Socket(read_data) + c.set_callback(lambda topic, msg: None) + return c + + +def test_subscribe_short_topic(): + # Remaining Length = 5 + 4 = 9 -> single VBI byte 0x09, pid=1 + c = make_client(b"\x90\x03\x00\x01\x00") + c.subscribe(b"abcd", qos=0) + out = c.sock._write_buffer.getvalue() + assert out[:2] == b"\x82\x09", out + assert out[2:4] == b"\x00\x01", out + assert out[4:6] == b"\x00\x04", out + assert out[6:10] == b"abcd", out + assert out[10:11] == b"\x00", out + + +def test_subscribe_long_topic(): + # Remaining Length = 5 + 123 = 128 -> VBI 0x80 0x01 + topic = b"a" * 123 + c = make_client(b"\x90\x03\x00\x01\x00") + c.subscribe(topic, qos=0) + out = c.sock._write_buffer.getvalue() + assert out[:3] == b"\x82\x80\x01", out + assert out[3:5] == b"\x00\x01", out + assert out[5:7] == b"\x00\x7b", out + assert out[7:130] == topic, out + assert out[130:131] == b"\x00", out + + +def test_unsubscribe_long_topic(): + # Remaining Length = 4 + 123 = 127 -> still one VBI byte; use 124 for 128 + topic = b"a" * 124 + c = make_client(b"\xb0\x02\x00\x01") + c.unsubscribe(topic) + out = c.sock._write_buffer.getvalue() + assert out[:3] == b"\xa2\x80\x01", out + assert out[3:5] == b"\x00\x01", out + assert out[5:7] == b"\x00\x7c", out + assert out[7:131] == topic, out + + +test_subscribe_short_topic() +test_subscribe_long_topic() +test_unsubscribe_long_topic() diff --git a/micropython/umqtt.simple/umqtt/simple.py b/micropython/umqtt.simple/umqtt/simple.py index bfdaf503c..7539711e3 100644 --- a/micropython/umqtt.simple/umqtt/simple.py +++ b/micropython/umqtt.simple/umqtt/simple.py @@ -7,6 +7,16 @@ class MQTTException(Exception): pass +def _encode_len(pkt, sz): + i = 1 + while sz > 0x7F: + pkt[i] = (sz & 0x7F) | 0x80 + sz >>= 7 + i += 1 + pkt[i] = sz + return i + + class MQTTClient: def __init__( self, @@ -93,12 +103,7 @@ def connect(self, clean_session=True, timeout=None): msg[6] |= 0x4 | (self.lw_qos & 0x1) << 3 | (self.lw_qos & 0x2) << 3 msg[6] |= self.lw_retain << 5 - i = 1 - while sz > 0x7F: - premsg[i] = (sz & 0x7F) | 0x80 - sz >>= 7 - i += 1 - premsg[i] = sz + i = _encode_len(premsg, sz) self.sock.write(premsg, i + 2) self.sock.write(msg) @@ -130,12 +135,7 @@ def publish(self, topic, msg, retain=False, qos=0): if qos > 0: sz += 2 assert sz < 2097152 - i = 1 - while sz > 0x7F: - pkt[i] = (sz & 0x7F) | 0x80 - sz >>= 7 - i += 1 - pkt[i] = sz + i = _encode_len(pkt, sz) # print(hex(len(pkt)), hexlify(pkt, ":")) self.sock.write(pkt, i + 1) self._send_str(topic) @@ -158,37 +158,33 @@ def publish(self, topic, msg, retain=False, qos=0): elif qos == 2: assert 0 - def subscribe(self, topic, qos=0): - assert self.cb is not None, "Subscribe callback is not set" - pkt = bytearray(b"\x82\0\0\0") + def _send_subunsub(self, topic, typ, ack_op, ack_n, qos=0): + pkt = bytearray(4) + pkt[0] = typ self.pid += 1 - struct.pack_into("!BH", pkt, 1, 2 + 2 + len(topic) + 1, self.pid) - # print(hex(len(pkt)), hexlify(pkt, ":")) - self.sock.write(pkt) + pid = self.pid + i = _encode_len(pkt, 4 + len(topic) + (ack_n > 3)) + self.sock.write(pkt, i + 1) + struct.pack_into("!H", pkt, 0, pid) + self.sock.write(pkt, 2) self._send_str(topic) - self.sock.write(qos.to_bytes(1, "little")) + if ack_n > 3: + self.sock.write(bytes((qos,))) while 1: op = self.wait_msg() - if op == 0x90: - resp = self.sock.read(4) - # print(resp) - assert resp[1] == pkt[2] and resp[2] == pkt[3] - if resp[3] == 0x80: + if op == ack_op: + resp = self.sock.read(ack_n) + assert (resp[1] << 8 | resp[2]) == pid + if ack_n > 3 and resp[3] == 0x80: raise MQTTException(resp[3]) return + def subscribe(self, topic, qos=0): + assert self.cb is not None, "Subscribe callback is not set" + self._send_subunsub(topic, 0x82, 0x90, 4, qos) + def unsubscribe(self, topic): - pkt = bytearray(b"\xa2\0\0\0") - self.pid += 1 - struct.pack_into("!BH", pkt, 1, 2 + 2 + len(topic), self.pid) - self.sock.write(pkt) - self._send_str(topic) - while 1: - op = self.wait_msg() - if op == 0xB0: - resp = self.sock.read(3) - assert resp[1] == pkt[2] and resp[2] == pkt[3] - return + self._send_subunsub(topic, 0xA2, 0xB0, 3) # Wait for a single incoming MQTT message and process it. # Subscribed messages are delivered to a callback previously