mqtt: various fixes

1 files changed, 21 insertions, 0 deletions

diff --git a/include/py/homekit/mqtt/_wrapper.py b/include/py/homekit/mqtt/_wrapper.py
index 3c2774c..5fc33fe 100644
--- a/include/py/homekit/mqtt/_wrapper.py
+++ b/include/py/homekit/mqtt/_wrapper.py
@@ -7,6 +7,8 @@ from ..util import strgen
 
 class MqttWrapper(Mqtt):
     _nodes: list[MqttNode]
+    _connect_callbacks: list[callable]
+    _disconnect_callbacks: list[callable]
 
     def __init__(self,
                  client_id: str,
@@ -18,17 +20,30 @@ class MqttWrapper(Mqtt):
         super().__init__(clean_session=clean_session,
                          client_id=client_id)
         self._nodes = []
+        self._connect_callbacks = []
+        self._disconnect_callbacks = []
         self._topic_prefix = topic_prefix
 
     def on_connect(self, client: mqtt.Client, userdata, flags, rc):
         super().on_connect(client, userdata, flags, rc)
         for node in self._nodes:
             node.on_connect(self)
+        for f in self._connect_callbacks:
+            try:
+                f()
+            except Exception as e:
+                self._logger.exception(e)
 
     def on_disconnect(self, client: mqtt.Client, userdata, rc):
         super().on_disconnect(client, userdata, rc)
         for node in self._nodes:
             node.on_disconnect()
+        for f in self._disconnect_callbacks:
+            try:
+                f()
+            except Exception as e:
+                self._logger.exception(e)
+
 
     def on_message(self, client: mqtt.Client, userdata, msg):
         try:
@@ -40,6 +55,12 @@ class MqttWrapper(Mqtt):
         except Exception as e:
             self._logger.exception(str(e))
 
+    def add_connect_callback(self, f: callable):
+        self._connect_callbacks.append(f)
+
+    def add_disconnect_callback(self, f: callable):
+       self._disconnect_callbacks.append(f)
+
     def add_node(self, node: MqttNode):
         self._nodes.append(node)
         if self._connected: