mirror of
https://github.com/home-assistant/core.git
synced 2025-07-21 12:17:07 +00:00
Handle changes to source entity in derivative helper (#146407)
* Handle changes to source entity in derivative helper * Rename helper function, improve docstring * Add tests * Improve derivative tests * Deduplicate tests * Rename helpers/helper_entity.py to helpers/helper_integration.py * Rename tests
This commit is contained in:
Erik Montnemery
GitHub
parent
4fdbb9c0e2
commit
16a0b7f44e
@ -2,12 +2,18 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from homeassistant.components.sensor import DOMAIN as SENSOR_DOMAIN
|
||||
from homeassistant.config_entries import ConfigEntry
|
||||
from homeassistant.const import CONF_SOURCE, Platform
|
||||
from homeassistant.core import HomeAssistant
|
||||
from homeassistant.helpers import entity_registry as er
|
||||
from homeassistant.helpers.device import (
|
||||
async_entity_id_to_device_id,
|
||||
async_remove_stale_devices_links_keep_entity_device,
|
||||
)
|
||||
from homeassistant.helpers.helper_integration import async_handle_source_entity_changes
|
||||
|
||||
from .const import DOMAIN
|
||||
|
||||
|
||||
async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
|
||||
@ -17,6 +23,27 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
|
||||
hass, entry.entry_id, entry.options[CONF_SOURCE]
|
||||
)
|
||||
|
||||
def set_source_entity_id_or_uuid(source_entity_id: str) -> None:
|
||||
hass.config_entries.async_update_entry(
|
||||
entry,
|
||||
options={**entry.options, CONF_SOURCE: source_entity_id},
|
||||
)
|
||||
|
||||
entity_registry = er.async_get(hass)
|
||||
entry.async_on_unload(
|
||||
async_handle_source_entity_changes(
|
||||
hass,
|
||||
helper_config_entry_id=entry.entry_id,
|
||||
get_helper_entity_id=lambda: entity_registry.async_get_entity_id(
|
||||
SENSOR_DOMAIN, DOMAIN, entry.entry_id
|
||||
),
|
||||
set_source_entity_id_or_uuid=set_source_entity_id_or_uuid,
|
||||
source_device_id=async_entity_id_to_device_id(
|
||||
hass, entry.options[CONF_SOURCE]
|
||||
),
|
||||
source_entity_id_or_uuid=entry.options[CONF_SOURCE],
|
||||
)
|
||||
)
|
||||
await hass.config_entries.async_forward_entry_setups(entry, (Platform.SENSOR,))
|
||||
entry.async_on_unload(entry.add_update_listener(config_entry_update_listener))
|
||||
return True
|
||||
|
||||
@ -9,9 +9,9 @@ import voluptuous as vol
|
||||
from homeassistant.components.homeassistant import exposed_entities
|
||||
from homeassistant.config_entries import ConfigEntry
|
||||
from homeassistant.const import CONF_ENTITY_ID
|
||||
from homeassistant.core import Event, HomeAssistant, callback, valid_entity_id
|
||||
from homeassistant.core import HomeAssistant, callback
|
||||
from homeassistant.helpers import device_registry as dr, entity_registry as er
|
||||
from homeassistant.helpers.event import async_track_entity_registry_updated_event
|
||||
from homeassistant.helpers.helper_integration import async_handle_source_entity_changes
|
||||
|
||||
from .const import CONF_INVERT, CONF_TARGET_DOMAIN, DOMAIN
|
||||
|
||||
@ -42,7 +42,6 @@ def async_add_to_device(
|
||||
async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
|
||||
"""Set up a config entry."""
|
||||
entity_registry = er.async_get(hass)
|
||||
device_registry = dr.async_get(hass)
|
||||
try:
|
||||
entity_id = er.async_validate_entity_id(
|
||||
entity_registry, entry.options[CONF_ENTITY_ID]
|
||||
@ -55,72 +54,26 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
|
||||
)
|
||||
return False
|
||||
|
||||
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(entry.entry_id)
|
||||
|
||||
if data["action"] != "update":
|
||||
return
|
||||
|
||||
if "entity_id" in data["changes"]:
|
||||
# Entity_id changed, update or reload the config entry
|
||||
if valid_entity_id(entry.options[CONF_ENTITY_ID]):
|
||||
# If the entity is pointed to by an entity ID, update the entry
|
||||
hass.config_entries.async_update_entry(
|
||||
entry,
|
||||
options={**entry.options, CONF_ENTITY_ID: data["entity_id"]},
|
||||
)
|
||||
else:
|
||||
await hass.config_entries.async_reload(entry.entry_id)
|
||||
|
||||
if device_id and "device_id" in data["changes"]:
|
||||
# Handle the wrapped switch being moved to a different device or removed
|
||||
# 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
|
||||
|
||||
# The wrapped switch has been moved to a different device, update the
|
||||
# switch_as_x entity and the device entry to include our config entry
|
||||
switch_as_x_entity_id = entity_registry.async_get_entity_id(
|
||||
entry.options[CONF_TARGET_DOMAIN], DOMAIN, entry.entry_id
|
||||
)
|
||||
if switch_as_x_entity_id:
|
||||
# Update the switch_as_x entity to point to the new device (or no device)
|
||||
entity_registry.async_update_entity(
|
||||
switch_as_x_entity_id, device_id=entity_entry.device_id
|
||||
)
|
||||
|
||||
if entity_entry.device_id is not None:
|
||||
device_registry.async_update_device(
|
||||
entity_entry.device_id, add_config_entry_id=entry.entry_id
|
||||
)
|
||||
|
||||
device_registry.async_update_device(
|
||||
device_id, remove_config_entry_id=entry.entry_id
|
||||
)
|
||||
|
||||
# Reload the config entry so the switch_as_x entity is recreated with
|
||||
# correct device info
|
||||
await hass.config_entries.async_reload(entry.entry_id)
|
||||
def set_source_entity_id_or_uuid(source_entity_id: str) -> None:
|
||||
hass.config_entries.async_update_entry(
|
||||
entry,
|
||||
options={**entry.options, CONF_ENTITY_ID: source_entity_id},
|
||||
)
|
||||
|
||||
entry.async_on_unload(
|
||||
async_track_entity_registry_updated_event(
|
||||
hass, entity_id, async_registry_updated
|
||||
async_handle_source_entity_changes(
|
||||
hass,
|
||||
helper_config_entry_id=entry.entry_id,
|
||||
get_helper_entity_id=lambda: entity_registry.async_get_entity_id(
|
||||
entry.options[CONF_TARGET_DOMAIN], DOMAIN, entry.entry_id
|
||||
),
|
||||
set_source_entity_id_or_uuid=set_source_entity_id_or_uuid,
|
||||
source_device_id=async_add_to_device(hass, entry, entity_id),
|
||||
source_entity_id_or_uuid=entry.options[CONF_ENTITY_ID],
|
||||
)
|
||||
)
|
||||
entry.async_on_unload(entry.add_update_listener(config_entry_update_listener))
|
||||
|
||||
device_id = async_add_to_device(hass, entry, entity_id)
|
||||
|
||||
await hass.config_entries.async_forward_entry_setups(
|
||||
entry, (entry.options[CONF_TARGET_DOMAIN],)
|
||||
)
|
||||
|
||||
105
homeassistant/helpers/helper_integration.py
Normal file
105
homeassistant/helpers/helper_integration.py
Normal file
@ -0,0 +1,105 @@
|
||||
"""Helpers for helper integrations."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Callable
|
||||
|
||||
from homeassistant.core import CALLBACK_TYPE, Event, HomeAssistant, valid_entity_id
|
||||
|
||||
from . import device_registry as dr, entity_registry as er
|
||||
from .event import async_track_entity_registry_updated_event
|
||||
|
||||
|
||||
def async_handle_source_entity_changes(
|
||||
hass: HomeAssistant,
|
||||
*,
|
||||
helper_config_entry_id: str,
|
||||
get_helper_entity_id: Callable[[], str | None],
|
||||
set_source_entity_id_or_uuid: Callable[[str], None],
|
||||
source_device_id: str | None,
|
||||
source_entity_id_or_uuid: str,
|
||||
) -> CALLBACK_TYPE:
|
||||
"""Handle changes to a helper entity's source entity.
|
||||
|
||||
The following changes are handled:
|
||||
- Entity removal: If the source entity is removed, the helper config entry
|
||||
is removed, and the helper entity is cleaned up.
|
||||
- Entity ID changed: If the source entity's entity ID changes and the source
|
||||
entity is identified by an entity ID, the set_source_entity_id_or_uuid is
|
||||
called. If the source entity is identified by a UUID, the helper config entry
|
||||
is reloaded.
|
||||
- Source entity moved to another device: The helper entity is updated to link
|
||||
to the new device, and the helper config entry removed from the old device
|
||||
and added to the new device. Then the helper config entry is reloaded.
|
||||
- Source entity removed from the device: The helper entity is updated to link
|
||||
to no device, and the helper config entry removed from the old device. Then
|
||||
the helper config entry is reloaded.
|
||||
"""
|
||||
|
||||
async def async_registry_updated(
|
||||
event: Event[er.EventEntityRegistryUpdatedData],
|
||||
) -> None:
|
||||
"""Handle entity registry update."""
|
||||
nonlocal source_device_id
|
||||
|
||||
data = event.data
|
||||
if data["action"] == "remove":
|
||||
await hass.config_entries.async_remove(helper_config_entry_id)
|
||||
|
||||
if data["action"] != "update":
|
||||
return
|
||||
|
||||
if "entity_id" in data["changes"]:
|
||||
# Entity_id changed, update or reload the config entry
|
||||
if valid_entity_id(source_entity_id_or_uuid):
|
||||
# If the entity is pointed to by an entity ID, update the entry
|
||||
set_source_entity_id_or_uuid(data["entity_id"])
|
||||
else:
|
||||
await hass.config_entries.async_reload(helper_config_entry_id)
|
||||
|
||||
if not source_device_id or "device_id" not in data["changes"]:
|
||||
return
|
||||
|
||||
# Handle the source entity being moved to a different device or removed
|
||||
# from the device
|
||||
if (
|
||||
not (source_entity_entry := entity_registry.async_get(data["entity_id"]))
|
||||
or not device_registry.async_get(source_device_id)
|
||||
or source_entity_entry.device_id == source_device_id
|
||||
):
|
||||
# No need to do any cleanup
|
||||
return
|
||||
|
||||
# The source entity has been moved to a different device, update the helper
|
||||
# helper entity to link to the new device and the helper device to include
|
||||
# the helper config entry
|
||||
helper_entity_id = get_helper_entity_id()
|
||||
if helper_entity_id:
|
||||
# Update the helper entity to link to the new device (or no device)
|
||||
entity_registry.async_update_entity(
|
||||
helper_entity_id, device_id=source_entity_entry.device_id
|
||||
)
|
||||
|
||||
if source_entity_entry.device_id is not None:
|
||||
device_registry.async_update_device(
|
||||
source_entity_entry.device_id,
|
||||
add_config_entry_id=helper_config_entry_id,
|
||||
)
|
||||
|
||||
device_registry.async_update_device(
|
||||
source_device_id, remove_config_entry_id=helper_config_entry_id
|
||||
)
|
||||
source_device_id = source_entity_entry.device_id
|
||||
|
||||
# Reload the config entry so the helper entity is recreated with
|
||||
# correct device info
|
||||
await hass.config_entries.async_reload(helper_config_entry_id)
|
||||
|
||||
device_registry = dr.async_get(hass)
|
||||
entity_registry = er.async_get(hass)
|
||||
source_entity_id = er.async_validate_entity_id(
|
||||
entity_registry, source_entity_id_or_uuid
|
||||
)
|
||||
return async_track_entity_registry_updated_event(
|
||||
hass, source_entity_id, async_registry_updated
|
||||
)
|
||||
@ -1,23 +1,103 @@
|
||||
"""Test the Derivative integration."""
|
||||
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
|
||||
from homeassistant.components import derivative
|
||||
from homeassistant.components.derivative.config_flow import ConfigFlowHandler
|
||||
from homeassistant.components.derivative.const import DOMAIN
|
||||
from homeassistant.core import HomeAssistant
|
||||
from homeassistant.config_entries import ConfigEntry
|
||||
from homeassistant.core import Event, HomeAssistant
|
||||
from homeassistant.helpers import device_registry as dr, entity_registry as er
|
||||
from homeassistant.helpers.event import async_track_entity_registry_updated_event
|
||||
|
||||
from tests.common import MockConfigEntry
|
||||
|
||||
|
||||
@pytest.mark.parametrize("platform", ["sensor"])
|
||||
@pytest.fixture
|
||||
def sensor_config_entry(hass: HomeAssistant) -> er.RegistryEntry:
|
||||
"""Fixture to create a sensor config entry."""
|
||||
sensor_config_entry = MockConfigEntry()
|
||||
sensor_config_entry.add_to_hass(hass)
|
||||
return sensor_config_entry
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sensor_device(
|
||||
device_registry: dr.DeviceRegistry, sensor_config_entry: ConfigEntry
|
||||
) -> dr.DeviceEntry:
|
||||
"""Fixture to create a sensor device."""
|
||||
return device_registry.async_get_or_create(
|
||||
config_entry_id=sensor_config_entry.entry_id,
|
||||
connections={(dr.CONNECTION_NETWORK_MAC, "12:34:56:AB:CD:EF")},
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sensor_entity_entry(
|
||||
entity_registry: er.EntityRegistry,
|
||||
sensor_config_entry: ConfigEntry,
|
||||
sensor_device: dr.DeviceEntry,
|
||||
) -> er.RegistryEntry:
|
||||
"""Fixture to create a sensor entity entry."""
|
||||
return entity_registry.async_get_or_create(
|
||||
"sensor",
|
||||
"test",
|
||||
"unique",
|
||||
config_entry=sensor_config_entry,
|
||||
device_id=sensor_device.id,
|
||||
original_name="ABC",
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def derivative_config_entry(
|
||||
hass: HomeAssistant,
|
||||
sensor_entity_entry: er.RegistryEntry,
|
||||
) -> MockConfigEntry:
|
||||
"""Fixture to create a derivative config entry."""
|
||||
config_entry = MockConfigEntry(
|
||||
data={},
|
||||
domain=DOMAIN,
|
||||
options={
|
||||
"name": "My derivative",
|
||||
"round": 1.0,
|
||||
"source": sensor_entity_entry.entity_id,
|
||||
"time_window": {"seconds": 0.0},
|
||||
"unit_prefix": "k",
|
||||
"unit_time": "min",
|
||||
},
|
||||
title="My derivative",
|
||||
version=ConfigFlowHandler.VERSION,
|
||||
minor_version=ConfigFlowHandler.MINOR_VERSION,
|
||||
)
|
||||
|
||||
config_entry.add_to_hass(hass)
|
||||
|
||||
return config_entry
|
||||
|
||||
|
||||
def track_entity_registry_actions(hass: HomeAssistant, entity_id: str) -> list[str]:
|
||||
"""Track entity registry actions for an entity."""
|
||||
events = []
|
||||
|
||||
def add_event(event: Event[er.EventEntityRegistryUpdatedData]) -> None:
|
||||
"""Add entity registry updated event to the list."""
|
||||
events.append(event.data["action"])
|
||||
|
||||
async_track_entity_registry_updated_event(hass, entity_id, add_event)
|
||||
|
||||
return events
|
||||
|
||||
|
||||
async def test_setup_and_remove_config_entry(
|
||||
hass: HomeAssistant,
|
||||
entity_registry: er.EntityRegistry,
|
||||
platform: str,
|
||||
) -> None:
|
||||
"""Test setting up and removing a config entry."""
|
||||
input_sensor_entity_id = "sensor.input"
|
||||
derivative_entity_id = f"{platform}.my_derivative"
|
||||
derivative_entity_id = "sensor.my_derivative"
|
||||
|
||||
# Setup the config entry
|
||||
config_entry = MockConfigEntry(
|
||||
@ -147,3 +227,194 @@ async def test_device_cleaning(
|
||||
derivative_config_entry.entry_id
|
||||
)
|
||||
assert len(devices_after_reload) == 1
|
||||
|
||||
|
||||
async def test_async_handle_source_entity_changes_source_entity_removed(
|
||||
hass: HomeAssistant,
|
||||
device_registry: dr.DeviceRegistry,
|
||||
entity_registry: er.EntityRegistry,
|
||||
derivative_config_entry: MockConfigEntry,
|
||||
sensor_config_entry: ConfigEntry,
|
||||
sensor_device: dr.DeviceEntry,
|
||||
sensor_entity_entry: er.RegistryEntry,
|
||||
) -> None:
|
||||
"""Test the derivative config entry is removed when the source entity is removed."""
|
||||
# Add another config entry to the sensor device
|
||||
other_config_entry = MockConfigEntry()
|
||||
other_config_entry.add_to_hass(hass)
|
||||
device_registry.async_update_device(
|
||||
sensor_device.id, add_config_entry_id=other_config_entry.entry_id
|
||||
)
|
||||
|
||||
assert await hass.config_entries.async_setup(derivative_config_entry.entry_id)
|
||||
await hass.async_block_till_done()
|
||||
|
||||
derivative_entity_entry = entity_registry.async_get("sensor.my_derivative")
|
||||
assert derivative_entity_entry.device_id == sensor_entity_entry.device_id
|
||||
|
||||
sensor_device = device_registry.async_get(sensor_device.id)
|
||||
assert derivative_config_entry.entry_id in sensor_device.config_entries
|
||||
|
||||
events = track_entity_registry_actions(hass, derivative_entity_entry.entity_id)
|
||||
|
||||
# Remove the source sensor's config entry from the device, this removes the
|
||||
# source sensor
|
||||
with patch(
|
||||
"homeassistant.components.derivative.async_unload_entry",
|
||||
wraps=derivative.async_unload_entry,
|
||||
) as mock_unload_entry:
|
||||
device_registry.async_update_device(
|
||||
sensor_device.id, remove_config_entry_id=sensor_config_entry.entry_id
|
||||
)
|
||||
await hass.async_block_till_done()
|
||||
await hass.async_block_till_done()
|
||||
mock_unload_entry.assert_called_once()
|
||||
|
||||
# Check that the derivative config entry is removed from the device
|
||||
sensor_device = device_registry.async_get(sensor_device.id)
|
||||
assert derivative_config_entry.entry_id not in sensor_device.config_entries
|
||||
|
||||
# Check that the derivative config entry is removed
|
||||
assert derivative_config_entry.entry_id not in hass.config_entries.async_entry_ids()
|
||||
|
||||
# Check we got the expected events
|
||||
assert events == ["remove"]
|
||||
|
||||
|
||||
async def test_async_handle_source_entity_changes_source_entity_removed_from_device(
|
||||
hass: HomeAssistant,
|
||||
device_registry: dr.DeviceRegistry,
|
||||
entity_registry: er.EntityRegistry,
|
||||
derivative_config_entry: MockConfigEntry,
|
||||
sensor_device: dr.DeviceEntry,
|
||||
sensor_entity_entry: er.RegistryEntry,
|
||||
) -> None:
|
||||
"""Test the source entity removed from the source device."""
|
||||
assert await hass.config_entries.async_setup(derivative_config_entry.entry_id)
|
||||
await hass.async_block_till_done()
|
||||
|
||||
derivative_entity_entry = entity_registry.async_get("sensor.my_derivative")
|
||||
assert derivative_entity_entry.device_id == sensor_entity_entry.device_id
|
||||
|
||||
sensor_device = device_registry.async_get(sensor_device.id)
|
||||
assert derivative_config_entry.entry_id in sensor_device.config_entries
|
||||
|
||||
events = track_entity_registry_actions(hass, derivative_entity_entry.entity_id)
|
||||
|
||||
# Remove the source sensor from the device
|
||||
with patch(
|
||||
"homeassistant.components.derivative.async_unload_entry",
|
||||
wraps=derivative.async_unload_entry,
|
||||
) as mock_unload_entry:
|
||||
entity_registry.async_update_entity(
|
||||
sensor_entity_entry.entity_id, device_id=None
|
||||
)
|
||||
await hass.async_block_till_done()
|
||||
mock_unload_entry.assert_called_once()
|
||||
|
||||
# Check that the derivative config entry is removed from the device
|
||||
sensor_device = device_registry.async_get(sensor_device.id)
|
||||
assert derivative_config_entry.entry_id not in sensor_device.config_entries
|
||||
|
||||
# Check that the derivative config entry is not removed
|
||||
assert derivative_config_entry.entry_id in hass.config_entries.async_entry_ids()
|
||||
|
||||
# Check we got the expected events
|
||||
assert events == ["update"]
|
||||
|
||||
|
||||
async def test_async_handle_source_entity_changes_source_entity_moved_other_device(
|
||||
hass: HomeAssistant,
|
||||
device_registry: dr.DeviceRegistry,
|
||||
entity_registry: er.EntityRegistry,
|
||||
derivative_config_entry: MockConfigEntry,
|
||||
sensor_config_entry: ConfigEntry,
|
||||
sensor_device: dr.DeviceEntry,
|
||||
sensor_entity_entry: er.RegistryEntry,
|
||||
) -> None:
|
||||
"""Test the source entity is moved to another device."""
|
||||
sensor_device_2 = device_registry.async_get_or_create(
|
||||
config_entry_id=sensor_config_entry.entry_id,
|
||||
connections={(dr.CONNECTION_NETWORK_MAC, "12:34:56:AB:CD:FF")},
|
||||
)
|
||||
|
||||
assert await hass.config_entries.async_setup(derivative_config_entry.entry_id)
|
||||
await hass.async_block_till_done()
|
||||
|
||||
derivative_entity_entry = entity_registry.async_get("sensor.my_derivative")
|
||||
assert derivative_entity_entry.device_id == sensor_entity_entry.device_id
|
||||
|
||||
sensor_device = device_registry.async_get(sensor_device.id)
|
||||
assert derivative_config_entry.entry_id in sensor_device.config_entries
|
||||
sensor_device_2 = device_registry.async_get(sensor_device_2.id)
|
||||
assert derivative_config_entry.entry_id not in sensor_device_2.config_entries
|
||||
|
||||
events = track_entity_registry_actions(hass, derivative_entity_entry.entity_id)
|
||||
|
||||
# Move the source sensor to another device
|
||||
with patch(
|
||||
"homeassistant.components.derivative.async_unload_entry",
|
||||
wraps=derivative.async_unload_entry,
|
||||
) as mock_unload_entry:
|
||||
entity_registry.async_update_entity(
|
||||
sensor_entity_entry.entity_id, device_id=sensor_device_2.id
|
||||
)
|
||||
await hass.async_block_till_done()
|
||||
mock_unload_entry.assert_called_once()
|
||||
|
||||
# Check that the derivative config entry is moved to the other device
|
||||
sensor_device = device_registry.async_get(sensor_device.id)
|
||||
assert derivative_config_entry.entry_id not in sensor_device.config_entries
|
||||
sensor_device_2 = device_registry.async_get(sensor_device_2.id)
|
||||
assert derivative_config_entry.entry_id in sensor_device_2.config_entries
|
||||
|
||||
# Check that the derivative config entry is not removed
|
||||
assert derivative_config_entry.entry_id in hass.config_entries.async_entry_ids()
|
||||
|
||||
# Check we got the expected events
|
||||
assert events == ["update"]
|
||||
|
||||
|
||||
async def test_async_handle_source_entity_new_entity_id(
|
||||
hass: HomeAssistant,
|
||||
device_registry: dr.DeviceRegistry,
|
||||
entity_registry: er.EntityRegistry,
|
||||
derivative_config_entry: MockConfigEntry,
|
||||
sensor_device: dr.DeviceEntry,
|
||||
sensor_entity_entry: er.RegistryEntry,
|
||||
) -> None:
|
||||
"""Test the source entity's entity ID is changed."""
|
||||
assert await hass.config_entries.async_setup(derivative_config_entry.entry_id)
|
||||
await hass.async_block_till_done()
|
||||
|
||||
derivative_entity_entry = entity_registry.async_get("sensor.my_derivative")
|
||||
assert derivative_entity_entry.device_id == sensor_entity_entry.device_id
|
||||
|
||||
sensor_device = device_registry.async_get(sensor_device.id)
|
||||
assert derivative_config_entry.entry_id in sensor_device.config_entries
|
||||
|
||||
events = track_entity_registry_actions(hass, derivative_entity_entry.entity_id)
|
||||
|
||||
# Change the source entity's entity ID
|
||||
with patch(
|
||||
"homeassistant.components.derivative.async_unload_entry",
|
||||
wraps=derivative.async_unload_entry,
|
||||
) as mock_unload_entry:
|
||||
entity_registry.async_update_entity(
|
||||
sensor_entity_entry.entity_id, new_entity_id="sensor.new_entity_id"
|
||||
)
|
||||
await hass.async_block_till_done()
|
||||
mock_unload_entry.assert_called_once()
|
||||
|
||||
# Check that the derivative config entry is updated with the new entity ID
|
||||
assert derivative_config_entry.options["source"] == "sensor.new_entity_id"
|
||||
|
||||
# Check that the helper config is still in the device
|
||||
sensor_device = device_registry.async_get(sensor_device.id)
|
||||
assert derivative_config_entry.entry_id in sensor_device.config_entries
|
||||
|
||||
# Check that the derivative config entry is not removed
|
||||
assert derivative_config_entry.entry_id in hass.config_entries.async_entry_ids()
|
||||
|
||||
# Check we got the expected events
|
||||
assert events == []
|
||||
|
||||
424
tests/helpers/test_helper_integration.py
Normal file
424
tests/helpers/test_helper_integration.py
Normal file
@ -0,0 +1,424 @@
|
||||
"""Tests for the helper entity helpers."""
|
||||
|
||||
from collections.abc import Generator
|
||||
from unittest.mock import AsyncMock, Mock
|
||||
|
||||
import pytest
|
||||
|
||||
from homeassistant.config_entries import ConfigEntry
|
||||
from homeassistant.core import Event, HomeAssistant
|
||||
from homeassistant.helpers import device_registry as dr, entity_registry as er
|
||||
from homeassistant.helpers.event import async_track_entity_registry_updated_event
|
||||
from homeassistant.helpers.helper_integration import async_handle_source_entity_changes
|
||||
|
||||
from tests.common import (
|
||||
MockConfigEntry,
|
||||
MockModule,
|
||||
mock_config_flow,
|
||||
mock_integration,
|
||||
mock_platform,
|
||||
)
|
||||
|
||||
HELPER_DOMAIN = "helper"
|
||||
SOURCE_DOMAIN = "test"
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def source_config_entry(hass: HomeAssistant) -> er.RegistryEntry:
|
||||
"""Fixture to create a source config entry."""
|
||||
source_config_entry = MockConfigEntry()
|
||||
source_config_entry.add_to_hass(hass)
|
||||
return source_config_entry
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def source_device(
|
||||
device_registry: dr.DeviceRegistry,
|
||||
source_config_entry: ConfigEntry,
|
||||
) -> dr.DeviceEntry:
|
||||
"""Fixture to create a source device."""
|
||||
return device_registry.async_get_or_create(
|
||||
config_entry_id=source_config_entry.entry_id,
|
||||
connections={(dr.CONNECTION_NETWORK_MAC, "12:34:56:AB:CD:EF")},
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def source_entity_entry(
|
||||
entity_registry: er.EntityRegistry,
|
||||
source_config_entry: ConfigEntry,
|
||||
source_device: dr.DeviceEntry,
|
||||
) -> er.RegistryEntry:
|
||||
"""Fixture to create a source entity entry."""
|
||||
return entity_registry.async_get_or_create(
|
||||
"sensor",
|
||||
SOURCE_DOMAIN,
|
||||
"unique",
|
||||
config_entry=source_config_entry,
|
||||
device_id=source_device.id,
|
||||
original_name="ABC",
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def helper_config_entry(
|
||||
hass: HomeAssistant,
|
||||
source_entity_entry: er.RegistryEntry,
|
||||
use_entity_registry_id: bool,
|
||||
) -> MockConfigEntry:
|
||||
"""Fixture to create a helper config entry."""
|
||||
config_entry = MockConfigEntry(
|
||||
data={},
|
||||
domain=HELPER_DOMAIN,
|
||||
options={
|
||||
"name": "My helper",
|
||||
"round": 1.0,
|
||||
"source": source_entity_entry.id
|
||||
if use_entity_registry_id
|
||||
else source_entity_entry.entity_id,
|
||||
"time_window": {"seconds": 0.0},
|
||||
"unit_prefix": "k",
|
||||
"unit_time": "min",
|
||||
},
|
||||
title="My helper",
|
||||
)
|
||||
|
||||
config_entry.add_to_hass(hass)
|
||||
|
||||
return config_entry
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_helper_flow() -> Generator[None]:
|
||||
"""Mock helper config flow."""
|
||||
|
||||
class MockConfigFlow:
|
||||
"""Mock the helper config flow."""
|
||||
|
||||
VERSION = 1
|
||||
MINOR_VERSION = 1
|
||||
|
||||
with mock_config_flow(HELPER_DOMAIN, MockConfigFlow):
|
||||
yield
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def helper_entity_entry(
|
||||
entity_registry: er.EntityRegistry,
|
||||
helper_config_entry: ConfigEntry,
|
||||
source_device: dr.DeviceEntry,
|
||||
) -> er.RegistryEntry:
|
||||
"""Fixture to create a helper entity entry."""
|
||||
return entity_registry.async_get_or_create(
|
||||
"sensor",
|
||||
HELPER_DOMAIN,
|
||||
helper_config_entry.entry_id,
|
||||
config_entry=helper_config_entry,
|
||||
device_id=source_device.id,
|
||||
original_name="ABC",
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def async_remove_entry() -> AsyncMock:
|
||||
"""Fixture to mock async_remove_entry."""
|
||||
return AsyncMock(return_value=True)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def async_unload_entry() -> AsyncMock:
|
||||
"""Fixture to mock async_unload_entry."""
|
||||
return AsyncMock(return_value=True)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def set_source_entity_id_or_uuid() -> AsyncMock:
|
||||
"""Fixture to mock async_unload_entry."""
|
||||
return Mock()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_helper_integration(
|
||||
hass: HomeAssistant,
|
||||
entity_registry: er.EntityRegistry,
|
||||
helper_config_entry: MockConfigEntry,
|
||||
source_entity_entry: er.RegistryEntry,
|
||||
async_remove_entry: AsyncMock,
|
||||
async_unload_entry: AsyncMock,
|
||||
set_source_entity_id_or_uuid: Mock,
|
||||
) -> None:
|
||||
"""Mock the helper integration."""
|
||||
|
||||
def get_helper_entity_id() -> str | None:
|
||||
"""Get the helper entity ID."""
|
||||
return entity_registry.async_get_entity_id(
|
||||
"sensor", HELPER_DOMAIN, helper_config_entry.entry_id
|
||||
)
|
||||
|
||||
async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
|
||||
"""Mock setup entry."""
|
||||
async_handle_source_entity_changes(
|
||||
hass,
|
||||
helper_config_entry_id=helper_config_entry.entry_id,
|
||||
get_helper_entity_id=get_helper_entity_id,
|
||||
set_source_entity_id_or_uuid=set_source_entity_id_or_uuid,
|
||||
source_device_id=source_entity_entry.device_id,
|
||||
source_entity_id_or_uuid=helper_config_entry.options["source"],
|
||||
)
|
||||
return True
|
||||
|
||||
mock_integration(
|
||||
hass,
|
||||
MockModule(
|
||||
HELPER_DOMAIN,
|
||||
async_remove_entry=async_remove_entry,
|
||||
async_setup_entry=async_setup_entry,
|
||||
async_unload_entry=async_unload_entry,
|
||||
),
|
||||
)
|
||||
mock_platform(hass, f"{HELPER_DOMAIN}.config_flow", None)
|
||||
|
||||
|
||||
def track_entity_registry_actions(hass: HomeAssistant, entity_id: str) -> list[str]:
|
||||
"""Track entity registry actions for an entity."""
|
||||
events = []
|
||||
|
||||
def add_event(event: Event[er.EventEntityRegistryUpdatedData]) -> None:
|
||||
"""Add entity registry updated event to the list."""
|
||||
events.append(event.data["action"])
|
||||
|
||||
async_track_entity_registry_updated_event(hass, entity_id, add_event)
|
||||
|
||||
return events
|
||||
|
||||
|
||||
@pytest.mark.parametrize("use_entity_registry_id", [True, False])
|
||||
@pytest.mark.usefixtures("mock_helper_flow", "mock_helper_integration")
|
||||
async def test_async_handle_source_entity_changes_source_entity_removed(
|
||||
hass: HomeAssistant,
|
||||
device_registry: dr.DeviceRegistry,
|
||||
entity_registry: er.EntityRegistry,
|
||||
helper_config_entry: MockConfigEntry,
|
||||
helper_entity_entry: er.RegistryEntry,
|
||||
source_config_entry: ConfigEntry,
|
||||
source_device: dr.DeviceEntry,
|
||||
source_entity_entry: er.RegistryEntry,
|
||||
async_remove_entry: AsyncMock,
|
||||
async_unload_entry: AsyncMock,
|
||||
set_source_entity_id_or_uuid: Mock,
|
||||
) -> None:
|
||||
"""Test the helper config entry is removed when the source entity is removed."""
|
||||
# Add the helper config entry to the source device
|
||||
device_registry.async_update_device(
|
||||
source_device.id, add_config_entry_id=helper_config_entry.entry_id
|
||||
)
|
||||
# Add another config entry to the source device
|
||||
other_config_entry = MockConfigEntry()
|
||||
other_config_entry.add_to_hass(hass)
|
||||
device_registry.async_update_device(
|
||||
source_device.id, add_config_entry_id=other_config_entry.entry_id
|
||||
)
|
||||
|
||||
assert await hass.config_entries.async_setup(helper_config_entry.entry_id)
|
||||
await hass.async_block_till_done()
|
||||
|
||||
# Check preconditions
|
||||
helper_entity_entry = entity_registry.async_get(helper_entity_entry.entity_id)
|
||||
assert helper_entity_entry.device_id == source_entity_entry.device_id
|
||||
source_device = device_registry.async_get(source_device.id)
|
||||
assert helper_config_entry.entry_id in source_device.config_entries
|
||||
|
||||
events = track_entity_registry_actions(hass, helper_entity_entry.entity_id)
|
||||
|
||||
# Remove the source entitys's config entry from the device, this removes the
|
||||
# source entity
|
||||
device_registry.async_update_device(
|
||||
source_device.id, remove_config_entry_id=source_config_entry.entry_id
|
||||
)
|
||||
await hass.async_block_till_done()
|
||||
await hass.async_block_till_done()
|
||||
|
||||
# Check that the helper config entry is unloaded and removed
|
||||
async_unload_entry.assert_called_once()
|
||||
async_remove_entry.assert_called_once()
|
||||
set_source_entity_id_or_uuid.assert_not_called()
|
||||
|
||||
# Check that the helper config entry is removed from the device
|
||||
source_device = device_registry.async_get(source_device.id)
|
||||
assert helper_config_entry.entry_id not in source_device.config_entries
|
||||
|
||||
# Check that the helper config entry is removed
|
||||
assert helper_config_entry.entry_id not in hass.config_entries.async_entry_ids()
|
||||
|
||||
# Check we got the expected events
|
||||
assert events == ["remove"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("use_entity_registry_id", [True, False])
|
||||
@pytest.mark.usefixtures("mock_helper_flow", "mock_helper_integration")
|
||||
async def test_async_handle_source_entity_changes_source_entity_removed_from_device(
|
||||
hass: HomeAssistant,
|
||||
device_registry: dr.DeviceRegistry,
|
||||
entity_registry: er.EntityRegistry,
|
||||
helper_config_entry: MockConfigEntry,
|
||||
helper_entity_entry: er.RegistryEntry,
|
||||
source_device: dr.DeviceEntry,
|
||||
source_entity_entry: er.RegistryEntry,
|
||||
async_remove_entry: AsyncMock,
|
||||
async_unload_entry: AsyncMock,
|
||||
set_source_entity_id_or_uuid: Mock,
|
||||
) -> None:
|
||||
"""Test the source entity removed from the source device."""
|
||||
# Add the helper config entry to the source device
|
||||
device_registry.async_update_device(
|
||||
source_device.id, add_config_entry_id=helper_config_entry.entry_id
|
||||
)
|
||||
|
||||
assert await hass.config_entries.async_setup(helper_config_entry.entry_id)
|
||||
await hass.async_block_till_done()
|
||||
|
||||
# Check preconditions
|
||||
helper_entity_entry = entity_registry.async_get(helper_entity_entry.entity_id)
|
||||
assert helper_entity_entry.device_id == source_entity_entry.device_id
|
||||
|
||||
source_device = device_registry.async_get(source_device.id)
|
||||
assert helper_config_entry.entry_id in source_device.config_entries
|
||||
|
||||
events = track_entity_registry_actions(hass, helper_entity_entry.entity_id)
|
||||
|
||||
# Remove the source entity from the device
|
||||
entity_registry.async_update_entity(source_entity_entry.entity_id, device_id=None)
|
||||
await hass.async_block_till_done()
|
||||
async_remove_entry.assert_not_called()
|
||||
async_unload_entry.assert_called_once()
|
||||
set_source_entity_id_or_uuid.assert_not_called()
|
||||
|
||||
# Check that the helper config entry is removed from the device
|
||||
source_device = device_registry.async_get(source_device.id)
|
||||
assert helper_config_entry.entry_id not in source_device.config_entries
|
||||
|
||||
# Check that the helper config entry is not removed
|
||||
assert helper_config_entry.entry_id in hass.config_entries.async_entry_ids()
|
||||
|
||||
# Check we got the expected events
|
||||
assert events == ["update"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("use_entity_registry_id", [True, False])
|
||||
@pytest.mark.usefixtures("mock_helper_flow", "mock_helper_integration")
|
||||
async def test_async_handle_source_entity_changes_source_entity_moved_other_device(
|
||||
hass: HomeAssistant,
|
||||
device_registry: dr.DeviceRegistry,
|
||||
entity_registry: er.EntityRegistry,
|
||||
helper_config_entry: MockConfigEntry,
|
||||
helper_entity_entry: er.RegistryEntry,
|
||||
source_config_entry: ConfigEntry,
|
||||
source_device: dr.DeviceEntry,
|
||||
source_entity_entry: er.RegistryEntry,
|
||||
async_remove_entry: AsyncMock,
|
||||
async_unload_entry: AsyncMock,
|
||||
set_source_entity_id_or_uuid: Mock,
|
||||
) -> None:
|
||||
"""Test the source entity is moved to another device."""
|
||||
# Add the helper config entry to the source device
|
||||
device_registry.async_update_device(
|
||||
source_device.id, add_config_entry_id=helper_config_entry.entry_id
|
||||
)
|
||||
|
||||
# Create another device to move the source entity to
|
||||
source_device_2 = device_registry.async_get_or_create(
|
||||
config_entry_id=source_config_entry.entry_id,
|
||||
connections={(dr.CONNECTION_NETWORK_MAC, "12:34:56:AB:CD:FF")},
|
||||
)
|
||||
|
||||
assert await hass.config_entries.async_setup(helper_config_entry.entry_id)
|
||||
await hass.async_block_till_done()
|
||||
|
||||
# Check preconditions
|
||||
helper_entity_entry = entity_registry.async_get(helper_entity_entry.entity_id)
|
||||
assert helper_entity_entry.device_id == source_entity_entry.device_id
|
||||
|
||||
source_device = device_registry.async_get(source_device.id)
|
||||
assert helper_config_entry.entry_id in source_device.config_entries
|
||||
source_device_2 = device_registry.async_get(source_device_2.id)
|
||||
assert helper_config_entry.entry_id not in source_device_2.config_entries
|
||||
|
||||
events = track_entity_registry_actions(hass, helper_entity_entry.entity_id)
|
||||
|
||||
# Move the source entity to another device
|
||||
entity_registry.async_update_entity(
|
||||
source_entity_entry.entity_id, device_id=source_device_2.id
|
||||
)
|
||||
await hass.async_block_till_done()
|
||||
async_remove_entry.assert_not_called()
|
||||
async_unload_entry.assert_called_once()
|
||||
set_source_entity_id_or_uuid.assert_not_called()
|
||||
|
||||
# Check that the helper config entry is moved to the other device
|
||||
source_device = device_registry.async_get(source_device.id)
|
||||
assert helper_config_entry.entry_id not in source_device.config_entries
|
||||
source_device_2 = device_registry.async_get(source_device_2.id)
|
||||
assert helper_config_entry.entry_id in source_device_2.config_entries
|
||||
|
||||
# Check that the helper config entry is not removed
|
||||
assert helper_config_entry.entry_id in hass.config_entries.async_entry_ids()
|
||||
|
||||
# Check we got the expected events
|
||||
assert events == ["update"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("use_entity_registry_id", "unload_calls", "set_source_entity_id_calls"),
|
||||
[(True, 1, 0), (False, 0, 1)],
|
||||
)
|
||||
@pytest.mark.usefixtures("mock_helper_flow", "mock_helper_integration")
|
||||
async def test_async_handle_source_entity_new_entity_id(
|
||||
hass: HomeAssistant,
|
||||
device_registry: dr.DeviceRegistry,
|
||||
entity_registry: er.EntityRegistry,
|
||||
helper_config_entry: MockConfigEntry,
|
||||
helper_entity_entry: er.RegistryEntry,
|
||||
source_device: dr.DeviceEntry,
|
||||
source_entity_entry: er.RegistryEntry,
|
||||
async_remove_entry: AsyncMock,
|
||||
async_unload_entry: AsyncMock,
|
||||
set_source_entity_id_or_uuid: Mock,
|
||||
unload_calls: int,
|
||||
set_source_entity_id_calls: int,
|
||||
) -> None:
|
||||
"""Test the source entity's entity ID is changed."""
|
||||
# Add the helper config entry to the source device
|
||||
device_registry.async_update_device(
|
||||
source_device.id, add_config_entry_id=helper_config_entry.entry_id
|
||||
)
|
||||
|
||||
assert await hass.config_entries.async_setup(helper_config_entry.entry_id)
|
||||
await hass.async_block_till_done()
|
||||
|
||||
# Check preconditions
|
||||
helper_entity_entry = entity_registry.async_get(helper_entity_entry.entity_id)
|
||||
assert helper_entity_entry.device_id == source_entity_entry.device_id
|
||||
|
||||
source_device = device_registry.async_get(source_device.id)
|
||||
assert helper_config_entry.entry_id in source_device.config_entries
|
||||
|
||||
events = track_entity_registry_actions(hass, helper_entity_entry.entity_id)
|
||||
|
||||
# Change the source entity's entity ID
|
||||
entity_registry.async_update_entity(
|
||||
source_entity_entry.entity_id, new_entity_id="sensor.new_entity_id"
|
||||
)
|
||||
await hass.async_block_till_done()
|
||||
async_remove_entry.assert_not_called()
|
||||
assert len(async_unload_entry.mock_calls) == unload_calls
|
||||
assert len(set_source_entity_id_or_uuid.mock_calls) == set_source_entity_id_calls
|
||||
|
||||
# Check that the helper config is still in the device
|
||||
source_device = device_registry.async_get(source_device.id)
|
||||
assert helper_config_entry.entry_id in source_device.config_entries
|
||||
|
||||
# Check that the helper config entry is not removed
|
||||
assert helper_config_entry.entry_id in hass.config_entries.async_entry_ids()
|
||||
|
||||
# Check we got the expected events
|
||||
assert events == []
|
||||
Loading…
x
Reference in New Issue
Block a user