Add a new option "poll_only" to force polling of devices.

This may help with devices that stopped working properly on recent
versions, such as #338

- also temporarily switch to polling when testing config changes.

Devices seem to be at their connection limit, and won't return data
when being tested, and in some cases will not respond to commands even
though they are connected and showing up to date status.

custom_components/tuya_local/__init__.py

@@ -16,6 +16,7 @@ from homeassistant.helpers.entity_registry import async_migrate_entries
 from .const import (
     CONF_DEVICE_ID,
     CONF_LOCAL_KEY,
+    CONF_POLL_ONLY,
     CONF_PROTOCOL_VERSION,
     CONF_TYPE,
     DOMAIN,
@@ -151,6 +152,19 @@ async def async_migrate_entry(hass, entry: ConfigEntry):
         entry.options = {}
         entry.version = 10
 
+    if entry.version <= 10:
+        conf = entry.data | entry.options
+        entry.data = {
+            CONF_DEVICE_ID: conf[CONF_DEVICE_ID],
+            CONF_LOCAL_KEY: conf[CONF_LOCAL_KEY],
+            CONF_HOST: conf[CONF_HOST],
+            CONF_TYPE: conf[CONF_TYPE],
+            CONF_PROTOCOL_VERSION: "auto",
+            CONF_POLL_ONLY: False,
+        }
+        entry.options = {}
+        entry.version = 11
+
     return True
 
 

custom_components/tuya_local/config_flow.py

@@ -11,8 +11,9 @@ from .const import (
     API_PROTOCOL_VERSIONS,
     CONF_DEVICE_ID,
     CONF_LOCAL_KEY,
-    CONF_TYPE,
+    CONF_POLL_ONLY,
     CONF_PROTOCOL_VERSION,
+    CONF_TYPE,
 )
 from .helpers.device_config import get_config
 
@@ -20,8 +21,8 @@ _LOGGER = logging.getLogger(__name__)
 
 
 class ConfigFlowHandler(config_entries.ConfigFlow, domain=DOMAIN):
-    VERSION = 10
-    CONNECTION_CLASS = config_entries.CONN_CLASS_LOCAL_POLL
+    VERSION = 11
+    CONNECTION_CLASS = config_entries.CONN_CLASS_LOCAL_PUSH
     device = None
     data = {}
 
@@ -31,6 +32,7 @@ class ConfigFlowHandler(config_entries.ConfigFlow, domain=DOMAIN):
         host_opts = {}
         key_opts = {}
         proto_opts = {"default": "auto"}
+        polling_opts = {"default": False}
 
         if user_input is not None:
             await self.async_set_unique_id(user_input[CONF_DEVICE_ID])
@@ -46,6 +48,7 @@ class ConfigFlowHandler(config_entries.ConfigFlow, domain=DOMAIN):
                 host_opts["default"] = user_input[CONF_HOST]
                 key_opts["default"] = user_input[CONF_LOCAL_KEY]
                 proto_opts["default"] = user_input[CONF_PROTOCOL_VERSION]
+                polling_opts["default"] = user_input[CONF_POLL_ONLY]
 
         return self.async_show_form(
             step_id="user",
@@ -58,6 +61,7 @@ class ConfigFlowHandler(config_entries.ConfigFlow, domain=DOMAIN):
                         CONF_PROTOCOL_VERSION,
                         **proto_opts,
                     ): vol.In(["auto"] + API_PROTOCOL_VERSIONS),
+                    vol.Required(CONF_POLL_ONLY, **polling_opts): bool,
                 }
             ),
             errors=errors,
@@ -147,6 +151,9 @@ class OptionsFlowHandler(config_entries.OptionsFlow):
             vol.Required(
                 CONF_PROTOCOL_VERSION, default=config.get(CONF_PROTOCOL_VERSION, "auto")
             ): vol.In(["auto"] + API_PROTOCOL_VERSIONS),
+            vol.Required(
+                CONF_POLL_ONLY, default=config.get(CONF_POLL_ONLY, False)
+            ): bool,
         }
         cfg = get_config(config[CONF_TYPE])
         if cfg is None:
@@ -160,6 +167,11 @@ class OptionsFlowHandler(config_entries.OptionsFlow):
 
 
 async def async_test_connection(config: dict, hass: HomeAssistant):
