rework mqtt sink according to update aiomqtt library

Author: raymar9

smartmeter_datacollector/sinks/mqtt_sink.py

@@ -5,20 +5,18 @@
 # SPDX-License-Identifier: GPL-2.0-only
 # See LICENSES/README.md for more information.
 #
+import asyncio
 import json
 import logging
 import ssl
 from configparser import SectionProxy
 from dataclasses import dataclass
 from typing import Optional
 
-from asyncio_mqtt import Client
-from asyncio_mqtt.client import ProtocolVersion
-from asyncio_mqtt.error import MqttCodeError, MqttError
-from paho.mqtt.client import MQTT_ERR_NO_CONN
+from aiomqtt import Client, MqttCodeError, MqttError
 
-from smartmeter_datacollector.smartmeter.meter_data import MeterDataPoint
 from smartmeter_datacollector.sinks.data_sink import DataSink
+from smartmeter_datacollector.smartmeter.meter_data import MeterDataPoint
 
 LOGGER = logging.getLogger("sink")
 
@@ -31,13 +29,14 @@ class MqttConfig:
     use_tls: bool = False
     username: Optional[str] = None
     password: Optional[str] = None
-    client_id: Optional[str] = None
     ca_cert_path: Optional[str] = None
     check_hostname: bool = True
     client_cert_path: Optional[str] = None
     client_key_path: Optional[str] = None
 
-    def with_tls(self, ca_cert_path: Optional[str] = None, check_hostname: bool = True) -> "MqttConfig":
+    def with_tls(
+        self, ca_cert_path: Optional[str] = None, check_hostname: bool = True
+    ) -> "MqttConfig":
         self.use_tls = True
         self.ca_cert_path = ca_cert_path
         self.check_hostname = check_hostname
@@ -56,23 +55,29 @@ def with_client_cert_auth(self, cert_path: str, key_path: str) -> "MqttConfig":
 
     @staticmethod
     def from_sink_config(config: SectionProxy) -> "MqttConfig":
-        mqtt_cfg = MqttConfig(config.get("host"), config.getint("port", 1883))
+        host = config.get("host")
+        if host is None:
+            raise ValueError("MQTT config: 'host' must be set")
+        port = config.getint("port", 1883)
+        mqtt_cfg = MqttConfig(host, port)
         if config.getboolean("tls", fallback=False):
-            mqtt_cfg.with_tls(config.get("ca_file_path"), config.getboolean("check_hostname", True))
-        if config.get("username") and config.get("password"):
-            mqtt_cfg.with_user_pass_auth(
-                config.get("username"),
-                config.get("password"))
-        if config.get("client_cert_path") and config.get("client_key_path"):
-            mqtt_cfg.with_client_cert_auth(
-                config.get("client_cert_path"),
-                config.get("client_key_path")
+            mqtt_cfg.with_tls(
+                config.get("ca_file_path"), config.getboolean("check_hostname", True)
             )
+        username = config.get("username")
+        password = config.get("password")
+        if username is not None and password is not None:
+            mqtt_cfg.with_user_pass_auth(username, password)
+        client_cert_path = config.get("client_cert_path")
+        client_key_path = config.get("client_key_path")
+        if client_cert_path is not None and client_key_path is not None:
+            mqtt_cfg.with_client_cert_auth(str(client_cert_path), str(client_key_path))
         return mqtt_cfg
 
 
 class MqttDataSink(DataSink):
     TIMEOUT = 3
+    RETRIES = 2
 
     def __init__(self, config: MqttConfig) -> None:
         tls_context = None
@@ -81,28 +86,26 @@ def __init__(self, config: MqttConfig) -> None:
                 config.ca_cert_path,
                 config.check_hostname,
                 config.client_cert_path,
-                config.client_key_path)
-
-        user_pass_auth = {}
-        if config.username and config.password:
-            user_pass_auth["username"] = config.username
-            user_pass_auth["password"] = config.password
+                config.client_key_path,
+            )
 
         self._client = Client(
             hostname=config.broker_host,
             port=config.port,
-            client_id=config.client_id,
+            username=config.username,
+            password=config.password,
+            timeout=self.TIMEOUT,
             tls_context=tls_context,
-            protocol=ProtocolVersion.V311,
-            clean_session=True,
-            **user_pass_auth)
+        )
+
+        self._client_task: Optional[asyncio.Task] = None
 
     @staticmethod
     def _build_ssl_context(
         ca_file_path: Optional[str] = None,
         check_hostname: bool = True,
         client_cert_path: Optional[str] = None,
-        client_key_path: Optional[str] = None
+        client_key_path: Optional[str] = None,
     ) -> ssl.SSLContext:
         context = ssl.create_default_context(cafile=ca_file_path)
 
