multi-agent-mux/tests/_local_mqtt.py

"""Issue #26 - a test-only, in-process MQTT 3.1.1 broker.

The suite used to round-trip through the public default broker
(mqtt_common.BrokerConfig.host). conftest starts one of these per session on
127.0.0.1 and points MQTT_BROKER/MQTT_PORT/MQTT_TLS at it, so every real
producer (publish_event, job_subscriber, the wrapper, the remote runner, the
monitor daemon) talks to a broker the test process owns.

Scope is what the repo uses: paho-mqtt 2.x clients with its defaults
(protocol level 4 = 3.1.1, clean session, no will), QoS 0/1, retained
messages. No TLS, no auth checks, no persistent sessions, no QoS 2, no
redelivery of unacknowledged QoS 1 messages. The name does not start with
`test_`, so pytest never collects it.
"""
from __future__ import annotations

import collections
import socket
import struct
import threading

CONNECT, CONNACK, PUBLISH, PUBACK = 1, 2, 3, 4
SUBSCRIBE, SUBACK, UNSUBSCRIBE, UNSUBACK = 8, 9, 10, 11
PINGREQ, PINGRESP, DISCONNECT = 12, 13, 14
NAMES = {CONNECT: "CONNECT", PUBLISH: "PUBLISH", PUBACK: "PUBACK", SUBSCRIBE: "SUBSCRIBE",
         UNSUBSCRIBE: "UNSUBSCRIBE", PINGREQ: "PINGREQ", DISCONNECT: "DISCONNECT"}

_CURRENT = None


def topic_matches(filt: str, topic: str) -> bool:
    f, t = filt.split("/"), topic.split("/")
    if topic.startswith("$") and f[0] in ("+", "#"):
        return False
    for i, part in enumerate(f):
        if part == "#":
            return True
        if i >= len(t) or (part != "+" and part != t[i]):
            return False
    return len(f) == len(t)


def _varint(n: int) -> bytes:
    out = bytearray()
    while True:
        n, digit = divmod(n, 128)
        out.append(digit | (0x80 if n else 0))
        if not n:
            return bytes(out)


def _packet(first: int, body: bytes) -> bytes:
    return bytes([first]) + _varint(len(body)) + body


def _str(data: bytes, i: int):
    (n,) = struct.unpack_from("!H", data, i)
    return data[i + 2:i + 2 + n], i + 2 + n


class _Conn:
    def __init__(self, sock):
        self.sock = sock
        self.lock = threading.Lock()
        self.subs = {}          # filter -> granted qos
        self.next_pid = 0
        self.client_id = None

    def send(self, data: bytes) -> None:
        with self.lock:
            try:
                self.sock.sendall(data)
            except OSError:
                pass

    def recv_exact(self, n: int) -> bytes:
        buf = b""
        while len(buf) < n:
            chunk = self.sock.recv(n - len(buf))
            if not chunk:
                raise ConnectionError("closed")
            buf += chunk
        return buf


