Passively discover devices and update IP addresses

custom_components/localtuya/__init__.py

@@ -63,13 +63,16 @@ from homeassistant.const import (
     CONF_DEVICE_ID,
     CONF_PLATFORM,
     CONF_ENTITIES,
+    CONF_HOST,
+    EVENT_HOMEASSISTANT_STOP,
     SERVICE_RELOAD,
 )
 from homeassistant.helpers.reload import async_integration_yaml_config
 
-from .const import DOMAIN, TUYA_DEVICE
+from .const import DATA_DISCOVERY, DOMAIN, TUYA_DEVICE
 from .config_flow import config_schema
 from .common import TuyaDevice
+from .discovery import TuyaDiscovery
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -92,6 +95,8 @@ async def async_setup(hass: HomeAssistant, config: dict):
     """Set up the LocalTuya integration component."""
     hass.data.setdefault(DOMAIN, {})
 
+    device_cache = {}
+
     async def _handle_reload(service):
         """Handle reload service call."""
         config = await async_integration_yaml_config(hass, DOMAIN)
@@ -112,6 +117,51 @@ async def async_setup(hass: HomeAssistant, config: dict):
 
         await asyncio.gather(*reload_tasks)
 
+    def _entry_by_device_id(device_id):
+        """Look up config entry by device id."""
+        current_entries = hass.config_entries.async_entries(DOMAIN)
+        for entry in current_entries:
+            if entry.data[CONF_DEVICE_ID] == device_id:
+                return entry
+        return None
+
+    def _device_discovered(device):
+        """Update address of device if it has changed."""
+        device_ip = device["ip"]
+        device_id = device["gwId"]
+
+        # If device is not in cache, check if a config entry exists
+        if device_id not in device_cache:
+            entry = _entry_by_device_id(device_id)
+            if entry:
+                # Save address from config entry in cache to trigger
+                # potential update below
+                device_cache[device_id] = entry.data[CONF_HOST]
+
+        # If device is in cache and address changed...
+        if device_id in device_cache and device_cache[device_id] != device_ip:
+            _LOGGER.debug("Device %s changed IP to %s", device_id, device_ip)
+
+            entry = _entry_by_device_id(device_id)
+            if entry:
+                hass.config_entries.async_update_entry(
+                    entry, data={**entry.data, CONF_HOST: device_ip}
+                )
+                device_cache[device_id] = device_ip
+
+    discovery = TuyaDiscovery(_device_discovered)
+
+    def _shutdown(event):
+        """Clean up resources when shutting down."""
+        discovery.close()
+
+    try:
+        await discovery.start()
+        hass.data[DOMAIN][DATA_DISCOVERY] = discovery
+        hass.bus.async_listen_once(EVENT_HOMEASSISTANT_STOP, _shutdown)
+    except Exception:
+        _LOGGER.exception("failed to set up discovery")
+
     hass.helpers.service.async_register_admin_service(
         DOMAIN,
         SERVICE_RELOAD,

custom_components/localtuya/config_flow.py

@@ -21,15 +21,13 @@ from .const import (  # pylint: disable=unused-import
     CONF_LOCAL_KEY,
     CONF_PROTOCOL_VERSION,
     CONF_DPS_STRINGS,
+    DATA_DISCOVERY,
     DOMAIN,
     PLATFORMS,
 )
-from .discovery import discover
 
 _LOGGER = logging.getLogger(__name__)
 
-DISCOVER_TIMEOUT = 6.0
-
 PLATFORM_TO_ADD = "platform_to_add"
 NO_ADDITIONAL_PLATFORMS = "no_additional_platforms"
 DISCOVERED_DEVICE = "discovered_device"
@@ -211,16 +209,12 @@ class LocaltuyaConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
                 self.selected_device = user_input[DISCOVERED_DEVICE].split(" ")[0]
             return await self.async_step_basic_info()
 
-        try:
-            devices = await discover(DISCOVER_TIMEOUT, self.hass.loop)
-            self.devices = {
-                ip: dev
-                for ip, dev in devices.items()
-                if dev["gwId"] not in self._async_current_ids()
-            }
-        except Exception:  # pylint: disable=broad-except
-            _LOGGER.exception("discovery failed")
-            errors["base"] = "discovery_failed"
+        devices = self.hass.data[DOMAIN][DATA_DISCOVERY].devices
+        self.devices = {
+            ip: dev
+            for ip, dev in devices.items()
+            if dev["gwId"] not in self._async_current_ids()
+        }
 
         return self.async_show_form(
             step_id="user", errors=errors, data_schema=user_schema(self.devices)

custom_components/localtuya/const.py

@@ -28,6 +28,8 @@ CONF_SPAN_TIME = "span_time"
 # sensor
 CONF_SCALING = "scaling"
 
+DATA_DISCOVERY = "discovery"
+
 DOMAIN = "localtuya"
 
 # Platforms in this list must support config flows

custom_components/localtuya/discovery.py

@@ -31,9 +31,29 @@ def decrypt_udp(message):
 class TuyaDiscovery(asyncio.DatagramProtocol):
     """Datagram handler listening for Tuya broadcast messages."""
 
-    def __init__(self, found_devices):
-        """Initialize a new TuyaDiscovery instance."""
-        self.found_devices = found_devices
+    def __init__(self, callback):
+        """Initialize a new BaseDiscovery."""
+        self.devices = {}
+        self._listeners = []
+        self._callback = callback
+
+    async def start(self):
+        """Start discovery by listening to broadcasts."""
+        loop = asyncio.get_running_loop()
+        listener = loop.create_datagram_endpoint(
+            lambda: self, local_addr=("0.0.0.0", 6666)
+        )
+        encrypted_listener = loop.create_datagram_endpoint(
+            lambda: self, local_addr=("0.0.0.0", 6667)
+        )
+
+        self.listeners = await asyncio.gather(listener, encrypted_listener)
+        _LOGGER.debug("Listening to broadcasts on UDP port 6666 and 6667")
+
+    def close(self):
+        """Stop discovery."""
+        for transport, _ in self.listeners:
+            transport.close()
 
     def datagram_received(self, data, addr):
         """Handle received broadcast message."""
@@ -44,43 +64,12 @@ class TuyaDiscovery(asyncio.DatagramProtocol):
             data = data.decode()
 
         decoded = json.loads(data)
-        if decoded.get("ip") not in self.found_devices:
-            self.found_devices[decoded.get("ip")] = decoded
-            _LOGGER.debug("Discovered device: %s", decoded)
+        self.device_found(decoded)
 
+    def device_found(self, device):
+        """Discover a new device."""
+        if device.get("ip") not in self.devices:
+            self.devices[device.get("ip")] = device
+            _LOGGER.debug("Discovered device: %s", device)
 
-async def discover(timeout, loop):
-    """Discover and return Tuya devices on the network."""
-    found_devices = {}
-
-    def proto_factory():
-        return TuyaDiscovery(found_devices)
-
-    listener = loop.create_datagram_endpoint(
-        proto_factory, local_addr=("0.0.0.0", 6666)
-    )
-    encrypted_listener = loop.create_datagram_endpoint(
-        proto_factory, local_addr=("0.0.0.0", 6667)
-    )
-
-    listeners = await asyncio.gather(listener, encrypted_listener)
-    _LOGGER.debug("Listening to broadcasts on UDP port 6666 and 6667")
-
-    try:
-        await asyncio.sleep(timeout)
-    finally:
-        for transport, _ in listeners:
-            transport.close()
-
-    return found_devices
-
-
-def main():
-    """Run discovery and print result."""
-    loop = asyncio.get_event_loop()
-    res = loop.run_until_complete(discover(5, loop))
-    print(res)
-
-
-if __name__ == "__main__":
-    main()
+        self._callback(device)

custom_components/localtuya/translations/en.json

@@ -7,8 +7,7 @@
             "cannot_connect": "Cannot connect to device. Verify that address is correct and try again.",
             "invalid_auth": "Failed to authenticate with device. Verify that device id and local key are correct.",
             "unknown": "An unknown error occurred. See log for details.",
-            "entity_already_configured": "Entity with this ID has already been configured.",
-            "discovery_failed": "Failed to discover devices. You can still add a device manually."
+            "entity_already_configured": "Entity with this ID has already been configured."
         },
         "step": {
             "user": {