Add device linking and stale device link clean up helpers (#119761)
This commit is contained in:
parent
e2a34d209f
commit
a8ba22f6bb
6 changed files with 358 additions and 80 deletions
|
@ -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)
|
||||
|
|
|
@ -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,
|
||||
|
|
|
@ -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]
|
||||
|
|
75
homeassistant/helpers/device.py
Normal file
75
homeassistant/helpers/device.py
Normal 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)
|
56
tests/components/utility_meter/test_select.py
Normal file
56
tests/components/utility_meter/test_select.py
Normal 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
|
211
tests/helpers/test_device.py
Normal file
211
tests/helpers/test_device.py
Normal 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
|
Loading…
Add table
Reference in a new issue