class LocalBroker:
    def __init__(self):
        self._lsock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self._lsock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        self._lsock.bind(("127.0.0.1", 0))
        self._lsock.listen(128)
        self.host, self.port = self._lsock.getsockname()
        self._lock = threading.Lock()
        self._conns = set()
        self.retained = {}                     # topic -> (payload, qos)
        self.connections = 0
        self.protocol_levels = []              # one entry per CONNECT
        self.client_ids = []
        self.packets = collections.Counter()   # inbound packets by name
        self._stopped = False

    # -- lifecycle -----------------------------------------------------------
    def start(self) -> "LocalBroker":
        threading.Thread(target=self._accept_loop, name="local-mqtt-accept", daemon=True).start()
        return self

    def stop(self) -> None:
        self._stopped = True
        try:
            self._lsock.close()
        except OSError:
            pass
        with self._lock:
            conns = list(self._conns)
        for c in conns:
            try:
                c.sock.shutdown(socket.SHUT_RDWR)
                c.sock.close()
            except OSError:
                pass

    def env(self) -> dict:
        return {"MQTT_BROKER": self.host, "MQTT_PORT": str(self.port), "MQTT_TLS": "0"}

    def stats(self) -> dict:
        with self._lock:
            return {"connections": self.connections, "protocol_levels": list(self.protocol_levels),
                    "levels_by_client": dict(zip(self.client_ids, self.protocol_levels)),
                    "packets": dict(self.packets), "open": len(self._conns)}

    # -- network -------------------------------------------------------------
    def _accept_loop(self) -> None:
        while not self._stopped:
            try:
                sock, _ = self._lsock.accept()
            except OSError:
                return
            sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
            threading.Thread(target=self._serve, args=(_Conn(sock),), name="local-mqtt-conn", daemon=True).start()

    def _serve(self, conn: _Conn) -> None:
        with self._lock:
            self._conns.add(conn)
        try:
            while True:
                first = conn.recv_exact(1)[0]
                length, mult = 0, 1
                while True:
                    b = conn.recv_exact(1)[0]
                    length += (b & 0x7F) * mult
                    mult *= 128
                    if not b & 0x80:
                        break
                body = conn.recv_exact(length) if length else b""
                ptype = first >> 4
                with self._lock:
                    self.packets[NAMES.get(ptype, str(ptype))] += 1
                if conn.client_id is None and ptype != CONNECT:
                    return
                if not self._handle(conn, ptype, first & 0x0F, body):
                    return
        except (ConnectionError, OSError, struct.error, IndexError):
            return
        finally:
            with self._lock:
                self._conns.discard(conn)
            try:
                conn.sock.close()
            except OSError:
                pass

    def _handle(self, conn: _Conn, ptype: int, flags: int, body: bytes) -> bool:
        if ptype == CONNECT:
            _name, i = _str(body, 0)
            level, cflags = body[i], body[i + 1]
            client_id, i = _str(body, i + 4)
            with self._lock:
                self.connections += 1
                self.protocol_levels.append(level)
                self.client_ids.append(client_id.decode("utf-8", "replace"))
            if level not in (3, 4):   # 3.1 / 3.1.1 only (no MQTT v5 in the repo)
                conn.send(_packet(CONNACK << 4, b"\x00\x01"))
                return False
            conn.client_id = client_id.decode("utf-8", "replace")
            conn.send(_packet(CONNACK << 4, b"\x00\x00"))
            return True
        if ptype == PUBLISH:
            qos, retain = (flags >> 1) & 3, bool(flags & 1)
            topic, i = _str(body, 0)
            if qos == 2:              # the repo publishes QoS 0/1 only
                return False
            pid = b""
            if qos == 1:
                pid = body[i:i + 2]
                i += 2
            # Route (and store a retained message) BEFORE the PUBACK: a publisher
            # that has its ack may rely on the broker state it asked for.
            self._route(topic.decode("utf-8"), body[i:], qos, retain)
            if qos == 1:
                conn.send(_packet(PUBACK << 4, pid))
            return True
        if ptype == SUBSCRIBE:
            pid, i, granted, new = body[:2], 2, bytearray(), []
            while i < len(body):
                filt, i = _str(body, i)
                q = min(body[i], 1)
                i += 1
                conn.subs[filt.decode("utf-8")] = q
                granted.append(q)
                new.append((filt.decode("utf-8"), q))
            conn.send(_packet(SUBACK << 4, pid + bytes(granted)))
            with self._lock:
                retained = list(self.retained.items())
            for filt, q in new:
                for topic, (payload, rq) in retained:
                    if topic_matches(filt, topic):
                        self._deliver(conn, topic, payload, min(rq, q), True)
            return True
        if ptype == UNSUBSCRIBE:
            pid, i = body[:2], 2
            while i < len(body):
                filt, i = _str(body, i)
                conn.subs.pop(filt.decode("utf-8"), None)
            conn.send(_packet(UNSUBACK << 4, pid))
            return True
        if ptype == PINGREQ:
            conn.send(_packet(PINGRESP << 4, b""))
            return True
        if ptype == PUBACK:           # a subscriber acknowledging our QoS 1 delivery
            return True
        return False                  # DISCONNECT, or anything unsupported

    def _route(self, topic: str, payload: bytes, qos: int, retain: bool) -> None:
        with self._lock:
            if retain:
                if payload:
                    self.retained[topic] = (payload, qos)
                else:
                    self.retained.pop(topic, None)
            conns = list(self._conns)
        for c in conns:
            granted = [q for f, q in list(c.subs.items()) if topic_matches(f, topic)]
            if granted:
                self._deliver(c, topic, payload, min(qos, max(granted)), False)

    def _deliver(self, conn: _Conn, topic: str, payload: bytes, qos: int, retain: bool) -> None:
        t = topic.encode("utf-8")
        body = struct.pack("!H", len(t)) + t
        if qos:
            with conn.lock:
                conn.next_pid = conn.next_pid % 65535 + 1
                pid = conn.next_pid
            body += struct.pack("!H", pid)
        conn.send(_packet((PUBLISH << 4) | (qos << 1) | int(retain), body + payload))


def start_broker() -> LocalBroker:
    global _CURRENT
    _CURRENT = LocalBroker().start()
    return _CURRENT


def stop_broker() -> None:
    global _CURRENT
    if _CURRENT is not None:
        _CURRENT.stop()
        _CURRENT = None


def current_broker() -> LocalBroker:
    assert _CURRENT is not None, "the session's local MQTT broker is not running (conftest)"
    return _CURRENT


def local_broker_env() -> dict:
    """MQTT_BROKER / MQTT_PORT / MQTT_TLS for this session's local broker."""
    return current_broker().env()