+    domain_data = hass.data.get(DOMAIN)
+    existing = domain_data.get(config[CONF_DEVICE_ID]) if domain_data else None
+    if existing:
+        existing["device"].pause()
+
     device = TuyaLocalDevice(
         "Test",
         config[CONF_DEVICE_ID],
@@ -167,6 +179,10 @@ async def async_test_connection(config: dict, hass: HomeAssistant):
         config[CONF_LOCAL_KEY],
         config[CONF_PROTOCOL_VERSION],
         hass,
+        True,
     )
     await device.async_refresh()
+    if existing:
+        existing["device"].resume()
+
     return device if device.has_returned_state else None

custom_components/tuya_local/const.py

@@ -5,5 +5,6 @@ DOMAIN = "tuya_local"
 CONF_DEVICE_ID = "device_id"
 CONF_LOCAL_KEY = "local_key"
 CONF_TYPE = "type"
+CONF_POLL_ONLY = "poll_only"
 CONF_PROTOCOL_VERSION = "protocol_version"
 API_PROTOCOL_VERSIONS = [3.3, 3.1, 3.2, 3.4, 3.5]

custom_components/tuya_local/device.py

@@ -22,6 +22,7 @@ from .const import (
     API_PROTOCOL_VERSIONS,
     CONF_DEVICE_ID,
     CONF_LOCAL_KEY,
+    CONF_POLL_ONLY,
     CONF_PROTOCOL_VERSION,
     DOMAIN,
 )
@@ -45,6 +46,7 @@ class TuyaLocalDevice(object):
         local_key,
         protocol_version,
         hass: HomeAssistant,
+        poll_only=False,
     ):
         """
         Represents a Tuya-based device.
@@ -54,6 +56,8 @@ class TuyaLocalDevice(object):
             address (str): The network address.
             local_key (str): The encryption key.
             protocol_version (str | number): The protocol version.
+            hass (HomeAssistant): The Home Assistant instance.
+            poll_only (bool): True if the device should be polled only
         """
         self._name = name
         self._children = []
@@ -65,7 +69,8 @@ class TuyaLocalDevice(object):
         self._api = tinytuya.Device(dev_id, address, local_key)
         self._refresh_task = None
         self._protocol_configured = protocol_version
-
+        self._poll_only = poll_only
+        self._temporary_poll = False
         self._reset_cached_state()
 
         self._hass = hass
@@ -178,22 +183,32 @@ class TuyaLocalDevice(object):
                 f"{self.name} receive loop terminated by exception {t}",
             )
 
+    @property
+    def should_poll(self):
+        return self._poll_only or self._temporary_poll or not self.has_returned_state
+
+    def pause(self):
+        self._temporary_poll = True
+
+    def resume(self):
+        self._temporary_poll = False
+
     async def async_receive(self):
         """Receive messages from a persistent connection asynchronously."""
         # If we didn't yet get any state from the device, we may need to
         # negotiate the protocol before making the connection persistent
-        persist = self.has_returned_state
+        persist = not self.should_poll
         self._api.set_socketPersistent(persist)
         while self._running:
             try:
                 last_cache = self._cached_state["updated_at"]
                 now = time()
-                if persist != self.has_returned_state:
+                if persist == self.should_poll:
                     # use persistent connections after initial communication
                     # has been established.  Until then, we need to rotate
                     # the protocol version, which seems to require a fresh
                     # connection.
-                    persist = self.has_returned_state
+                    persist = not self.should_poll
                     self._api.set_socketPersistent(persist)
 
                 if now - last_cache > self._CACHE_TIMEOUT:
@@ -485,6 +500,7 @@ def setup_device(hass: HomeAssistant, config: dict):
         config[CONF_LOCAL_KEY],
         config[CONF_PROTOCOL_VERSION],
         hass,
+        config[CONF_POLL_ONLY],
     )
     hass.data[DOMAIN][config[CONF_DEVICE_ID]] = {"device": device}
 

custom_components/tuya_local/translations/en.json

@@ -9,7 +9,8 @@
                     "host": "IP address or hostname",
                     "device_id": "Device ID (uuid)",
                     "local_key": "Local key",
