Merge pull request #11 from andrew-codechimp/device-association

Device association

Author: andrew-codechimp

custom_components/periodic_min_max/sensor.py

@@ -5,6 +5,7 @@
 from datetime import datetime
 from typing import Any
 
+import voluptuous as vol
 from homeassistant.components.sensor import (
     SensorDeviceClass,
     SensorEntity,
@@ -20,14 +21,18 @@
     STATE_UNKNOWN,
 )
 from homeassistant.core import Event, EventStateChangedData, HomeAssistant, callback
+from homeassistant.helpers import device_registry as dr
 from homeassistant.helpers import entity_platform
 from homeassistant.helpers import entity_registry as er
 from homeassistant.helpers.device import async_device_info_to_link_from_entity
 from homeassistant.helpers.entity_platform import (
     AddConfigEntryEntitiesCallback,
     AddEntitiesCallback,
 )
-from homeassistant.helpers.event import async_track_state_change_event
+from homeassistant.helpers.event import (
+    async_track_entity_registry_updated_event,
+    async_track_state_change_event,
+)
 from homeassistant.helpers.reload import async_setup_reload_service
 from homeassistant.helpers.restore_state import RestoreEntity
 from homeassistant.helpers.typing import ConfigType, DiscoveryInfoType, StateType
@@ -49,18 +54,95 @@
 SENSOR_TYPE_TO_ATTR = {v: k for k, v in SENSOR_TYPES.items()}
 
 
+@callback
+def async_add_to_device(
+    hass: HomeAssistant, entry: ConfigEntry, entity_id: str
+) -> str | None:
+    """Add our config entry to the tracked entity's device."""
+    registry = er.async_get(hass)
+    device_registry = dr.async_get(hass)
+    device_id = None
+
+    if (
+        not (source_device := registry.async_get(entity_id))
+        or not (device_id := source_device.device_id)
+        or not (device_registry.async_get(device_id))
+    ):
+        return device_id
+
+    device_registry.async_update_device(device_id, add_config_entry_id=entry.entry_id)
+
+    return device_id
+
+
+async def config_entry_update_listener(hass: HomeAssistant, entry: ConfigEntry) -> None:
+    """Update listener, called when the config entry options are changed."""
+    await hass.config_entries.async_reload(entry.entry_id)
+
+
 async def async_setup_entry(
     hass: HomeAssistant,
     config_entry: ConfigEntry,
     async_add_entities: AddConfigEntryEntitiesCallback,
 ) -> None:
     """Initialize min/max/mean config entry."""
-    registry = er.async_get(hass)
-    entity_id = er.async_validate_entity_id(
-        registry, config_entry.options[CONF_ENTITY_ID]
-    )
+    entity_registry = er.async_get(hass)
+    device_registry = dr.async_get(hass)
+    try:
+        entity_id = er.async_validate_entity_id(
+            entity_registry, config_entry.options[CONF_ENTITY_ID]
+        )
+    except vol.Invalid:
+        # The entity is identified by an unknown entity registry ID
+        LOGGER.error(
+            "Failed to setup periodic_min_max for unknown entity %s",
+            config_entry.options[CONF_ENTITY_ID],
+        )
+        return False
+
     sensor_type = config_entry.options[CONF_TYPE]
 
+    async def async_registry_updated(
+        event: Event[er.EventEntityRegistryUpdatedData],
+    ) -> None:
+        """Handle entity registry update."""
+        data = event.data
+        if data["action"] == "remove":
+            await hass.config_entries.async_remove(config_entry.entry_id)
+
+        if data["action"] != "update":
+            return
+
+        if "entity_id" in data["changes"]:
+            # Entity_id changed, reload the config entry
+            await hass.config_entries.async_reload(config_entry.entry_id)
+
+        if device_id and "device_id" in data["changes"]:
+            # If the tracked entity is no longer in the device, remove our config entry
+            # from the device
+            if (
+                not (entity_entry := entity_registry.async_get(data[CONF_ENTITY_ID]))
+                or not device_registry.async_get(device_id)
+                or entity_entry.device_id == device_id
+            ):
+                # No need to do any cleanup
+                return
+
+            device_registry.async_update_device(
+                device_id, remove_config_entry_id=config_entry.entry_id
+            )
+
+    config_entry.async_on_unload(
+        async_track_entity_registry_updated_event(
+            hass, entity_id, async_registry_updated
+        )
+    )
+    config_entry.async_on_unload(
+        config_entry.add_update_listener(config_entry_update_listener)
+    )
+
+    device_id = async_add_to_device(hass, config_entry, entity_id)
+
     async_add_entities(
         [
             PeriodicMinMaxSensor(

scripts/develop

@@ -11,13 +11,10 @@ if [[ ! -d "${PWD}/config" ]]; then
 fi
 
 # Set the path to custom_components
-## This let's us have the structure we want <root>/custom_components/periodic_min_max
+## This let's us have the structure we want <root>/custom_components/integration_battery_types
 ## while at the same time have Home Assistant configuration inside <root>/config
 ## without resulting to symlinks.
 export PYTHONPATH="${PYTHONPATH}:${PWD}/custom_components"
 
 # Start Home Assistant
-# debugpy logging, if required: --log-to debugpy
-python3 -Xfrozen_modules=off \
-    -m debugpy --listen localhost:5678 \
-    -m homeassistant --config "${PWD}/config"
+hass --config "${PWD}/config" --debug