From 59dfc3575949d521cb6d859d4f20baa4af0692e2 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Pierre=20St=C3=A5hl?= <pierre.staahl@gmail.com>
Date: Mon, 5 Oct 2020 12:48:45 +0200
Subject: [PATCH] Fix bug when unloading an entry

---
 custom_components/localtuya/__init__.py        |  3 +--
 custom_components/localtuya/common.py          | 16 ++++++++++++++--
 custom_components/localtuya/pytuya/__init__.py |  8 +++++---
 3 files changed, 20 insertions(+), 7 deletions(-)

diff --git a/custom_components/localtuya/__init__.py b/custom_components/localtuya/__init__.py
index 623c2de..3f7b0c6 100644
--- a/custom_components/localtuya/__init__.py
+++ b/custom_components/localtuya/__init__.py
@@ -56,12 +56,11 @@ import logging
 from homeassistant.config_entries import SOURCE_IMPORT, ConfigEntry
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.const import (
+    CONF_DEVICE_ID,
     CONF_PLATFORM,
     CONF_ENTITIES,
     SERVICE_RELOAD,
 )
-from homeassistant.helpers.event import async_track_time_interval
-from homeassistant.helpers.dispatcher import async_dispatcher_send
 from homeassistant.helpers.reload import async_integration_yaml_config
 
 from .const import DOMAIN, TUYA_DEVICE
diff --git a/custom_components/localtuya/common.py b/custom_components/localtuya/common.py
index 556128b..dcbfadd 100644
--- a/custom_components/localtuya/common.py
+++ b/custom_components/localtuya/common.py
@@ -101,6 +101,7 @@ class TuyaDevice(pytuya.TuyaListener):
         self._interface = None
         self._status = {}
         self._dps_to_request = {}
+        self._is_closing = False
         self._connect_task = None
         self._connection_attempts = 0
 
@@ -110,7 +111,7 @@ class TuyaDevice(pytuya.TuyaListener):
 
     def connect(self):
         """Connet to device if not already connected."""
-        if self._connect_task is None or self._interface:
+        if not self._is_closing and self._connect_task is None and not self._interface:
             self._connect_task = asyncio.ensure_future(self._make_connection())
 
     async def _make_connection(self):
@@ -148,8 +149,16 @@ class TuyaDevice(pytuya.TuyaListener):
             self._hass.loop.call_soon(self.connect)
         self._connect_task = None
 
+    def close(self):
+        """Close connection and stop re-connect loop."""
+        self._is_closing = True
+        if self._connect_task:
+            self._connect_task.cancel()
+        if self._interface:
+            self._interface.close()
+
     async def set_dps(self, state, dps_index):
-        """Change value of a DP of the Tuya device and update the cached status."""
+        """Change value of a DP of the Tuya device."""
         if self._interface is not None:
             try:
                 await self._interface.set_dps(state, dps_index)
@@ -172,6 +181,9 @@ class TuyaDevice(pytuya.TuyaListener):
     @callback
     def disconnected(self, exc):
         """Device disconnected."""
+        if exc is not None:
+            _LOGGER.error("Disconnected from %: %s", exc)
+
         signal = f"localtuya_{self._config_entry[CONF_DEVICE_ID]}"
         async_dispatcher_send(self._hass, signal, None)
 
diff --git a/custom_components/localtuya/pytuya/__init__.py b/custom_components/localtuya/pytuya/__init__.py
index 6e01106..a5c4a80 100644
--- a/custom_components/localtuya/pytuya/__init__.py
+++ b/custom_components/localtuya/pytuya/__init__.py
@@ -359,7 +359,11 @@ class TuyaProtocol(asyncio.Protocol):
         dev_type = self.dev_type
 
         # Wait for special sequence number if heartbeat
-        seqno = MessageDispatcher.HEARTBEAT_SEQNO if command == HEARTBEAT else (self.seqno - 1)
+        seqno = (
+            MessageDispatcher.HEARTBEAT_SEQNO
+            if command == HEARTBEAT
+            else (self.seqno - 1)
+        )
 
         self.transport.write(payload)
         msg = await self.dispatcher.wait_for(seqno)
@@ -387,8 +391,6 @@ class TuyaProtocol(asyncio.Protocol):
 
     async def heartbeat(self):
         """Send a heartbeat message."""
-        # We don't expect a response to this, just send blindly
-        #self.transport.write(self._generate_payload(HEARTBEAT))
         return await self.exchange(HEARTBEAT)
 
     async def set_dps(self, value, dps_index):