Refactor: Prevent heartbeater loop duplicates

xZetsubou · Jan 9, 2024 · 0cb5c95 · 0cb5c95
1 parent 42b6569
commit 0cb5c95
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 6 deletions.
diff --git a/custom_components/localtuya/common.py b/custom_components/localtuya/common.py
@@ -280,8 +280,7 @@ async def _make_connection(self):
                 status = await self._interface.status(cid=self._node_id)
                 if status is None:  # and not self.is_subdevice
                     raise Exception("Failed to retrieve status")
-                if not self._interface.heartbeater:
-                    self._interface.start_heartbeat()
+                self._interface.start_heartbeat()
                 self.status_updated(status)
 
             except UnicodeDecodeError as e:

diff --git a/custom_components/localtuya/core/pytuya/__init__.py b/custom_components/localtuya/core/pytuya/__init__.py
@@ -571,12 +571,14 @@ class MessageDispatcher(ContextualLogger):
     RESET_SEQNO = -101
     SESS_KEY_SEQNO = -102
 
-    def __init__(self, dev_id, listener, protocol_version, local_key, enable_debug):
+    def __init__(
+        self, dev_id, callback_status_update, protocol_version, local_key, enable_debug
+    ):
         """Initialize a new MessageBuffer."""
         super().__init__()
         self.buffer = b""
         self.listeners: dict[str, asyncio.Semaphore] = {}
-        self.listener = listener
+        self.callback_status_update = callback_status_update
         self.version = protocol_version
         self.local_key = local_key
         self.set_logger(_LOGGER, dev_id, enable_debug)
@@ -702,7 +704,7 @@ def _dispatch(self, msg):
                     )
             else:
                 self.debug("Got status update")
-                self.listener(msg)
+                self.callback_status_update(msg)
                 # workdaround for >= v3.4 devices until find prper way to wait seqno correctly.
                 if msg.seqno in self.listeners:
                     sem = self.listeners[msg.seqno]
@@ -881,7 +883,9 @@ async def heartbeat_loop():
             self.transport = None
             transport.close()
 
-        self.heartbeater = self.loop.create_task(heartbeat_loop())
+        if self.heartbeater is None:
+            # Prevent duplicates heartbeat task
+            self.heartbeater = self.loop.create_task(heartbeat_loop())
 
     def data_received(self, data):
         """Received data from device."""