run linters

Author: ovv

quixstreams/sinks/community/mqtt.py

@@ -1,9 +1,9 @@
-from quixstreams.sinks.base.sink import BaseSink
-from quixstreams.sinks.base.exceptions import SinkBackpressureError
-from typing import List, Tuple, Any
-from quixstreams.models.types import HeaderValue
-from datetime import datetime
 import json
+from datetime import datetime
+from typing import Any, List, Tuple
+
+from quixstreams.models.types import HeaderValue
+from quixstreams.sinks.base.sink import BaseSink
 
 try:
     import paho.mqtt.client as paho
@@ -14,21 +14,24 @@
         "run pip install quixstreams[paho-mqtt] to fix it"
     ) from exc
 
+
 class MQTTSink(BaseSink):
     """
     A sink that publishes messages to an MQTT broker.
     """
 
-    def __init__(self,
-                 mqtt_client_id: str,
-                 mqtt_server: str,
-                 mqtt_port: int,
-                 mqtt_topic_root: str,
-                 mqtt_username: str = None,
-                 mqtt_password: str = None,
-                 mqtt_version: str = "3.1.1",
-                 tls_enabled: bool = True,
-                 qos: int = 1):
+    def __init__(
+        self,
+        mqtt_client_id: str,
+        mqtt_server: str,
+        mqtt_port: int,
+        mqtt_topic_root: str,
+        mqtt_username: str = None,
+        mqtt_password: str = None,
+        mqtt_version: str = "3.1.1",
+        tls_enabled: bool = True,
+        qos: int = 1,
+    ):
         """
         Initialize the MQTTSink.
 
@@ -42,21 +45,27 @@ def __init__(self,
         :param tls_enabled: Whether to use TLS encryption. Defaults to True
         :param qos: Quality of Service level (0, 1, or 2). Defaults to 1
         """
-        
+
         super().__init__()
-        
+
         self.mqtt_version = mqtt_version
         self.mqtt_username = mqtt_username
         self.mqtt_password = mqtt_password
         self.mqtt_topic_root = mqtt_topic_root
         self.tls_enabled = tls_enabled
         self.qos = qos
 
-        self.mqtt_client = paho.Client(callback_api_version=paho.CallbackAPIVersion.VERSION2,
-                                client_id = mqtt_client_id, userdata = None, protocol = self._mqtt_protocol_version())
+        self.mqtt_client = paho.Client(
+            callback_api_version=paho.CallbackAPIVersion.VERSION2,
+            client_id=mqtt_client_id,
+            userdata=None,
+            protocol=self._mqtt_protocol_version(),
+        )
 
         if self.tls_enabled:
-            self.mqtt_client.tls_set(tls_version = mqtt.client.ssl.PROTOCOL_TLS)  # we'll be using tls now
+            self.mqtt_client.tls_set(
+                tls_version=mqtt.client.ssl.PROTOCOL_TLS
+            )  # we'll be using tls now
 
         self.mqtt_client.reconnect_delay_set(5, 60)
         self._configure_authentication()