@@ -113,61 +116,78 @@ def _build_ssl_context(
             try:
                 context.load_cert_chain(client_cert_path, client_key_path, None)
             except ssl.SSLError as ex:
-                LOGGER.error("Client certificate does not match with the key and is ignored. '%s'", ex)
+                LOGGER.error("Client certificate does not match with the key and is ignored ('%s')", ex)
         return context
 
     async def start(self) -> None:
-        if await self._connect_to_server():
-            LOGGER.info("Connected to MQTT broker.")
+        if self._client_task is not None:
+            LOGGER.warning("MQTT client task is already running")
+            return
+        self._client_task = asyncio.create_task(self._connection_handler())
+        LOGGER.info("Connecting to MQTT broker...")
 
     async def stop(self) -> None:
-        await self._disconnect_from_server()
-        LOGGER.info("Disconnected from MQTT broker.")
+        if self._client_task is None:
+            LOGGER.warning("MQTT client task is not running")
+            return
+        LOGGER.info("Disconnecting from MQTT broker...")
+        self._client_task.cancel()
+        try:
+            await self._client_task
+        except asyncio.CancelledError:
+            pass
+        self._client_task = None
+        LOGGER.info("Disconnected from MQTT broker")
 
     async def send(self, data_point: MeterDataPoint) -> None:
         topic = MqttDataSink.get_topic_name_for_datapoint(data_point)
         dp_json = self.data_point_to_mqtt_json(data_point)
-        try:
-            await self._client.publish(topic, dp_json)
-            LOGGER.debug("%s sent to MQTT broker.", dp_json)
-        except ValueError as ex:
-            LOGGER.error("MQTT payload or topic is invalid: '%s'", ex)
-        except MqttCodeError as ex:
-            if ex.rc == MQTT_ERR_NO_CONN:
-                if await self._connect_to_server():
-                    LOGGER.info("Reconnected to MQTT broker.")
-                    try:
-                        await self._client.publish(topic, dp_json)
-                        LOGGER.debug("%s sent to MQTT broker.", dp_json)
-                    except MqttError:
-                        LOGGER.warning("MQTT message not sent.")
-            else:
-                LOGGER.error("MQTT message sending error: '%s'", ex)
-        except MqttError as ex:
-            LOGGER.error("MQTT message sending error: '%s'", ex)
-
-    async def _connect_to_server(self) -> bool:
-        try:
-            await self._client.connect(timeout=self.TIMEOUT)
-        except MqttError as ex:
-            LOGGER.error(ex)
-            return False
-        return True
 
-    async def _disconnect_from_server(self) -> None:
-        try:
-            await self._client.disconnect(timeout=self.TIMEOUT)
-        except MqttError as ex:
-            LOGGER.error(ex)
-            await self._client.force_disconnect()
+        await self._publish_with_retries(topic, dp_json, retries=self.RETRIES)
+
+    async def _connection_handler(self) -> None:
+        while True:
+            try:
+                async with self._client:
+                    LOGGER.info("Connected to MQTT broker using client ID '%s'", self._client.identifier)
+                    async for msg in self._client.messages:
+                        # just used to keep the connection alive; no messages are expected
+                        LOGGER.debug("Received message on topic '%s': %s", msg.topic, msg.payload)
+            except MqttError:
+                LOGGER.warning("No connection to MQTT broker. Reconnecting in 1 second...")
+                await asyncio.sleep(1)
+            except asyncio.CancelledError:
+                LOGGER.debug("MQTT client task cancelled")
+                break
+
+    async def _publish_with_retries(self, topic: str, payload: str, retries: int = 1) -> None:
+        for attempt in range(retries + 1):
+            try:
+                await self._client.publish(topic, payload)
+                LOGGER.debug("Published MQTT msg with topic '%s': %s", topic, payload)
+                return
+            except ValueError as ex:
+                LOGGER.error("MQTT payload or topic is invalid: '%s'", ex)
+                return
+            except TimeoutError:
+                LOGGER.warning("MQTT publish attempt %d failed: not connected to broker", attempt + 1)
+            except (MqttCodeError, MqttError) as ex:
+                LOGGER.warning("MQTT publish attempt %d failed: %s", attempt + 1, ex)
+            if attempt < retries:
+                await asyncio.sleep(1)
+
+        LOGGER.error("Failed to publish MQTT msg with topic '%s' after %d attempts. Discarding msg",
+                     topic, retries + 1)
 
     @staticmethod
     def get_topic_name_for_datapoint(data_point: MeterDataPoint) -> str:
         return f"smartmeter/{data_point.source}/{data_point.type.identifier}"
 
     @staticmethod
     def data_point_to_mqtt_json(data_point: MeterDataPoint) -> str:
-        return json.dumps({
-            "value": data_point.value,
-            "timestamp": int(data_point.timestamp.timestamp())
-        })
+        return json.dumps(
+            {
+                "value": data_point.value,
+                "timestamp": int(data_point.timestamp.timestamp()),
+            }
+        )

tests/test_mqtt_sink.py

@@ -7,68 +7,74 @@
 #
 import configparser
 import json
-import sys
 from datetime import datetime, timezone
 from unittest import mock
 
 import pytest
-from asyncio_mqtt.error import MqttCodeError
+import pytest_mock
+from aiomqtt import MqttCodeError
 from paho.mqtt.client import MQTT_ERR_NO_CONN
-from pytest_mock.plugin import MockerFixture
 
 from smartmeter_datacollector.sinks.mqtt_sink import MqttConfig, MqttDataSink
 from smartmeter_datacollector.smartmeter.meter_data import MeterDataPoint, MeterDataPointType
 
-TEST_TYPE = MeterDataPointType("TEST_TYPE", "test type", "unit")
+TEST_DATA_POINT_TYPE = MeterDataPointType("TEST_TYPE", "test type", "unit")
 
 
-@pytest.mark.asyncio
-async def test_mqtt_sink_start_stop(mocker: MockerFixture):
-    config = MqttConfig("localhost")
-    sink = MqttDataSink(config)
-    client_mock = mocker.patch.object(sink, "_client", autospec=True)
-
-    await sink.start()
-    client_mock.connect.assert_awaited_once()
-
-    await sink.stop()
-    client_mock.disconnect.assert_awaited_once()
+@pytest.fixture(autouse=True)
+def mocked_mqtt_client(mocker: pytest_mock.MockerFixture):
+    mock_client_class = mocker.patch("smartmeter_datacollector.sinks.mqtt_sink.Client", autospec=True)
+    return mock_client_class.return_value
 
 
 @pytest.mark.asyncio
-async def test_mqtt_sink_send_point_when_started(mocker: MockerFixture):
+async def test_mqtt_sink_send_datapoint(mocked_mqtt_client):
     config = MqttConfig("localhost")
     sink = MqttDataSink(config)
-    client_mock = mocker.patch.object(sink, "_client", autospec=True)
-    data_point = MeterDataPoint(TEST_TYPE, 1.0, "test_source", datetime.now(timezone.utc))
-    expected_topic = f"smartmeter/test_source/{TEST_TYPE.identifier}"
+
+    data_point = MeterDataPoint(TEST_DATA_POINT_TYPE, 1.0, "test_source", datetime.now(timezone.utc))
+    expected_topic = f"smartmeter/test_source/{TEST_DATA_POINT_TYPE.identifier}"
     expected_payload = json.dumps({
         "value": data_point.value,
         "timestamp": int(data_point.timestamp.timestamp()),
     })
 
-    await sink.start()
     await sink.send(data_point)
 
-    client_mock.publish.assert_awaited_once_with(expected_topic, expected_payload)
+    mocked_mqtt_client.publish.assert_awaited_with(expected_topic, expected_payload)
 
 
 @pytest.mark.asyncio
-async def test_mqtt_sink_send_reconnect_when_not_started(mocker: MockerFixture):
+async def test_mqtt_sink_retry_sending_datapoint(mocked_mqtt_client: mock.MagicMock):
     config = MqttConfig("localhost")
     sink = MqttDataSink(config)
-    client_mock = mocker.patch.object(sink, "_client", autospec=True)
-    data_point = MeterDataPoint(TEST_TYPE, 1.0, "test_source", datetime.now(timezone.utc))
+    data_point = MeterDataPoint(TEST_DATA_POINT_TYPE, 1.0, "test_source", datetime.now(timezone.utc))
 
-    client_mock.publish.side_effect = MqttCodeError(MQTT_ERR_NO_CONN)
-    await sink.send(data_point)
+    # Simulate publish raising MqttCodeError the first time, then succeeding
+    mocked_mqtt_client.publish.side_effect = [MqttCodeError(MQTT_ERR_NO_CONN), None]
+
+    with mock.patch("smartmeter_datacollector.sinks.mqtt_sink.asyncio.sleep"):
+        await sink.send(data_point)
+
+    assert mocked_mqtt_client.publish.await_count == 2
+
+
+@pytest.mark.asyncio
+async def test_mqtt_sink_only_retry_sending_3_times(mocked_mqtt_client: mock.MagicMock):
+    config = MqttConfig("localhost")
+    sink = MqttDataSink(config)
+    data_point = MeterDataPoint(TEST_DATA_POINT_TYPE, 1.0, "test_source", datetime.now(timezone.utc))
+
+    mocked_mqtt_client.publish.side_effect = [MqttCodeError(MQTT_ERR_NO_CONN)]*4
+
+    with mock.patch("smartmeter_datacollector.sinks.mqtt_sink.asyncio.sleep"):
+        await sink.send(data_point)
 
-    client_mock.publish.assert_awaited()
-    assert client_mock.publish.await_count == 2
-    client_mock.connect.assert_awaited_once()
+    assert mocked_mqtt_client.publish.await_count == 3
 
 
-def test_mqtt_config_unencrypted_unauthorized():
+@mock.patch("smartmeter_datacollector.sinks.mqtt_sink.Client", new_callable=mock.MagicMock)
+def test_mqtt_config_unencrypted_unauthorized(_: mock.MagicMock):
     cfg_parser = configparser.ConfigParser()
     cfg_parser.read_dict({
         "sink": {
@@ -91,7 +97,8 @@ def test_mqtt_config_unencrypted_unauthorized():
     sink = MqttDataSink(cfg)
 
 
-def test_mqtt_config_encrypted_unauthorized():
+@mock.patch("smartmeter_datacollector.sinks.mqtt_sink.Client", new_callable=mock.MagicMock)
+def test_mqtt_config_encrypted_unauthorized(_: mock.MagicMock):
     cfg_parser = configparser.ConfigParser()
     cfg_parser.read_dict({
         "sink": {
@@ -117,7 +124,8 @@ def test_mqtt_config_encrypted_unauthorized():
         sink = MqttDataSink(cfg)
 
 
-def test_mqtt_config_encrypted_authorized_user_pass():
+@mock.patch("smartmeter_datacollector.sinks.mqtt_sink.Client", new_callable=mock.MagicMock)
+def test_mqtt_config_encrypted_authorized_user_pass(_: mock.MagicMock):
     cfg_parser = configparser.ConfigParser()
     cfg_parser.read_dict({
         "sink": {
@@ -146,7 +154,8 @@ def test_mqtt_config_encrypted_authorized_user_pass():
     sink = MqttDataSink(cfg)
 
 
-def test_mqtt_config_encrypted_authorized_client_cert():
+@mock.patch("smartmeter_datacollector.sinks.mqtt_sink.Client", new_callable=mock.MagicMock)
+def test_mqtt_config_encrypted_authorized_client_cert(_: mock.MagicMock):
     cfg_parser = configparser.ConfigParser()
     cfg_parser.read_dict({
         "sink": {