From d2454ef8eb08e1b6ccb5e5b880d9d8b62de366e1 Mon Sep 17 00:00:00 2001 From: Tom Kimber Date: Sat, 10 Oct 2026 00:20:53 +0100 Subject: [PATCH] feat(amgi-paho-mqtt): add message send manager --- packages/amgi-paho-mqtt/README.md | 21 + packages/amgi-paho-mqtt/pyproject.toml | 3 +- .../src/amgi_paho_mqtt/__init__.py | 642 ++++++++++++-- .../tests_amgi_paho_mqtt/test_mqtt_message.py | 811 ++++++++++++++++++ .../test_mqtt_message_integration.py | 146 ++++ uv.lock | 2 + 6 files changed, 1535 insertions(+), 90 deletions(-) create mode 100644 packages/amgi-paho-mqtt/tests_amgi_paho_mqtt/test_mqtt_message.py diff --git a/packages/amgi-paho-mqtt/README.md b/packages/amgi-paho-mqtt/README.md index ffe56403..c978cd22 100644 --- a/packages/amgi-paho-mqtt/README.md +++ b/packages/amgi-paho-mqtt/README.md @@ -43,6 +43,27 @@ Or the application could be run via the commandline: asyncfast run amgi-paho-mqtt main:app order-topic ``` +## Sending Messages + +Messages can be sent from the application via the `message.send` event. The `mqtt` binding of the event controls the +send: + +```python +await send( + { + "type": "message.send", + "address": "order-topic", + "headers": [(b"trace-id", b"abc123")], + "payload": b"payload", + "bindings": {"mqtt": {"qos": 1, "retain": False}}, + } +) +``` + +`qos` sets the MQTT QoS level (defaults to `0`), and `retain` marks the message as retained (defaults to `False`). With +MQTT v5, message headers are sent as user properties, with `content-type` mapped to the MQTT content type property. +MQTT v3.1.1 has no way to represent headers, so they are dropped. + ## Contact For questions or suggestions, please contact [jack.burridge@mail.com](mailto:jack.burridge@mail.com). diff --git a/packages/amgi-paho-mqtt/pyproject.toml b/packages/amgi-paho-mqtt/pyproject.toml index 67856061..0cf28eaf 100644 --- a/packages/amgi-paho-mqtt/pyproject.toml +++ b/packages/amgi-paho-mqtt/pyproject.toml @@ -29,8 +29,9 @@ dependencies = [ "amgi-common==0.48.0", "amgi-types==0.48.0", "paho-mqtt>=2.1.0", + "typing-extensions>=4.15.0; python_full_version<'3.11'", ] -entry-points.amgi_server.amgi-paho-mqtt = "amgi_paho_mqtt:run" +entry-points.amgi_server.amgi-paho-mqtt = "amgi_paho_mqtt:_run_cli" [dependency-groups] dev = [ diff --git a/packages/amgi-paho-mqtt/src/amgi_paho_mqtt/__init__.py b/packages/amgi-paho-mqtt/src/amgi_paho_mqtt/__init__.py index 34b7274e..a0c0a6f9 100644 --- a/packages/amgi-paho-mqtt/src/amgi_paho_mqtt/__init__.py +++ b/packages/amgi-paho-mqtt/src/amgi_paho_mqtt/__init__.py @@ -1,12 +1,23 @@ import asyncio +import logging +import sys +import threading from asyncio import AbstractEventLoop from asyncio import Event from asyncio import Future from asyncio import Task +from collections.abc import Awaitable +from collections.abc import Callable +from collections.abc import Sequence +from contextlib import suppress from functools import cached_property +from select import select from socket import SO_SNDBUF +from socket import socketpair from socket import SOL_SOCKET +from types import TracebackType from typing import Any +from typing import AsyncContextManager from typing import TYPE_CHECKING from weakref import WeakValueDictionary @@ -16,20 +27,36 @@ from amgi_types import AMGIReceiveEvent from amgi_types import AMGISendEvent from amgi_types import MessageScope +from amgi_types import MessageSendEvent from paho.mqtt.client import Client from paho.mqtt.client import ConnectFlags from paho.mqtt.client import DisconnectFlags +from paho.mqtt.client import error_string from paho.mqtt.client import MQTT_ERR_SUCCESS from paho.mqtt.client import MQTTMessage from paho.mqtt.client import MQTTv311 +from paho.mqtt.client import MQTTv5 from paho.mqtt.enums import CallbackAPIVersion from paho.mqtt.enums import MQTTProtocolVersion +from paho.mqtt.packettypes import PacketTypes from paho.mqtt.properties import Properties from paho.mqtt.reasoncodes import ReasonCode +if sys.version_info >= (3, 11): # pragma: no cover + from typing import Self +else: + from typing_extensions import Self + if TYPE_CHECKING: from paho.mqtt.client import SocketLike +_MessageSendT = Callable[[MessageSendEvent], Awaitable[None]] +_MessageSendManagerT = AsyncContextManager[_MessageSendT] + +logger = logging.getLogger("amgi-paho-mqtt.error") + +_RECONNECT_DELAY = 1.0 + def run( app: AMGIApplication, @@ -37,11 +64,31 @@ def run( host: str = "localhost", port: int = 1883, client_id: str | None = None, + message_send: _MessageSendManagerT | None = None, + connect_timeout: float = 10.0, ) -> None: - server = Server(app, topic, host, port, client_id) + server = Server( + app, + topic, + host, + port, + client_id, + message_send=message_send, + connect_timeout=connect_timeout, + ) server_serve(server) +def _run_cli( + app: AMGIApplication, + topic: str, + host: str = "localhost", + port: int = 1883, + client_id: str | None = None, +) -> None: + run(app, topic, host=host, port=port, client_id=client_id) + + class PublishError(OSError): """Raised when publishing fails.""" @@ -50,6 +97,484 @@ async def _receive() -> AMGIReceiveEvent: raise RuntimeError("Receive should not be called") +def _create_properties(headers: Sequence[tuple[bytes, bytes]]) -> Properties: + # paho-mqtt does not provide type annotations for the Properties constructor. + properties_factory: Any = Properties + properties: Properties = properties_factory(PacketTypes.PUBLISH) + for key, value in headers: + if key.lower() == b"content-type": + properties.ContentType = value.decode() + else: + properties.UserProperty = (key.decode(), value.decode()) + return properties + + +class _SocketWatcher: + """ + Watches a socket for readiness from a background thread, invoking the read, + and write callbacks on the event loop. This is used when the event loop does + not support add_reader, and add_writer, such as the proactor event loop on + Windows. As the callbacks are invoked on the event loop, paho is only ever + driven from the event loop thread. + """ + + def __init__( + self, + loop: AbstractEventLoop, + socket: "SocketLike", + on_read: Callable[[], Any], + ) -> None: + self._loop = loop + self._socket = socket + self._on_read = on_read + self._on_write: Callable[[], Any] | None = None + self._closed = False + # Each readiness is handled by the event loop before the socket is selected + # again, as select is level triggered. + self._handled = threading.Event() + self._waker_receive, self._waker_send = socketpair() + self._waker_receive.setblocking(False) + self._waker_send.setblocking(False) + self._thread = threading.Thread( + target=self._run, name="amgi-paho-mqtt-socket-watcher", daemon=True + ) + self._thread.start() + + def set_writer(self, on_write: Callable[[], Any] | None) -> None: + self._on_write = on_write + self._wake() + + def close(self) -> None: + self._closed = True + self._handled.set() + self._wake() + + def _wake(self) -> None: + with suppress(OSError): + self._waker_send.send(b"\0") + + def _run(self) -> None: + try: + while not self._closed: + writers = [self._socket] if self._on_write is not None else [] + try: + readable, writable, _ = select( + [self._socket, self._waker_receive], writers, [] + ) + except (OSError, ValueError): + # The socket was closed by paho. + return + if self._waker_receive in readable: + with suppress(OSError): + self._waker_receive.recv(4096) + read = self._socket in readable + write = self._socket in writable + if self._closed or not (read or write): + continue + self._handled.clear() + try: + self._loop.call_soon_threadsafe(self._dispatch, read, write) + except RuntimeError: + # The event loop is closed. + return + self._handled.wait() + finally: + self._waker_receive.close() + self._waker_send.close() + + def _dispatch(self, read: bool, write: bool) -> None: + try: + if self._closed: + return + if read: + self._on_read() + on_write = self._on_write + if write and on_write is not None and not self._closed: + on_write() + finally: + self._handled.set() + + +class _ClientLoop: + """ + Drives a paho client from the running event loop, so the network loop, and + publish acknowledgements are handled asynchronously, rather than by a paho + managed thread. + """ + + def __init__(self, client: Client) -> None: + self._client = client + self._disconnected_event = Event() + self._publish_futures = WeakValueDictionary[int, Future[None]]() + self._misc_task: Task[None] | None = None + self._socket_watcher: _SocketWatcher | None = None + self._reconnect_task: Task[None] | None = None + self._closed = False + + client.on_socket_open = self._on_socket_open + client.on_socket_close = self._on_socket_close + client.on_socket_register_write = self._on_socket_register_write + client.on_socket_unregister_write = self._on_socket_unregister_write + client.on_publish = self._on_publish + client.on_disconnect = self._on_disconnect + + @cached_property + def loop(self) -> AbstractEventLoop: + return asyncio.get_running_loop() + + def _on_socket_open( + self, client: Client, userdata: Any, socket: "SocketLike" + ) -> None: + def _register() -> None: + if self._closed: + return + loop = self.loop + try: + loop.add_reader(socket, client.loop_read) + except NotImplementedError: + self._socket_watcher = _SocketWatcher(loop, socket, client.loop_read) + self._misc_task = loop.create_task(self._misc_loop(client)) + + self._call_on_loop(_register) + + def _on_socket_close( + self, client: Client, userdata: Any, socket: "SocketLike" + ) -> None: + misc_task = self._misc_task + self._misc_task = None + + def _unregister() -> None: + socket_watcher = self._socket_watcher + if socket_watcher is not None: + self._socket_watcher = None + socket_watcher.close() + else: + self.loop.remove_reader(socket) + if misc_task is not None and not misc_task.done(): + misc_task.cancel() + + self._call_on_loop(_unregister) + + def _on_socket_register_write( + self, client: Client, userdata: Any, socket: "SocketLike" + ) -> None: + def _register_write() -> None: + socket_watcher = self._socket_watcher + if socket_watcher is not None: + socket_watcher.set_writer(client.loop_write) + else: + self.loop.add_writer(socket, client.loop_write) + + self._call_on_loop(_register_write) + + def _on_socket_unregister_write( + self, client: Client, userdata: Any, socket: "SocketLike" + ) -> None: + def _unregister_write() -> None: + socket_watcher = self._socket_watcher + if socket_watcher is not None: + socket_watcher.set_writer(None) + else: + self.loop.remove_writer(socket) + + self._call_on_loop(_unregister_write) + + def _call_on_loop(self, callback: Callable[[], None]) -> None: + """ + Invoke a callback on the event loop. paho may invoke the socket + callbacks from a worker thread, when a blocking connect is offloaded, + in which case the callback is marshalled back onto the event loop. + """ + try: + asyncio.get_running_loop() + except RuntimeError: + self.loop.call_soon_threadsafe(callback) + else: + callback() + + def _on_publish( + self, + client: Client, + userdata: Any, + mid: int, + reason_code: ReasonCode, + properties: Properties, + ) -> None: + message_future = self._publish_futures.get(mid) + + if message_future is not None and not message_future.done(): + if reason_code.is_failure: + message_future.set_exception(PublishError(reason_code.getName())) + else: + message_future.set_result(None) + + def _on_disconnect( + self, + client: Client, + userdata: Any, + disconnect_flags: DisconnectFlags, + reason_code: ReasonCode, + properties: Properties | None, + ) -> None: + self._disconnected_event.set() + error = PublishError(str(reason_code)) + for message_future in tuple(self._publish_futures.values()): + if not message_future.done(): + message_future.set_exception(error) + + if reason_code.value != 0: + self._schedule_reconnect() + + async def publish( + self, + topic: str, + payload: bytes | None, + qos: int, + retain: bool = False, + properties: Properties | None = None, + ) -> None: + if properties is None: + mqtt_message_info = self._client.publish( + topic, payload, qos=qos, retain=retain + ) + else: + mqtt_message_info = self._client.publish( + topic, payload, qos=qos, retain=retain, properties=properties + ) + + if mqtt_message_info.rc != MQTT_ERR_SUCCESS: + raise PublishError( + f"Failed to publish message: {error_string(mqtt_message_info.rc)}" + ) + + message_future = self.loop.create_future() + self._publish_futures[mqtt_message_info.mid] = message_future + await message_future + + async def wait_disconnected(self) -> None: + await self._disconnected_event.wait() + + async def close(self) -> None: + """ + Stop driving the client: any in-flight reconnection is cancelled, and no + further reconnections are scheduled. + """ + self._closed = True + + reconnect_task = self._reconnect_task + self._reconnect_task = None + if reconnect_task is not None: + reconnect_task.cancel() + with suppress(asyncio.CancelledError): + await reconnect_task + + misc_task = self._misc_task + if misc_task is not None and not misc_task.done(): + misc_task.cancel() + + def _schedule_reconnect(self) -> None: + if self._closed: + return + reconnect_task = self._reconnect_task + if reconnect_task is not None and not reconnect_task.done(): + return + self._reconnect_task = self.loop.create_task(self._reconnect()) + + async def _reconnect(self) -> None: + while True: + await asyncio.sleep(_RECONNECT_DELAY) + worker = asyncio.ensure_future(asyncio.to_thread(self._client.reconnect)) + try: + rc = await asyncio.shield(worker) + except asyncio.CancelledError: + # The worker thread cannot be cancelled, so a connection it goes on + # to establish must be closed once it finishes. + worker.add_done_callback(self._disconnect_late_reconnect) + raise + except OSError as error: + logger.warning("MQTT reconnection failed: %s", error) + continue + if rc != MQTT_ERR_SUCCESS: + logger.warning("MQTT reconnection failed: %s", error_string(rc)) + continue + return + + def _disconnect_late_reconnect(self, worker: "Future[int]") -> None: + if worker.cancelled() or worker.exception() is not None: + return + if worker.result() == MQTT_ERR_SUCCESS: + self._client.disconnect() + + async def _misc_loop(self, client: Client) -> None: + while client.loop_misc() == MQTT_ERR_SUCCESS: + try: + await asyncio.sleep(1) + except asyncio.CancelledError: + break + + +class MessageSend: + def __init__( + self, + host: str = "localhost", + port: int = 1883, + protocol: MQTTProtocolVersion = MQTTv311, + connect_timeout: float = 10.0, + ) -> None: + self._host = host + self._port = port + self._protocol = protocol + self._connect_timeout = connect_timeout + self._client: Client | None = None + self._client_loop: _ClientLoop | None = None + self._connected_event = Event() + self._connect_error: ConnectionError | None = None + self._late_cleanups: set[Task[None]] = set() + + def _on_connect( + self, + client: Client, + userdata: Any, + connect_flags: ConnectFlags, + reason_code: ReasonCode, + properties: Properties | None, + ) -> None: + if reason_code.is_failure: + self._connect_error = ConnectionError( + f"MQTT connection failed: {reason_code}" + ) + else: + self._connect_error = None + self._connected_event.set() + + async def __aenter__(self) -> Self: + """ + Connect to the broker, and wait for the connection to be acknowledged. + + A failed or timed out connection raises, leaving the manager free to be + entered again to start a new connection. + """ + self._connected_event.clear() + self._connect_error = None + + client = Client(CallbackAPIVersion.VERSION2, protocol=self._protocol) + client.on_connect = self._on_connect + client_loop = _ClientLoop(client) + self._client_loop = client_loop + # Resolve the event loop before the blocking connect is offloaded to a thread. + client_loop.loop + + worker = asyncio.ensure_future( + asyncio.to_thread(client.connect, self._host, self._port, 60) + ) + try: + await asyncio.wait_for(asyncio.shield(worker), self._connect_timeout) + except asyncio.TimeoutError: + self._cleanup_after_worker(worker, client, client_loop) + await self._abort_connect(client, client_loop) + raise TimeoutError("Timed out connecting to MQTT broker") from None + except asyncio.CancelledError: + self._cleanup_after_worker(worker, client, client_loop) + await self._abort_connect(client, client_loop) + raise + except OSError: + await self._abort_connect(client, client_loop) + raise + + self._client = client + try: + await asyncio.wait_for(self._connected_event.wait(), self._connect_timeout) + except asyncio.TimeoutError: + await self._abort_connect(client, client_loop) + raise TimeoutError("Timed out waiting for MQTT connection") from None + except asyncio.CancelledError: + await self._abort_connect(client, client_loop) + raise + + connect_error = self._connect_error + if connect_error is not None: + await self._abort_connect(client, client_loop) + raise connect_error + + return self + + async def __call__(self, event: MessageSendEvent) -> None: + if self._client is None or self._client_loop is None: + raise RuntimeError("MessageSend not initialized") + + mqtt_bindings = event.get("bindings", {}).get("mqtt", {}) + qos = mqtt_bindings.get("qos", 0) + retain = mqtt_bindings.get("retain", False) + + properties: Properties | None = None + headers = event.get("headers") + if self._protocol == MQTTv5 and headers: + # MQTT v3.1.1 has no way to represent headers, so they are dropped. + properties = _create_properties(headers) + + await self._client_loop.publish( + event["address"], + event.get("payload"), + qos, + retain=retain, + properties=properties, + ) + + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc_val: BaseException | None, + exc_tb: TracebackType | None, + ) -> None: + if self._client is None or self._client_loop is None: + return + + client = self._client + client_loop = self._client_loop + await client_loop.close() + client.disconnect() + await client_loop.wait_disconnected() + + self._client = None + self._client_loop = None + + async def _abort_connect(self, client: Client, client_loop: _ClientLoop) -> None: + self._client = None + self._client_loop = None + await client_loop.close() + client.disconnect() + + def _cleanup_after_worker( + self, worker: "Future[None]", client: Client, client_loop: _ClientLoop + ) -> None: + """ + The blocking connect cannot be cancelled, so a connection it goes on to + establish after being abandoned is closed once the worker finishes. + """ + + def _cleanup(_: "Future[None]") -> None: + if worker.cancelled() or worker.exception() is not None: + return + task = client_loop.loop.create_task(self._discard_late(client, client_loop)) + self._late_cleanups.add(task) + task.add_done_callback(self._late_cleanups.discard) + + worker.add_done_callback(_cleanup) + + async def _discard_late(self, client: Client, client_loop: _ClientLoop) -> None: + await client_loop.close() + client.disconnect() + + +class _Send: + def __init__(self, message_send: _MessageSendT) -> None: + self._message_send = message_send + + async def __call__(self, event: AMGISendEvent) -> None: + if event["type"] == "message.send": + await self._message_send(event) + + class Server: def __init__( self, @@ -59,6 +584,8 @@ def __init__( port: int, client_id: str | None, protocol: MQTTProtocolVersion = MQTTv311, + message_send: _MessageSendManagerT | None = None, + connect_timeout: float = 10.0, ) -> None: self._app = app self._topic = topic @@ -68,27 +595,21 @@ def __init__( self._client = Client( CallbackAPIVersion.VERSION2, client_id=client_id, protocol=protocol ) + self._client_loop = _ClientLoop(self._client) self._client.on_connect = self._on_connect self._client.on_message = self._on_message - self._client.on_disconnect = self._on_disconnect - self._client.on_socket_open = self._on_socket_open - self._client.on_socket_close = self._on_socket_close - self._client.on_socket_register_write = self._on_socket_register_write - self._client.on_socket_unregister_write = self._on_socket_unregister_write self._client.on_subscribe = self._on_subscribe - self._client.on_publish = self._on_publish + + self._message_send_context = message_send or MessageSend( + host, port, protocol=protocol, connect_timeout=connect_timeout + ) + self._message_send: _MessageSendT | None = None self._subscribe_event = Event() - self._disconnected_event = Event() self._stop_event = Event() self._tasks = set[Task[None]]() - self._publish_futures = WeakValueDictionary[int, Future[None]]() self._state: dict[str, Any] = {} - @cached_property - def _loop(self) -> AbstractEventLoop: - return asyncio.get_running_loop() - def _on_connect( self, client: Client, @@ -100,11 +621,13 @@ def _on_connect( client.subscribe(self._topic) def _on_message(self, client: Client, userdata: Any, message: MQTTMessage) -> None: - task = self._loop.create_task(self._handle_message(message)) + task = self._client_loop.loop.create_task(self._handle_message(message)) self._tasks.add(task) task.add_done_callback(self._tasks.discard) async def _handle_message(self, message: MQTTMessage) -> None: + assert self._message_send is not None + scope: MessageScope = { "type": "message", "amgi": {"version": "2.0", "spec_version": "2.0"}, @@ -113,39 +636,7 @@ async def _handle_message(self, message: MQTTMessage) -> None: "payload": message.payload, "state": self._state.copy(), } - await self._app(scope, _receive, self._send) - - def _on_disconnect( - self, - client: Client, - userdata: Any, - disconnect_flags: DisconnectFlags, - reason_code: ReasonCode, - properties: Properties | None, - ) -> None: - self._disconnected_event.set() - - def _on_socket_open( - self, client: Client, userdata: Any, socket: "SocketLike" - ) -> None: - self._loop.add_reader(socket, client.loop_read) - self._misc_task = self._loop.create_task(self._misc_loop()) - - def _on_socket_close( - self, client: Client, userdata: Any, socket: "SocketLike" - ) -> None: - self._loop.remove_reader(socket) - self._misc_task.cancel() - - def _on_socket_register_write( - self, client: Client, userdata: Any, socket: "SocketLike" - ) -> None: - self._loop.add_writer(socket, client.loop_write) - - def _on_socket_unregister_write( - self, client: Client, userdata: Any, socket: "SocketLike" - ) -> None: - self._loop.remove_writer(socket) + await self._app(scope, _receive, _Send(self._message_send)) def _on_subscribe( self, @@ -157,51 +648,24 @@ def _on_subscribe( ) -> None: self._subscribe_event.set() - def _on_publish( - self, - client: Client, - userdata: Any, - mid: int, - reason_code: ReasonCode, - properties: Properties, - ) -> None: - message_future = self._publish_futures.get(mid) - - if message_future is not None: - if reason_code.is_failure: - message_future.set_exception(PublishError(reason_code.getName())) - else: - message_future.set_result(None) - - async def _send(self, message: AMGISendEvent) -> None: - if message["type"] == "message.send": - qos = message.get("bindings", {}).get("mqtt", {}).get("qos", 0) - mqtt_message_info = self._client.publish( - message["address"], message["payload"], qos=qos - ) - message_future = self._loop.create_future() - self._publish_futures[mqtt_message_info.mid] = message_future - await message_future - - async def _misc_loop(self) -> None: - while self._client.loop_misc() == MQTT_ERR_SUCCESS: - try: - await asyncio.sleep(1) - except asyncio.CancelledError: - break - async def serve(self) -> None: - self._client.connect(self._host, self._port, 60) - self._client.socket().setsockopt(SOL_SOCKET, SO_SNDBUF, 2048) + async with self._message_send_context as message_send: + self._message_send = message_send - await self._subscribe_event.wait() + # Resolve the event loop before the blocking connect is offloaded to a thread. + self._client_loop.loop + await asyncio.to_thread(self._client.connect, self._host, self._port, 60) + self._client.socket().setsockopt(SOL_SOCKET, SO_SNDBUF, 2048) - async with Lifespan(self._app, self._state): - await self._stop_event.wait() - self._client.unsubscribe(self._topic) - await asyncio.gather(*self._tasks) - self._client.disconnect() - await self._disconnected_event.wait() + await self._subscribe_event.wait() + + async with Lifespan(self._app, self._state): + await self._stop_event.wait() + self._client.unsubscribe(self._topic) + await asyncio.gather(*self._tasks) + await self._client_loop.close() + self._client.disconnect() + await self._client_loop.wait_disconnected() def stop(self) -> None: self._stop_event.set() diff --git a/packages/amgi-paho-mqtt/tests_amgi_paho_mqtt/test_mqtt_message.py b/packages/amgi-paho-mqtt/tests_amgi_paho_mqtt/test_mqtt_message.py new file mode 100644 index 00000000..e98f17f3 --- /dev/null +++ b/packages/amgi-paho-mqtt/tests_amgi_paho_mqtt/test_mqtt_message.py @@ -0,0 +1,811 @@ +import asyncio +import socket +import threading +from collections.abc import Callable +from collections.abc import Coroutine +from typing import Any +from unittest.mock import AsyncMock +from unittest.mock import Mock + +import amgi_paho_mqtt +import pytest +from amgi_paho_mqtt import _ClientLoop +from amgi_paho_mqtt import _Send +from amgi_paho_mqtt import _SocketWatcher +from amgi_paho_mqtt import MessageSend +from amgi_paho_mqtt import PublishError +from amgi_types import MessageAckEvent +from amgi_types import MessageSendEvent +from paho.mqtt.client import Client +from paho.mqtt.client import ConnectFlags +from paho.mqtt.client import DisconnectFlags +from paho.mqtt.client import MQTT_ERR_NO_CONN +from paho.mqtt.client import MQTT_ERR_SUCCESS +from paho.mqtt.client import MQTTMessageInfo +from paho.mqtt.client import MQTTv311 +from paho.mqtt.client import MQTTv5 +from paho.mqtt.enums import CallbackAPIVersion +from paho.mqtt.enums import MQTTProtocolVersion +from paho.mqtt.packettypes import PacketTypes +from paho.mqtt.properties import Properties +from paho.mqtt.reasoncodes import ReasonCode + + +def _patch_client( + monkeypatch: pytest.MonkeyPatch, connect: Callable[..., Any] | None = None +) -> list[Mock]: + """Stub the paho client constructed by MessageSend, returning the clients.""" + clients: list[Mock] = [] + + def _create_client(*args: Any, **kwargs: Any) -> Mock: + client = Mock(spec=Client) + if connect is not None: + client.connect.side_effect = connect + clients.append(client) + return client + + monkeypatch.setattr(amgi_paho_mqtt, "Client", _create_client) + return clients + + +def _run_on_selector_loop(coroutine: Coroutine[Any, Any, None]) -> None: + """ + Run a coroutine on a selector event loop, as the default proactor event loop + on Windows does not support add_reader, or add_writer. + """ + loop = asyncio.SelectorEventLoop() + try: + loop.run_until_complete(coroutine) + finally: + loop.close() + + +def _create_mock_message_send( + protocol: MQTTProtocolVersion = MQTTv311, +) -> tuple[MessageSend, Mock]: + client = Mock(spec=Client) + client.publish.return_value = MQTTMessageInfo(1) + client_loop = _ClientLoop(client) + message_send = MessageSend(protocol=protocol) + message_send._client = client + message_send._client_loop = client_loop + return message_send, client + + +async def _send( + message_send: MessageSend, client: Mock, event: MessageSendEvent +) -> None: + """Send the event, acknowledging the publish like the broker would.""" + publish_task = asyncio.create_task(message_send(event)) + client_loop = message_send._client_loop + assert client_loop is not None + await asyncio.sleep(0) + client_loop._on_publish( + client, None, 1, ReasonCode(PacketTypes.PUBACK), Mock(spec=Properties) + ) + await asyncio.wait_for(publish_task, 1) + + +async def test_send_message_send() -> None: + message_send = AsyncMock() + event: MessageSendEvent = { + "type": "message.send", + "address": "send-topic", + "headers": [], + "payload": b"test", + } + + await _Send(message_send)(event) + + message_send.assert_awaited_once_with(event) + + +async def test_send_ignores_other_events() -> None: + message_send = AsyncMock() + event: MessageAckEvent = {"type": "message.ack"} + + await _Send(message_send)(event) + + message_send.assert_not_awaited() + + +async def test_message_send_requires_context_manager() -> None: + message_send = MessageSend() + event: MessageSendEvent = { + "type": "message.send", + "address": "send-topic", + "headers": [], + } + + with pytest.raises(RuntimeError, match="MessageSend not initialized"): + await message_send(event) + + +@pytest.mark.parametrize("missing", ["client", "client_loop"]) +async def test_message_send_exit_returns_when_not_initialized(missing: str) -> None: + message_send = MessageSend() + if missing == "client": + message_send._client_loop = Mock() + else: + message_send._client = Mock() + + await message_send.__aexit__(None, None, None) + + +@pytest.mark.parametrize("qos", [1, 2]) +async def test_disconnect_fails_all_pending_publishes(qos: int) -> None: + client = Mock(spec=Client) + client.publish.side_effect = [MQTTMessageInfo(1), MQTTMessageInfo(2)] + client_loop = _ClientLoop(client) + tasks = [ + asyncio.create_task(client_loop.publish("send-topic", b"test", qos)) + for _ in range(2) + ] + await asyncio.sleep(0) + + reason_code = ReasonCode(PacketTypes.DISCONNECT, "Unspecified error") + client_loop._on_disconnect(client, None, DisconnectFlags(False), reason_code, None) + # A repeated disconnect must not try to fail completed futures again. + client_loop._on_disconnect(client, None, DisconnectFlags(False), reason_code, None) + await asyncio.wait_for(client_loop.wait_disconnected(), timeout=1) + + # A late acknowledgement may arrive before the awaiting tasks resume. + client_loop._on_publish( + client, None, 1, ReasonCode(PacketTypes.PUBACK), Mock(spec=Properties) + ) + + results = await asyncio.wait_for( + asyncio.gather(*tasks, return_exceptions=True), timeout=1 + ) + assert all(isinstance(result, PublishError) for result in results) + assert all(str(result) == "Unspecified error" for result in results) + + # The unexpected disconnect schedules a reconnection, which is stopped on close. + await client_loop.close() + + +async def test_disconnect_preserves_completed_and_cancelled_publishes() -> None: + client = Client(CallbackAPIVersion.VERSION2) + client_loop = _ClientLoop(client) + completed = asyncio.get_running_loop().create_future() + completed.set_result(None) + cancelled = asyncio.get_running_loop().create_future() + cancelled.cancel() + client_loop._publish_futures[1] = completed + client_loop._publish_futures[2] = cancelled + + client_loop._on_disconnect( + client, None, DisconnectFlags(False), ReasonCode(PacketTypes.DISCONNECT), None + ) + + assert completed.result() is None + assert cancelled.cancelled() + assert client_loop._reconnect_task is None + + +@pytest.mark.parametrize("qos", [0, 1]) +async def test_publish_while_disconnected_raises(qos: int) -> None: + client = Client(CallbackAPIVersion.VERSION2) + client_loop = _ClientLoop(client) + + with pytest.raises(PublishError, match="not currently connected"): + await asyncio.wait_for(client_loop.publish("send-topic", b"test", qos), 1) + + +async def test_publish_no_conn_does_not_register_future() -> None: + client = Mock(spec=Client) + mqtt_message_info = MQTTMessageInfo(1) + mqtt_message_info.rc = MQTT_ERR_NO_CONN + client.publish.return_value = mqtt_message_info + client_loop = _ClientLoop(client) + + with pytest.raises(PublishError): + await client_loop.publish("send-topic", b"test", 1) + + assert not client_loop._publish_futures + + +async def test_unexpected_disconnect_schedules_single_reconnect( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(amgi_paho_mqtt, "_RECONNECT_DELAY", 0.0) + client = Mock(spec=Client) + client.reconnect.return_value = MQTT_ERR_SUCCESS + client_loop = _ClientLoop(client) + + reason_code = ReasonCode(PacketTypes.DISCONNECT, "Unspecified error") + client_loop._on_disconnect(client, None, DisconnectFlags(False), reason_code, None) + # A second disconnect while reconnecting must not spawn a second task. + client_loop._on_disconnect(client, None, DisconnectFlags(False), reason_code, None) + + reconnect_task = client_loop._reconnect_task + assert reconnect_task is not None + await asyncio.wait_for(reconnect_task, 1) + + client.reconnect.assert_called_once() + + +async def test_reconnect_restores_publish( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(amgi_paho_mqtt, "_RECONNECT_DELAY", 0.0) + client = Mock(spec=Client) + client.reconnect.return_value = MQTT_ERR_SUCCESS + client.publish.return_value = MQTTMessageInfo(1) + client_loop = _ClientLoop(client) + + client_loop._on_disconnect( + client, + None, + DisconnectFlags(False), + ReasonCode(PacketTypes.DISCONNECT, "Unspecified error"), + None, + ) + reconnect_task = client_loop._reconnect_task + assert reconnect_task is not None + await asyncio.wait_for(reconnect_task, 1) + + publish_task = asyncio.create_task(client_loop.publish("send-topic", b"test", 1)) + await asyncio.sleep(0) + client_loop._on_publish( + client, None, 1, ReasonCode(PacketTypes.PUBACK), Mock(spec=Properties) + ) + + await asyncio.wait_for(publish_task, 1) + client.reconnect.assert_called_once() + + +async def test_reconnect_retries_until_close(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(amgi_paho_mqtt, "_RECONNECT_DELAY", 0.0) + client = Mock(spec=Client) + client.reconnect.side_effect = ConnectionRefusedError() + client_loop = _ClientLoop(client) + + client_loop._on_disconnect( + client, + None, + DisconnectFlags(False), + ReasonCode(PacketTypes.DISCONNECT, "Unspecified error"), + None, + ) + reconnect_task = client_loop._reconnect_task + assert reconnect_task is not None + + while client.reconnect.call_count < 2: + await asyncio.sleep(0) + + await client_loop.close() + assert reconnect_task.cancelled() + + # After close, no further reconnections are scheduled. + client_loop._on_disconnect( + client, + None, + DisconnectFlags(False), + ReasonCode(PacketTypes.DISCONNECT, "Unspecified error"), + None, + ) + assert client_loop._reconnect_task is None + + +async def test_graceful_disconnect_does_not_reconnect() -> None: + client = Mock(spec=Client) + client_loop = _ClientLoop(client) + + client_loop._on_disconnect( + client, None, DisconnectFlags(False), ReasonCode(PacketTypes.DISCONNECT), None + ) + + assert client_loop._reconnect_task is None + + +async def test_message_send_connect_refused() -> None: + message_send = MessageSend(port=1) + + with pytest.raises(ConnectionRefusedError): + await message_send.__aenter__() + + assert message_send._client is None + assert message_send._client_loop is None + + +async def test_message_send_connect_failure( + monkeypatch: pytest.MonkeyPatch, +) -> None: + clients = _patch_client(monkeypatch) + message_send = MessageSend() + + enter_task = asyncio.create_task(message_send.__aenter__()) + while not clients: + await asyncio.sleep(0) + + message_send._on_connect( + clients[0], + None, + ConnectFlags(False), + ReasonCode(PacketTypes.CONNACK, "Not authorized"), + None, + ) + + with pytest.raises(ConnectionError, match="Not authorized"): + await asyncio.wait_for(enter_task, 1) + + assert message_send._client is None + assert message_send._client_loop is None + + +async def test_message_send_connect_timeout( + monkeypatch: pytest.MonkeyPatch, +) -> None: + clients = _patch_client(monkeypatch) + + # Generous enough that the threaded connect finishes first, even on slow runners. + message_send = MessageSend(connect_timeout=0.2) + with pytest.raises(TimeoutError, match="Timed out waiting"): + await message_send.__aenter__() + + assert clients[0].disconnect.called + assert message_send._client is None + assert message_send._client_loop is None + + # A retry starts a fresh connection. + enter_task = asyncio.create_task(message_send.__aenter__()) + while len(clients) < 2: + await asyncio.sleep(0) + message_send._on_connect( + clients[1], None, ConnectFlags(False), ReasonCode(PacketTypes.CONNACK), None + ) + await asyncio.wait_for(enter_task, 1) + assert message_send._client is clients[1] + + +async def test_message_send_is_reenterable( + monkeypatch: pytest.MonkeyPatch, +) -> None: + clients = _patch_client(monkeypatch) + message_send = MessageSend() + + enter_task = asyncio.create_task(message_send.__aenter__()) + while not clients: + await asyncio.sleep(0) + message_send._on_connect( + clients[0], None, ConnectFlags(False), ReasonCode(PacketTypes.CONNACK), None + ) + await asyncio.wait_for(enter_task, 1) + + first_client = clients[0] + first_client_loop = message_send._client_loop + assert first_client_loop is not None + + exit_task = asyncio.create_task(message_send.__aexit__(None, None, None)) + while first_client.disconnect.call_count == 0: + await asyncio.sleep(0) + first_client_loop._on_disconnect( + first_client, + None, + DisconnectFlags(False), + ReasonCode(PacketTypes.DISCONNECT), + None, + ) + await asyncio.wait_for(exit_task, 1) + assert message_send._client is None + + # The second session must wait for its own CONNACK. + enter_task = asyncio.create_task(message_send.__aenter__()) + while len(clients) < 2: + await asyncio.sleep(0) + while message_send._client is not clients[1]: + await asyncio.sleep(0) + await asyncio.sleep(0) + assert not message_send._connected_event.is_set() + assert not enter_task.done() + message_send._on_connect( + clients[1], None, ConnectFlags(False), ReasonCode(PacketTypes.CONNACK), None + ) + await asyncio.wait_for(enter_task, 1) + assert message_send._client is clients[1] + + +@pytest.mark.parametrize("qos", [0, 1, 2]) +@pytest.mark.parametrize("retain", [True, False]) +async def test_message_send_extracts_qos_and_retain(qos: int, retain: bool) -> None: + message_send, client = _create_mock_message_send() + + event: MessageSendEvent = { + "type": "message.send", + "address": "send-topic", + "headers": [], + "payload": b"test", + "bindings": {"mqtt": {"qos": qos, "retain": retain}}, + } + await _send(message_send, client, event) + + client.publish.assert_called_once_with( + "send-topic", b"test", qos=qos, retain=retain + ) + + +async def test_message_send_maps_headers_to_user_properties() -> None: + message_send, client = _create_mock_message_send(MQTTv5) + + event: MessageSendEvent = { + "type": "message.send", + "address": "send-topic", + "headers": [(b"content-type", b"application/json"), (b"trace-id", b"abc")], + "payload": b"test", + } + await _send(message_send, client, event) + + client.publish.assert_called_once() + call = client.publish.call_args + assert call.args == ("send-topic", b"test") + assert set(call.kwargs) == {"qos", "retain", "properties"} + assert call.kwargs["qos"] == 0 + assert call.kwargs["retain"] is False + properties = call.kwargs["properties"] + assert properties.ContentType == "application/json" + assert properties.UserProperty == [("trace-id", "abc")] + + +async def test_message_send_drops_headers_for_mqtt_v3() -> None: + message_send, client = _create_mock_message_send() + + event: MessageSendEvent = { + "type": "message.send", + "address": "send-topic", + "headers": [(b"trace-id", b"abc")], + "payload": b"test", + } + await _send(message_send, client, event) + + client.publish.assert_called_once_with("send-topic", b"test", qos=0, retain=False) + + +async def test_message_send_v5_without_headers_sends_no_properties() -> None: + message_send, client = _create_mock_message_send(MQTTv5) + + event: MessageSendEvent = { + "type": "message.send", + "address": "send-topic", + "headers": [], + "payload": b"test", + } + await _send(message_send, client, event) + + client.publish.assert_called_once_with("send-topic", b"test", qos=0, retain=False) + + +async def test_cancelled_reconnect_disconnects_late_connection( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(amgi_paho_mqtt, "_RECONNECT_DELAY", 0.0) + release = threading.Event() + started = threading.Event() + + def _slow_reconnect() -> int: + started.set() + release.wait(5) + return int(MQTT_ERR_SUCCESS) + + client = Mock(spec=Client) + client.reconnect.side_effect = _slow_reconnect + client_loop = _ClientLoop(client) + client_loop._on_disconnect( + client, + None, + DisconnectFlags(False), + ReasonCode(PacketTypes.DISCONNECT, "Unspecified error"), + None, + ) + while not started.is_set(): + await asyncio.sleep(0) + + await client_loop.close() + client.disconnect.assert_not_called() + + release.set() + while not client.disconnect.called: + await asyncio.sleep(0) + + +async def test_late_reconnect_failure_does_not_disconnect() -> None: + client = Mock(spec=Client) + client_loop = _ClientLoop(client) + + failed: asyncio.Future[int] = asyncio.get_running_loop().create_future() + failed.set_exception(OSError()) + client_loop._disconnect_late_reconnect(failed) + + cancelled: asyncio.Future[int] = asyncio.get_running_loop().create_future() + cancelled.cancel() + client_loop._disconnect_late_reconnect(cancelled) + + refused: asyncio.Future[int] = asyncio.get_running_loop().create_future() + refused.set_result(MQTT_ERR_NO_CONN) + client_loop._disconnect_late_reconnect(refused) + + client.disconnect.assert_not_called() + + +async def test_message_send_connect_timeout_cleans_up_late_connection( + monkeypatch: pytest.MonkeyPatch, +) -> None: + release = threading.Event() + clients = _patch_client(monkeypatch, lambda *a: release.wait(5)) + + message_send = MessageSend(connect_timeout=0.01) + with pytest.raises(TimeoutError, match="Timed out connecting"): + await message_send.__aenter__() + assert clients[0].disconnect.call_count == 1 + + release.set() + while message_send._late_cleanups or clients[0].disconnect.call_count < 2: + await asyncio.sleep(0) + + +async def test_message_send_connect_worker_failure_skips_late_cleanup( + monkeypatch: pytest.MonkeyPatch, +) -> None: + clients = _patch_client(monkeypatch, Mock(side_effect=OSError("refused"))) + + message_send = MessageSend() + with pytest.raises(OSError): + await message_send.__aenter__() + assert clients[0].disconnect.call_count == 1 + assert not message_send._late_cleanups + + +async def test_message_send_cancel_while_connecting_aborts( + monkeypatch: pytest.MonkeyPatch, +) -> None: + clients = _patch_client(monkeypatch) + message_send = MessageSend() + + enter_task = asyncio.create_task(message_send.__aenter__()) + while message_send._client is None: + await asyncio.sleep(0) + enter_task.cancel() + with pytest.raises(asyncio.CancelledError): + await enter_task + + assert clients[0].disconnect.called + assert message_send._client is None + assert message_send._client_loop is None + + +async def test_message_send_cancel_during_connect_cleans_up( + monkeypatch: pytest.MonkeyPatch, +) -> None: + release = threading.Event() + started = threading.Event() + + def _connect(*a: Any) -> None: + started.set() + release.wait(5) + + created = _patch_client(monkeypatch, _connect) + + message_send = MessageSend() + enter_task = asyncio.create_task(message_send.__aenter__()) + while not started.is_set(): + await asyncio.sleep(0.001) + enter_task.cancel() + with pytest.raises(asyncio.CancelledError): + await enter_task + + release.set() + while created[0].disconnect.call_count < 2: + await asyncio.sleep(0.001) + + +async def test_reconnect_retries_on_error_return_code( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(amgi_paho_mqtt, "_RECONNECT_DELAY", 0.0) + client = Mock(spec=Client) + client.reconnect.side_effect = [MQTT_ERR_NO_CONN, MQTT_ERR_SUCCESS] + client_loop = _ClientLoop(client) + + client_loop._schedule_reconnect() + reconnect_task = client_loop._reconnect_task + assert reconnect_task is not None + await asyncio.wait_for(reconnect_task, 1) + + assert client.reconnect.call_count == 2 + + +def test_socket_open_registers_reader_and_close_unregisters() -> None: + _run_on_selector_loop(_socket_open_registers_reader_and_close_unregisters()) + + +async def _socket_open_registers_reader_and_close_unregisters() -> None: + client = Mock(spec=Client) + client.loop_misc.return_value = MQTT_ERR_SUCCESS + client_loop = _ClientLoop(client) + sock, peer = socket.socketpair() + try: + client_loop._on_socket_open(client, None, sock) + misc_task = client_loop._misc_task + assert misc_task is not None + + peer.send(b"x") + while not client.loop_read.called: + await asyncio.sleep(0) + + client_loop._on_socket_register_write(client, None, sock) + while not client.loop_write.called: + await asyncio.sleep(0) + client_loop._on_socket_unregister_write(client, None, sock) + + client_loop._on_socket_close(client, None, sock) + assert client_loop._misc_task is None + await asyncio.sleep(0) + assert misc_task.done() + finally: + sock.close() + peer.close() + + +async def test_socket_open_after_close_does_not_register() -> None: + client = Mock(spec=Client) + client_loop = _ClientLoop(client) + sock, peer = socket.socketpair() + try: + await client_loop.close() + client_loop._on_socket_open(client, None, sock) + assert client_loop._misc_task is None + finally: + sock.close() + peer.close() + + +def test_close_cancels_misc_task() -> None: + _run_on_selector_loop(_close_cancels_misc_task()) + + +async def _close_cancels_misc_task() -> None: + client = Mock(spec=Client) + client.loop_misc.return_value = MQTT_ERR_SUCCESS + client_loop = _ClientLoop(client) + sock, peer = socket.socketpair() + try: + client_loop._on_socket_open(client, None, sock) + misc_task = client_loop._misc_task + assert misc_task is not None + await client_loop.close() + await asyncio.sleep(0) + assert misc_task.done() + finally: + client_loop.loop.remove_reader(sock) + sock.close() + peer.close() + + +def test_call_on_loop_without_running_loop_marshals_to_loop() -> None: + client_loop = _ClientLoop(Mock(spec=Client)) + loop = Mock(spec=asyncio.AbstractEventLoop) + client_loop.__dict__["loop"] = loop + callback = Mock() + + client_loop._call_on_loop(callback) + + loop.call_soon_threadsafe.assert_called_once_with(callback) + callback.assert_not_called() + + +async def test_message_send_connect_timeout_then_worker_failure_skips_cleanup( + monkeypatch: pytest.MonkeyPatch, +) -> None: + release = threading.Event() + + def _connect(*a: Any) -> None: + release.wait(5) + raise OSError("refused") + + clients = _patch_client(monkeypatch, _connect) + + message_send = MessageSend(connect_timeout=0.01) + with pytest.raises(TimeoutError, match="Timed out connecting"): + await message_send.__aenter__() + + release.set() + await asyncio.sleep(0.1) + + assert clients[0].disconnect.call_count == 1 + assert not message_send._late_cleanups + + +async def test_socket_watcher_fallback_drives_client( + monkeypatch: pytest.MonkeyPatch, +) -> None: + # Like the proactor event loop on Windows, which does not support add_reader. + monkeypatch.setattr( + asyncio.get_running_loop(), + "add_reader", + Mock(side_effect=NotImplementedError), + ) + client = Mock(spec=Client) + client.loop_misc.return_value = MQTT_ERR_SUCCESS + callback_threads: set[int] = set() + client.loop_read.side_effect = lambda: callback_threads.add(threading.get_ident()) + client.loop_write.side_effect = lambda: callback_threads.add(threading.get_ident()) + client_loop = _ClientLoop(client) + sock, peer = socket.socketpair() + try: + client_loop._on_socket_open(client, None, sock) + socket_watcher = client_loop._socket_watcher + assert socket_watcher is not None + misc_task = client_loop._misc_task + assert misc_task is not None + + peer.send(b"x") + while not client.loop_read.called: + await asyncio.sleep(0) + + client_loop._on_socket_register_write(client, None, sock) + while not client.loop_write.called: + await asyncio.sleep(0) + client_loop._on_socket_unregister_write(client, None, sock) + + client_loop._on_socket_close(client, None, sock) + assert client_loop._socket_watcher is None + await asyncio.sleep(0) + assert misc_task.done() + await asyncio.to_thread(socket_watcher._thread.join, 1) + assert not socket_watcher._thread.is_alive() + assert callback_threads == {threading.get_ident()} + finally: + sock.close() + peer.close() + + +async def test_socket_watcher_exits_when_socket_closed() -> None: + # Windows may report a socket closed during select as readable, so the read + # callback can be invoked once before the watcher exits. + sock, peer = socket.socketpair() + try: + socket_watcher = _SocketWatcher(asyncio.get_running_loop(), sock, Mock()) + sock.close() + socket_watcher.set_writer(None) + await asyncio.to_thread(socket_watcher._thread.join, 1) + assert not socket_watcher._thread.is_alive() + finally: + peer.close() + + +def test_socket_watcher_exits_when_loop_closed() -> None: + loop = asyncio.new_event_loop() + on_read = Mock() + sock, peer = socket.socketpair() + try: + socket_watcher = _SocketWatcher(loop, sock, on_read) + loop.close() + peer.send(b"x") + socket_watcher._thread.join(1) + assert not socket_watcher._thread.is_alive() + on_read.assert_not_called() + + # Waking a stopped watcher is ignored. + socket_watcher.close() + finally: + sock.close() + peer.close() + + +async def test_socket_watcher_dispatch_after_close_is_ignored() -> None: + on_read = Mock() + on_write = Mock() + sock, peer = socket.socketpair() + try: + socket_watcher = _SocketWatcher(asyncio.get_running_loop(), sock, on_read) + socket_watcher.set_writer(on_write) + socket_watcher.close() + await asyncio.to_thread(socket_watcher._thread.join, 1) + + socket_watcher._dispatch(True, True) + + on_read.assert_not_called() + on_write.assert_not_called() + finally: + sock.close() + peer.close() diff --git a/packages/amgi-paho-mqtt/tests_amgi_paho_mqtt/test_mqtt_message_integration.py b/packages/amgi-paho-mqtt/tests_amgi_paho_mqtt/test_mqtt_message_integration.py index 1576e520..4fbeb28a 100644 --- a/packages/amgi-paho-mqtt/tests_amgi_paho_mqtt/test_mqtt_message_integration.py +++ b/packages/amgi-paho-mqtt/tests_amgi_paho_mqtt/test_mqtt_message_integration.py @@ -4,16 +4,22 @@ from threading import Event from typing import Any from typing import Generator +from unittest.mock import AsyncMock from uuid import uuid4 import pytest +from amgi_paho_mqtt import _run_cli +from amgi_paho_mqtt import MessageSend from amgi_paho_mqtt import PublishError from amgi_paho_mqtt import run from amgi_paho_mqtt import Server from paho.mqtt.client import Client +from paho.mqtt.client import ConnectFlags from paho.mqtt.client import MQTTMessage from paho.mqtt.client import MQTTv5 from paho.mqtt.enums import CallbackAPIVersion +from paho.mqtt.properties import Properties +from paho.mqtt.reasoncodes import ReasonCode from test_utils import assert_run_can_terminate from test_utils import MockApp from testcontainers.mqtt import MosquittoContainer @@ -152,6 +158,136 @@ async def test_lifespan(topic: str, mosquitto_container: MosquittoContainer) -> } +def _mosquitto_address(mosquitto_container: MosquittoContainer) -> tuple[str, int]: + return ( + mosquitto_container.get_container_host_ip(), + mosquitto_container.get_exposed_port(mosquitto_container.MQTT_PORT), + ) + + +@pytest.mark.integration +async def test_injected_message_send( + mosquitto_container: MosquittoContainer, +) -> None: + host, port = _mosquitto_address(mosquitto_container) + receive_topic = f"receive/{uuid4()}" + send_topic = f"send/{uuid4()}" + + subscribe_event = Event() + received_message: MQTTMessage + message_event = Event() + + client = Client(CallbackAPIVersion.VERSION2) + + def _on_connect( + _client: Client, + _userdata: Any, + _connect_flags: ConnectFlags, + _reason_code: ReasonCode, + _properties: Properties | None, + ) -> None: + client.subscribe(send_topic) + + client.on_connect = _on_connect + client.on_subscribe = lambda *_: subscribe_event.set() + + @client.message_callback() + def _message_callback( + _client: Client, _userdata: Any, _message: MQTTMessage + ) -> None: + nonlocal received_message + received_message = _message + message_event.set() + + client.loop_start() + client.connect(host, port, 60) + + app = MockApp() + server = Server( + app, + receive_topic, + host, + port, + str(uuid4()), + message_send=MessageSend(host, port, protocol=MQTTv5), + ) + + async with app.lifespan(server=server): + await asyncio.to_thread(subscribe_event.wait) + + mosquitto_container.publish_message(receive_topic, "") + + async with app.call() as (scope, receive, send): + await send( + { + "type": "message.send", + "address": send_topic, + "headers": [], + "payload": b"injected", + } + ) + + await asyncio.to_thread(message_event.wait) + assert received_message.topic == send_topic + assert received_message.payload == b"injected" + + client.disconnect() + + +@pytest.mark.integration +async def test_injected_message_send_manager( + mosquitto_container: MosquittoContainer, +) -> None: + host, port = _mosquitto_address(mosquitto_container) + receive_topic = f"receive/{uuid4()}" + + entered: list[bool] = [] + exited: list[bool] = [] + message_send = AsyncMock() + + class _MessageSendManager: + async def __aenter__(self) -> Any: + entered.append(True) + return message_send + + async def __aexit__(self, *exc_info: Any) -> None: + exited.append(True) + + app = MockApp() + server = Server( + app, + receive_topic, + host, + port, + str(uuid4()), + message_send=_MessageSendManager(), + ) + + async with app.lifespan(server=server): + mosquitto_container.publish_message(receive_topic, "") + + async with app.call() as (scope, receive, send): + await send( + { + "type": "message.send", + "address": "send/any", + "headers": [], + "payload": b"test", + } + ) + + assert entered == [True] + assert exited == [True] + message_send.assert_awaited_once_with( + { + "type": "message.send", + "address": "send/any", + "headers": [], + "payload": b"test", + } + ) + + @pytest.mark.integration async def test_message_send_deny( app: MockApp, topic: str, mosquitto_container: MosquittoContainer @@ -181,6 +317,16 @@ def test_run(topic: str, mosquitto_container: MosquittoContainer) -> None: ) +@pytest.mark.integration +def test_run_cli(topic: str, mosquitto_container: MosquittoContainer) -> None: + assert_run_can_terminate( + _run_cli, + topic, + host=mosquitto_container.get_container_host_ip(), + port=mosquitto_container.get_exposed_port(mosquitto_container.MQTT_PORT), + ) + + @pytest.mark.integration async def test_message_receive_not_callable( app: MockApp, topic: str, mosquitto_container: MosquittoContainer diff --git a/uv.lock b/uv.lock index e0024335..ebca77bb 100644 --- a/uv.lock +++ b/uv.lock @@ -680,6 +680,7 @@ dependencies = [ { name = "amgi-common" }, { name = "amgi-types" }, { name = "paho-mqtt" }, + { name = "typing-extensions", marker = "python_full_version < '3.11'" }, ] [package.dev-dependencies] @@ -697,6 +698,7 @@ requires-dist = [ { name = "amgi-common", editable = "packages/amgi-common" }, { name = "amgi-types", editable = "packages/amgi-types" }, { name = "paho-mqtt", specifier = ">=2.1.0" }, + { name = "typing-extensions", marker = "python_full_version < '3.11'", specifier = ">=4.15.0" }, ] [package.metadata.requires-dev]