diff --git a/custom_components/localtuya/config_flow.py b/custom_components/localtuya/config_flow.py index 5247c10..dea4d94 100644 --- a/custom_components/localtuya/config_flow.py +++ b/custom_components/localtuya/config_flow.py @@ -25,6 +25,7 @@ from .const import ( # pylint: disable=unused-import DOMAIN, PLATFORMS, ) +from .discovery import discover _LOGGER = logging.getLogger(__name__) @@ -209,7 +210,11 @@ class LocaltuyaConfigFlow(config_entries.ConfigFlow, domain=DOMAIN): self.selected_device = user_input[DISCOVERED_DEVICE].split(" ")[0] return await self.async_step_basic_info() - devices = self.hass.data[DOMAIN][DATA_DISCOVERY].devices + # Use cache if available or fallback to manual discovery + if DOMAIN in self.hass.data: + devices = self.hass.data[DOMAIN][DATA_DISCOVERY].devices + else: + devices = await discover() self.devices = { ip: dev for ip, dev in devices.items() diff --git a/custom_components/localtuya/discovery.py b/custom_components/localtuya/discovery.py index a75cc22..b019aad 100644 --- a/custom_components/localtuya/discovery.py +++ b/custom_components/localtuya/discovery.py @@ -16,6 +16,8 @@ _LOGGER = logging.getLogger(__name__) UDP_KEY = md5(b"yGAdlopoPVldABfn").digest() +DEFAULT_TIMEOUT = 6.0 + def decrypt_udp(message): """Decrypt encrypted UDP broadcasts.""" @@ -31,7 +33,7 @@ def decrypt_udp(message): class TuyaDiscovery(asyncio.DatagramProtocol): """Datagram handler listening for Tuya broadcast messages.""" - def __init__(self, callback): + def __init__(self, callback=None): """Initialize a new BaseDiscovery.""" self.devices = {} self._listeners = [] @@ -52,6 +54,7 @@ class TuyaDiscovery(asyncio.DatagramProtocol): def close(self): """Stop discovery.""" + self.callback = None for transport, _ in self.listeners: transport.close() @@ -72,4 +75,18 @@ class TuyaDiscovery(asyncio.DatagramProtocol): self.devices[device.get("ip")] = device _LOGGER.debug("Discovered device: %s", device) - self._callback(device) + if self._callback: + self._callback(device) + + +async def discover(): + """Discover and return devices on local network.""" + discover = TuyaDiscovery() + try: + await discover.start() + await asyncio.sleep(DEFAULT_TIMEOUT) + except Exception: + _LOGGER.exception("failed to discover devices") + finally: + discover.close() + return discover.devices