@@ -65,17 +74,31 @@ def __init__(self,
         self.mqtt_client.connect(mqtt_server, int(mqtt_port))
 
     # setting callbacks for different events to see if it works, print the message etc.
-    def _mqtt_on_connect_cb(self, client: paho.Client, userdata: any, connect_flags: paho.ConnectFlags,
-                    reason_code: paho.ReasonCode, properties: paho.Properties):
+    def _mqtt_on_connect_cb(
+        self,
+        client: paho.Client,
+        userdata: any,
+        connect_flags: paho.ConnectFlags,
+        reason_code: paho.ReasonCode,
+        properties: paho.Properties,
+    ):
         if reason_code == 0:
-            print("CONNECTED!") # required for Quix to know this has connected
+            print("CONNECTED!")  # required for Quix to know this has connected
         else:
             print(f"ERROR ({reason_code.value}). {reason_code.getName()}")
 
-    def _mqtt_on_disconnect_cb(self, client: paho.Client, userdata: any, disconnect_flags: paho.DisconnectFlags,
-                        reason_code: paho.ReasonCode, properties: paho.Properties):
-        print(f"DISCONNECTED! Reason code ({reason_code.value}) {reason_code.getName()}!")
-        
+    def _mqtt_on_disconnect_cb(
+        self,
+        client: paho.Client,
+        userdata: any,
+        disconnect_flags: paho.DisconnectFlags,
+        reason_code: paho.ReasonCode,
+        properties: paho.Properties,
+    ):
+        print(
+            f"DISCONNECTED! Reason code ({reason_code.value}) {reason_code.getName()}!"
+        )
+
     def _mqtt_protocol_version(self):
         if self.mqtt_version == "3.1":
             return paho.MQTTv31
@@ -90,38 +113,51 @@ def _configure_authentication(self):
         if self.mqtt_username:
             self.mqtt_client.username_pw_set(self.mqtt_username, self.mqtt_password)
 
-    def _publish_to_mqtt(self, data: str, key: bytes, timestamp: datetime, headers: List[Tuple[str, HeaderValue]]):
+    def _publish_to_mqtt(
+        self,
+        data: str,
+        key: bytes,
+        timestamp: datetime,
+        headers: List[Tuple[str, HeaderValue]],
+    ):
         if isinstance(data, bytes):
-            data = data.decode('utf-8')  # Decode bytes to string using utf-8
+            data = data.decode("utf-8")  # Decode bytes to string using utf-8
 
         json_data = json.dumps(data)
-        message_key_string = key.decode('utf-8')  # Convert to string using utf-8 encoding
+        message_key_string = key.decode(
+            "utf-8"
+        )  # Convert to string using utf-8 encoding
         # publish to MQTT
-        self.mqtt_client.publish(self.mqtt_topic_root + "/" + message_key_string, payload = json_data, qos = self.qos)
-
-
-    def add(self,
-            topic: str,
-            partition: int,
-            offset: int,
-            key: bytes,
-            value: bytes,
-            timestamp: datetime,
-            headers: List[Tuple[str, HeaderValue]],
-            **kwargs: Any):
+        self.mqtt_client.publish(
+            self.mqtt_topic_root + "/" + message_key_string,
+            payload=json_data,
+            qos=self.qos,
+        )
+
+    def add(
+        self,
+        topic: str,
+        partition: int,
+        offset: int,
+        key: bytes,
+        value: bytes,
+        timestamp: datetime,
+        headers: List[Tuple[str, HeaderValue]],
+        **kwargs: Any,
+    ):
         self._publish_to_mqtt(value, key, timestamp, headers)
 
     def _construct_topic(self, key):
         if key:
-            key_str = key.decode('utf-8') if isinstance(key, bytes) else str(key)
+            key_str = key.decode("utf-8") if isinstance(key, bytes) else str(key)
             return f"{self.mqtt_topic_root}/{key_str}"
         else:
             return self.mqtt_topic_root
 
     def on_paused(self, topic: str, partition: int):
         # not used
         pass
-    
+
     def flush(self, topic: str, partition: str):
         # not used
         pass

tests/test_quixstreams/test_sinks/test_community/test_mqtt_sink.py

@@ -1,8 +1,11 @@
-from unittest.mock import MagicMock, patch
-import pytest
 from datetime import datetime
+from unittest.mock import patch
+
+import pytest
+
 from quixstreams.sinks.community.mqtt import MQTTSink
 
+
 @pytest.fixture()
 def mqtt_sink_factory():
     def factory(
@@ -16,7 +19,7 @@ def factory(
         tls_enabled: bool = True,
         qos: int = 1,
     ) -> MQTTSink:
-        with patch('paho.mqtt.client.Client') as MockClient:
+        with patch("paho.mqtt.client.Client") as MockClient:
             mock_mqtt_client = MockClient.return_value
             sink = MQTTSink(
                 mqtt_client_id=mqtt_client_id,
@@ -27,13 +30,14 @@ def factory(
                 mqtt_password=mqtt_password,
                 mqtt_version=mqtt_version,
                 tls_enabled=tls_enabled,
-                qos=qos
+                qos=qos,
             )
             sink.mqtt_client = mock_mqtt_client
             return sink, mock_mqtt_client
 
     return factory
 
+
 class TestMQTTSink:
     def test_mqtt_connect(self, mqtt_sink_factory):
         sink, mock_mqtt_client = mqtt_sink_factory()
@@ -59,17 +63,19 @@ def test_mqtt_publish(self, mqtt_sink_factory):
             partition=0,
             offset=1,
             key=key,
-            value=data.encode('utf-8'),
+            value=data.encode("utf-8"),
             timestamp=timestamp,
-            headers=headers
+            headers=headers,
         )
 
         mock_mqtt_client.publish.assert_called_once_with(
             "test/topic/test_key", payload='"test_data"', qos=1
         )
 
     def test_mqtt_authentication(self, mqtt_sink_factory):
-        sink, mock_mqtt_client = mqtt_sink_factory(mqtt_username="user", mqtt_password="pass")
+        sink, mock_mqtt_client = mqtt_sink_factory(
+            mqtt_username="user", mqtt_password="pass"
+        )
         mock_mqtt_client.username_pw_set.assert_called_once_with("user", "pass")
 
     def test_mqtt_disconnect_on_delete(self, mqtt_sink_factory):