-                    "protocol_version": "Protocol version (try auto if not known)"
+                    "protocol_version": "Protocol version (try auto if not known)",
+                    "poll_only": "Poll only (try this if your device does not work fully)"
                 }
             },
             "select_type": {
@@ -43,7 +44,8 @@
                 "data": {
                     "host": "IP address or hostname",
                     "local_key": "Local key",
-                    "protocol_version": "Protocol version (try auto if not known)"
+                    "protocol_version": "Protocol version (try auto if not known)",
+                    "poll_only": "Poll only (try this if your device does not work fully)"
                 }
             }
         },

tests/test_config_flow.py

@@ -15,6 +15,7 @@ from custom_components.tuya_local import (
 from custom_components.tuya_local.const import (
     CONF_DEVICE_ID,
     CONF_LOCAL_KEY,
+    CONF_POLL_ONLY,
     CONF_PROTOCOL_VERSION,
     CONF_TYPE,
     DOMAIN,
@@ -49,12 +50,13 @@ async def test_init_entry(hass):
     """Test initialisation of the config flow."""
     entry = MockConfigEntry(
         domain=DOMAIN,
-        version=10,
+        version=11,
         title="test",
         data={
             CONF_DEVICE_ID: "deviceid",
             CONF_HOST: "hostname",
             CONF_LOCAL_KEY: "localkey",
+            CONF_POLL_ONLY: False,
             CONF_PROTOCOL_VERSION: "auto",
             CONF_TYPE: "kogan_kahtp_heater",
         },
@@ -249,6 +251,8 @@ async def test_async_test_connection_valid(mock_device, hass):
     mock_instance = AsyncMock()
     mock_instance.has_returned_state = True
     mock_device.return_value = mock_instance
+    hass.data[DOMAIN] = {"deviceid": {"device": mock_instance}}
+
     device = await config_flow.async_test_connection(
         {
             CONF_DEVICE_ID: "deviceid",
@@ -259,6 +263,8 @@ async def test_async_test_connection_valid(mock_device, hass):
         hass,
     )
     assert device == mock_instance
+    mock_instance.pause.assert_called_once()
+    mock_instance.resume.assert_called_once()
 
 
 @pytest.mark.asyncio
@@ -293,6 +299,7 @@ async def test_flow_user_init_invalid_config(mock_test, hass):
             CONF_HOST: "hostname",
             CONF_LOCAL_KEY: "badkey",
             CONF_PROTOCOL_VERSION: "auto",
+            CONF_POLL_ONLY: False,
         },
     )
     assert {"base": "connection"} == result["errors"]
@@ -436,6 +443,7 @@ async def test_flow_choose_entities_creates_config_entry(hass, bypass_setup):
             CONF_DEVICE_ID: "deviceid",
             CONF_LOCAL_KEY: "localkey",
             CONF_HOST: "hostname",
+            CONF_POLL_ONLY: False,
             CONF_PROTOCOL_VERSION: "auto",
             CONF_TYPE: "kogan_kahtp_heater",
         },
@@ -450,7 +458,7 @@ async def test_flow_choose_entities_creates_config_entry(hass, bypass_setup):
             },
         )
         expected = {
-            "version": 10,
+            "version": 11,
             "context": {"source": "choose_entities"},
             "type": "create_entry",
             "flow_id": ANY,
@@ -464,6 +472,7 @@ async def test_flow_choose_entities_creates_config_entry(hass, bypass_setup):
                 CONF_DEVICE_ID: "deviceid",
                 CONF_HOST: "hostname",
                 CONF_LOCAL_KEY: "localkey",
+                CONF_POLL_ONLY: False,
                 CONF_PROTOCOL_VERSION: "auto",
                 CONF_TYPE: "kogan_kahtp_heater",
             },
