Add device linking and stale device link clean up helpers (#119761)

This commit is contained in:
dougiteixeira 2024-06-21 09:46:39 -03:00 committed by GitHub
parent e2a34d209f
commit a8ba22f6bb
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 358 additions and 80 deletions

View file

@ -11,12 +11,11 @@ from homeassistant.components.sensor import DOMAIN as SENSOR_DOMAIN
from homeassistant.config_entries import ConfigEntry
from homeassistant.const import ATTR_ENTITY_ID, CONF_NAME, CONF_UNIQUE_ID, Platform
from homeassistant.core import HomeAssistant, split_entity_id
from homeassistant.helpers import (
device_registry as dr,
discovery,
entity_registry as er,
)
from homeassistant.helpers import discovery, entity_registry as er
import homeassistant.helpers.config_validation as cv
from homeassistant.helpers.device import (
async_remove_stale_devices_links_keep_entity_device,
)
from homeassistant.helpers.dispatcher import async_dispatcher_send
from homeassistant.helpers.typing import ConfigType
@ -192,7 +191,7 @@ async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
"""Set up Utility Meter from a config entry."""
await async_remove_stale_device_links(
async_remove_stale_devices_links_keep_entity_device(
hass, entry.entry_id, entry.options[CONF_SOURCE_SENSOR]
)
@ -266,27 +265,3 @@ async def async_migrate_entry(hass: HomeAssistant, config_entry: ConfigEntry) ->
_LOGGER.info("Migration to version %s successful", config_entry.version)
return True
async def async_remove_stale_device_links(
hass: HomeAssistant, entry_id: str, entity_id: str
) -> None:
"""Remove device link for entry, the source device may have changed."""
device_registry = dr.async_get(hass)
entity_registry = er.async_get(hass)
# Resolve source entity device
current_device_id = None
if ((source_entity := entity_registry.async_get(entity_id)) is not None) and (
source_entity.device_id is not None
):
current_device_id = source_entity.device_id
devices_in_entry = device_registry.devices.get_devices_for_config_entry_id(entry_id)
# Removes all devices from the config entry that are not the same as the current device
for device in devices_in_entry:
if device.id == current_device_id:
continue
device_registry.async_update_device(device.id, remove_config_entry_id=entry_id)

View file

@ -8,7 +8,7 @@ from homeassistant.components.select import SelectEntity
from homeassistant.config_entries import ConfigEntry
from homeassistant.const import CONF_UNIQUE_ID
from homeassistant.core import HomeAssistant
from homeassistant.helpers import device_registry as dr, entity_registry as er
from homeassistant.helpers.device import async_device_info_to_link_from_entity
from homeassistant.helpers.device_registry import DeviceInfo
from homeassistant.helpers.entity_platform import AddEntitiesCallback
from homeassistant.helpers.restore_state import RestoreEntity
@ -30,28 +30,10 @@ async def async_setup_entry(
unique_id = config_entry.entry_id
registry = er.async_get(hass)
source_entity = registry.async_get(config_entry.options[CONF_SOURCE_SENSOR])
dev_reg = dr.async_get(hass)
# Resolve source entity device
if (
(source_entity is not None)
and (source_entity.device_id is not None)
and (
(
device := dev_reg.async_get(
device_id=source_entity.device_id,
device_info = async_device_info_to_link_from_entity(
hass,
config_entry.options[CONF_SOURCE_SENSOR],
)
)
is not None
)
):
device_info = DeviceInfo(
identifiers=device.identifiers,
connections=device.connections,
)
else:
device_info = None
tariff_select = TariffSelect(
name,

View file

@ -37,12 +37,8 @@ from homeassistant.core import (
State,
callback,
)
from homeassistant.helpers import (
device_registry as dr,
entity_platform,
entity_registry as er,
)
from homeassistant.helpers.device_registry import DeviceInfo
from homeassistant.helpers import entity_platform, entity_registry as er
from homeassistant.helpers.device import async_device_info_to_link_from_entity
from homeassistant.helpers.dispatcher import async_dispatcher_connect
from homeassistant.helpers.entity_platform import AddEntitiesCallback
from homeassistant.helpers.event import (
@ -130,27 +126,10 @@ async def async_setup_entry(
registry, config_entry.options[CONF_SOURCE_SENSOR]
)
source_entity = registry.async_get(source_entity_id)
dev_reg = dr.async_get(hass)
# Resolve source entity device
if (
(source_entity is not None)
and (source_entity.device_id is not None)
and (
(
device := dev_reg.async_get(
device_id=source_entity.device_id,
device_info = async_device_info_to_link_from_entity(
hass,
source_entity_id,
)
)
is not None
)
):
device_info = DeviceInfo(
identifiers=device.identifiers,
connections=device.connections,
)
else:
device_info = None
cron_pattern = None
delta_values = config_entry.options[CONF_METER_DELTA_VALUES]

View file

@ -0,0 +1,75 @@
"""Provides useful helpers for handling devices."""
from homeassistant.core import HomeAssistant, callback
from . import device_registry as dr, entity_registry as er
@callback
def async_entity_id_to_device_id(
hass: HomeAssistant,
entity_id_or_uuid: str,
) -> str | None:
"""Resolve the device id to the entity id or entity uuid."""
ent_reg = er.async_get(hass)
entity_id = er.async_validate_entity_id(ent_reg, entity_id_or_uuid)
if (entity := ent_reg.async_get(entity_id)) is None:
return None
return entity.device_id
@callback
def async_device_info_to_link_from_entity(
hass: HomeAssistant,
entity_id_or_uuid: str,
) -> dr.DeviceInfo | None:
"""DeviceInfo with information to link a device to a configuration entry in the link category from a entity id or entity uuid."""
dev_reg = dr.async_get(hass)
if (device_id := async_entity_id_to_device_id(hass, entity_id_or_uuid)) is None or (
device := dev_reg.async_get(device_id=device_id)
) is None:
return None
return dr.DeviceInfo(
identifiers=device.identifiers,
connections=device.connections,
)
@callback
def async_remove_stale_devices_links_keep_entity_device(
hass: HomeAssistant,
entry_id: str,
source_entity_id_or_uuid: str,
) -> None:
"""Remove the link between stales devices and a configuration entry, keeping only the device that the informed entity is linked to."""
async_remove_stale_devices_links_keep_current_device(
hass=hass,
entry_id=entry_id,
current_device_id=async_entity_id_to_device_id(hass, source_entity_id_or_uuid),
)
@callback
def async_remove_stale_devices_links_keep_current_device(
hass: HomeAssistant,
entry_id: str,
current_device_id: str | None,
) -> None:
"""Remove the link between stales devices and a configuration entry, keeping only the device informed.
Device passed in the current_device_id parameter will be kept linked to the configuration entry.
"""
dev_reg = dr.async_get(hass)
# Removes all devices from the config entry that are not the same as the current device
for device in dev_reg.devices.get_devices_for_config_entry_id(entry_id):
if device.id == current_device_id:
continue
dev_reg.async_update_device(device.id, remove_config_entry_id=entry_id)

View file

@ -0,0 +1,56 @@
"""The tests for the utility_meter select platform."""
from homeassistant.components.utility_meter.const import DOMAIN
from homeassistant.core import HomeAssistant
from homeassistant.helpers import device_registry as dr, entity_registry as er
from tests.common import MockConfigEntry
async def test_device_id(
hass: HomeAssistant,
device_registry: dr.DeviceRegistry,
entity_registry: er.EntityRegistry,
) -> None:
"""Test for source entity device for Utility Meter."""
source_config_entry = MockConfigEntry()
source_config_entry.add_to_hass(hass)
source_device_entry = device_registry.async_get_or_create(
config_entry_id=source_config_entry.entry_id,
identifiers={("sensor", "identifier_test")},
connections={("mac", "30:31:32:33:34:35")},
)
source_entity = entity_registry.async_get_or_create(
"sensor",
"test",
"source",
config_entry=source_config_entry,
device_id=source_device_entry.id,
)
await hass.async_block_till_done()
assert entity_registry.async_get("sensor.test_source") is not None
utility_meter_config_entry = MockConfigEntry(
data={},
domain=DOMAIN,
options={
"cycle": "monthly",
"delta_values": False,
"name": "Energy",
"net_consumption": False,
"offset": 0,
"periodically_resetting": True,
"source": "sensor.test_source",
"tariffs": ["peak", "offpeak"],
},
title="Energy",
)
utility_meter_config_entry.add_to_hass(hass)
assert await hass.config_entries.async_setup(utility_meter_config_entry.entry_id)
await hass.async_block_till_done()
utility_meter_entity_select = entity_registry.async_get("select.energy")
assert utility_meter_entity_select is not None
assert utility_meter_entity_select.device_id == source_entity.device_id

View file

@ -0,0 +1,211 @@
"""Tests for the Device Utils."""
import pytest
import voluptuous as vol
from homeassistant.core import HomeAssistant
from homeassistant.helpers import device_registry as dr, entity_registry as er
from homeassistant.helpers.device import (
async_device_info_to_link_from_entity,
async_entity_id_to_device_id,
async_remove_stale_devices_links_keep_current_device,
async_remove_stale_devices_links_keep_entity_device,
)
from tests.common import MockConfigEntry
async def test_entity_id_to_device_id(
hass: HomeAssistant,
device_registry: dr.DeviceRegistry,
entity_registry: er.EntityRegistry,
) -> None:
"""Test returning an entity's device ID."""
config_entry = MockConfigEntry(domain="my")
config_entry.add_to_hass(hass)
device = device_registry.async_get_or_create(
identifiers={("test", "current_device")},
connections={("mac", "30:31:32:33:34:00")},
config_entry_id=config_entry.entry_id,
)
assert device is not None
# Entity registry
entity = entity_registry.async_get_or_create(
"sensor",
"test",
"source",
config_entry=config_entry,
device_id=device.id,
)
await hass.async_block_till_done()
assert entity_registry.async_get("sensor.test_source") is not None
device_id = async_entity_id_to_device_id(
hass,
entity_id_or_uuid=entity.entity_id,
)
assert device_id == device.id
with pytest.raises(vol.Invalid):
async_entity_id_to_device_id(
hass,
entity_id_or_uuid="unknown_uuid",
)
async def test_device_info_to_link(
hass: HomeAssistant,
device_registry: dr.DeviceRegistry,
entity_registry: er.EntityRegistry,
) -> None:
"""Test for returning device info with device link information."""
config_entry = MockConfigEntry(domain="my")
config_entry.add_to_hass(hass)
device = device_registry.async_get_or_create(
identifiers={("test", "my_device")},
connections={("mac", "30:31:32:33:34:00")},
config_entry_id=config_entry.entry_id,
)
assert device is not None
# Source entity registry
source_entity = entity_registry.async_get_or_create(
"sensor",
"test",
"source",
config_entry=config_entry,
device_id=device.id,
)
await hass.async_block_till_done()
assert entity_registry.async_get("sensor.test_source") is not None
result = async_device_info_to_link_from_entity(
hass, entity_id_or_uuid=source_entity.entity_id
)
assert result == {
"identifiers": {("test", "my_device")},
"connections": {("mac", "30:31:32:33:34:00")},
}
# With a non-existent entity id
result = async_device_info_to_link_from_entity(
hass, entity_id_or_uuid="sensor.invalid"
)
assert result is None
async def test_remove_stale_device_links_keep_entity_device(
hass: HomeAssistant,
device_registry: dr.DeviceRegistry,
entity_registry: er.EntityRegistry,
) -> None:
"""Test cleaning works for entity."""
config_entry = MockConfigEntry(domain="hue")
config_entry.add_to_hass(hass)
current_device = device_registry.async_get_or_create(
identifiers={("test", "current_device")},
connections={("mac", "30:31:32:33:34:00")},
config_entry_id=config_entry.entry_id,
)
assert current_device is not None
device_registry.async_get_or_create(
identifiers={("test", "stale_device_1")},
connections={("mac", "30:31:32:33:34:01")},
config_entry_id=config_entry.entry_id,
)
device_registry.async_get_or_create(
identifiers={("test", "stale_device_2")},
connections={("mac", "30:31:32:33:34:02")},
config_entry_id=config_entry.entry_id,
)
# Source entity registry
source_entity = entity_registry.async_get_or_create(
"sensor",
"test",
"source",
config_entry=config_entry,
device_id=current_device.id,
)
await hass.async_block_till_done()
assert entity_registry.async_get("sensor.test_source") is not None
devices_config_entry = device_registry.devices.get_devices_for_config_entry_id(
config_entry.entry_id
)
# 3 devices linked to the config entry are expected (1 current device + 2 stales)
assert len(devices_config_entry) == 3
# Manual cleanup should unlink stales devices from the config entry
async_remove_stale_devices_links_keep_entity_device(
hass,
entry_id=config_entry.entry_id,
source_entity_id_or_uuid=source_entity.entity_id,
)
devices_config_entry = device_registry.devices.get_devices_for_config_entry_id(
config_entry.entry_id
)
# After cleanup, only one device is expected to be linked to the configuration entry if at least source_entity_id_or_uuid or device_id was given, else zero
assert len(devices_config_entry) == 1
assert current_device in devices_config_entry
async def test_remove_stale_devices_links_keep_current_device(
hass: HomeAssistant,
device_registry: dr.DeviceRegistry,
) -> None:
"""Test cleanup works for device id."""
config_entry = MockConfigEntry(domain="hue")
config_entry.add_to_hass(hass)
current_device = device_registry.async_get_or_create(
identifiers={("test", "current_device")},
connections={("mac", "30:31:32:33:34:00")},
config_entry_id=config_entry.entry_id,
)
assert current_device is not None
device_registry.async_get_or_create(
identifiers={("test", "stale_device_1")},
connections={("mac", "30:31:32:33:34:01")},
config_entry_id=config_entry.entry_id,
)
device_registry.async_get_or_create(
identifiers={("test", "stale_device_2")},
connections={("mac", "30:31:32:33:34:02")},
config_entry_id=config_entry.entry_id,
)
devices_config_entry = device_registry.devices.get_devices_for_config_entry_id(
config_entry.entry_id
)
# 3 devices linked to the config entry are expected (1 current device + 2 stales)
assert len(devices_config_entry) == 3
# Manual cleanup should unlink stales devices from the config entry
async_remove_stale_devices_links_keep_current_device(
hass,
entry_id=config_entry.entry_id,
current_device_id=current_device.id,
)
devices_config_entry = device_registry.devices.get_devices_for_config_entry_id(
config_entry.entry_id
)
# After cleanup, only one device is expected to be linked to the configuration entry
assert len(devices_config_entry) == 1
assert current_device in devices_config_entry