From 3aa51397b8c163cb0860e088095d682e4a266a85 Mon Sep 17 00:00:00 2001 From: Alex X Date: Mon, 1 Apr 2024 17:21:14 +0300 Subject: [PATCH] Restore lazy entities after Hass restart --- .../xiaomi_gateway3/hass/add_entitites.py | 44 ++++++++++++++----- 1 file changed, 32 insertions(+), 12 deletions(-) diff --git a/custom_components/xiaomi_gateway3/hass/add_entitites.py b/custom_components/xiaomi_gateway3/hass/add_entitites.py index 3c29bcfd..22fd4160 100644 --- a/custom_components/xiaomi_gateway3/hass/add_entitites.py +++ b/custom_components/xiaomi_gateway3/hass/add_entitites.py @@ -24,6 +24,9 @@ def add_device(device: XDevice): return if device.did not in CONFIG_ENTRIES: + # connect all device entities to this gateway + CONFIG_ENTRIES[device.did] = gw + # instant setup all entities, except lazy for entity in get_entities(device, gw.stats_domain): gw.debug("add_entity", device=device, entity=entity.entity_id) @@ -32,9 +35,6 @@ def add_device(device: XDevice): # add listener for setup lazy entities (if device has them) if remove_listener := handle_lazy_entities(hass, config_entry, device): lazy_listeners[device.did] = remove_listener - - # connect all device entities to this gateway - CONFIG_ENTRIES[device.did] = gw else: # device already added to another config entry (gateway) # so we add device to the current config entry @@ -109,28 +109,48 @@ def handle_lazy_entities( hass: HomeAssistant, config_entry: ConfigEntry, device: XDevice ): """Create entities only when first data arrived.""" + # 1. Check if device has lazy entities lazy_attrs = { i.attr for i in device.converters if i.entity and i.entity.get("lazy") } + # 2. Exit if none if not lazy_attrs: return None - def on_device_update(data: dict): - for attr in data.keys() & lazy_attrs: - lazy_attrs.remove(attr) + def add_lazy_entity(attr: str) -> XEntity: + lazy_attrs.remove(attr) - conv = next(i for i in device.converters if i.attr == attr) - entity = create_entity(device, conv) - entity.on_device_update(data) + conv = next(i for i in device.converters if i.attr == attr) + entity = create_entity(device, conv) + + gw = CONFIG_ENTRIES.get(device.did) + gw.debug("add_lazy_entity", device=device, entity=entity.entity_id) + add_entity(hass, config_entry, entity) + return entity - gw = CONFIG_ENTRIES.get(device.did) - gw.debug("add_lazy_entity", device=device, entity=entity.entity_id) + # 3. Restore previous lazy entities from Hass entity registry + prefix = device.uid + "_" + reg = entity_registry.async_get(hass) + for entry in reg.entities.values(): + if entry.platform != DOMAIN or not entry.unique_id.startswith(prefix): + continue + _, attr = entry.unique_id.split("_", 1) + if attr in lazy_attrs: + add_lazy_entity(attr) + + # 4. Exit if none left + if not lazy_attrs: + return None - add_entity(hass, config_entry, entity) + def on_device_update(data: dict): + for attr in data.keys() & lazy_attrs: + entity = add_lazy_entity(attr) + entity.on_device_update(data) if not lazy_attrs: device.remove_listener(on_device_update) + # 5. Wait for rest lazy entities in every message from the device device.add_listener(on_device_update) return lambda: device.remove_listener(on_device_update)