@@ -476,13 +485,14 @@ async def test_options_flow_init(hass):
     """Test config flow options."""
     config_entry = MockConfigEntry(
         domain=DOMAIN,
-        version=10,
+        version=11,
         unique_id="uniqueid",
         data={
             CONF_DEVICE_ID: "deviceid",
             CONF_HOST: "hostname",
             CONF_LOCAL_KEY: "localkey",
             CONF_NAME: "test",
+            CONF_POLL_ONLY: False,
             CONF_PROTOCOL_VERSION: "auto",
             CONF_TYPE: "smartplugv1",
         },
@@ -513,13 +523,14 @@ async def test_options_flow_modifies_config(mock_test, hass):
 
     config_entry = MockConfigEntry(
         domain=DOMAIN,
-        version=10,
+        version=11,
         unique_id="uniqueid",
         data={
             CONF_DEVICE_ID: "deviceid",
             CONF_HOST: "hostname",
             CONF_LOCAL_KEY: "localkey",
             CONF_NAME: "test",
+            CONF_POLL_ONLY: False,
             CONF_PROTOCOL_VERSION: "auto",
             CONF_TYPE: "kogan_kahtp_heater",
         },
@@ -536,12 +547,14 @@ async def test_options_flow_modifies_config(mock_test, hass):
         user_input={
             CONF_HOST: "new_hostname",
             CONF_LOCAL_KEY: "new_key",
+            CONF_POLL_ONLY: False,
             CONF_PROTOCOL_VERSION: 3.3,
         },
     )
     expected = {
         CONF_HOST: "new_hostname",
         CONF_LOCAL_KEY: "new_key",
+        CONF_POLL_ONLY: False,
         CONF_PROTOCOL_VERSION: 3.3,
     }
     assert "create_entry" == result["type"]
@@ -557,13 +570,14 @@ async def test_options_flow_fails_when_connection_fails(mock_test, hass):
 
     config_entry = MockConfigEntry(
         domain=DOMAIN,
-        version=10,
+        version=11,
         unique_id="uniqueid",
         data={
             CONF_DEVICE_ID: "deviceid",
             CONF_HOST: "hostname",
             CONF_LOCAL_KEY: "localkey",
             CONF_NAME: "test",
+            CONF_POLL_ONLY: False,
             CONF_PROTOCOL_VERSION: "auto",
             CONF_TYPE: "smartplugv1",
         },
@@ -595,13 +609,14 @@ async def test_options_flow_fails_when_config_is_missing(mock_test, hass):
 
     config_entry = MockConfigEntry(
         domain=DOMAIN,
-        version=10,
+        version=11,
         unique_id="uniqueid",
         data={
             CONF_DEVICE_ID: "deviceid",
             CONF_HOST: "hostname",
             CONF_LOCAL_KEY: "localkey",
             CONF_NAME: "test",
+            CONF_POLL_ONLY: False,
             CONF_PROTOCOL_VERSION: "auto",
             CONF_TYPE: "non_existing",
         },
@@ -622,13 +637,14 @@ async def test_async_setup_entry_for_switch(mock_device, hass):
     """Test setting up based on a config entry.  Repeats test_init_entry."""
     config_entry = MockConfigEntry(
         domain=DOMAIN,
-        version=10,
+        version=11,
         unique_id="uniqueid",
         data={
             CONF_DEVICE_ID: "deviceid",
             CONF_HOST: "hostname",
             CONF_LOCAL_KEY: "localkey",
             CONF_NAME: "test",
+            CONF_POLL_ONLY: False,
             CONF_PROTOCOL_VERSION: "auto",
             CONF_TYPE: "smartplugv2",
         },

tests/test_device.py

@@ -594,3 +594,24 @@ class TestDevice(IsolatedAsyncioTestCase):
         except StopAsyncIteration:
             pass
         self.mock_api().set_socketPersistent.assert_called_once_with(False)
+
+    def test_should_poll(self):
+        self.subject._cached_state = {"1": "sample", "updated_at": time()}
+        self.subject._poll_only = False
+        self.subject._temporary_poll = False
+
+        # Test temporary poll via pause/resume
+        self.assertFalse(self.subject.should_poll)
+        self.subject.pause()
+        self.assertTrue(self.subject.should_poll)
+        self.subject.resume()
+        self.assertFalse(self.subject.should_poll)
+
+        # Test configured polling
+        self.subject._poll_only = True
+        self.assertTrue(self.subject.should_poll)
+        self.subject._poll_only = False
+
+        # Test initial polling
+        self.subject._cached_state = {}
+        self.assertTrue(self.subject.should_poll)