Add the ability to bind the template helper entity to a device (#117753)

This commit is contained in:
dougiteixeira 2024-06-22 07:45:06 -03:00 committed by GitHub
parent f258034f9c
commit 6e32a96ff3
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
11 changed files with 472 additions and 9 deletions

View File

@ -7,10 +7,13 @@ import logging
from homeassistant import config as conf_util
from homeassistant.config_entries import ConfigEntry
from homeassistant.const import CONF_UNIQUE_ID, SERVICE_RELOAD
from homeassistant.const import CONF_DEVICE_ID, CONF_UNIQUE_ID, SERVICE_RELOAD
from homeassistant.core import Event, HomeAssistant, ServiceCall
from homeassistant.exceptions import HomeAssistantError
from homeassistant.helpers import discovery
from homeassistant.helpers.device import (
async_remove_stale_devices_links_keep_current_device,
)
from homeassistant.helpers.reload import async_reload_integration_platforms
from homeassistant.helpers.service import async_register_admin_service
from homeassistant.helpers.typing import ConfigType
@ -57,6 +60,13 @@ async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
"""Set up a config entry."""
async_remove_stale_devices_links_keep_current_device(
hass,
entry.entry_id,
entry.options.get(CONF_DEVICE_ID),
)
await hass.config_entries.async_forward_entry_setups(
entry, (entry.options["template_type"],)
)

View File

@ -22,6 +22,7 @@ from homeassistant.const import (
ATTR_ENTITY_ID,
ATTR_FRIENDLY_NAME,
CONF_DEVICE_CLASS,
CONF_DEVICE_ID,
CONF_ENTITY_PICTURE_TEMPLATE,
CONF_FRIENDLY_NAME,
CONF_FRIENDLY_NAME_TEMPLATE,
@ -39,8 +40,9 @@ from homeassistant.const import (
)
from homeassistant.core import CALLBACK_TYPE, HomeAssistant, callback
from homeassistant.exceptions import TemplateError
from homeassistant.helpers import template
from homeassistant.helpers import selector, template
import homeassistant.helpers.config_validation as cv
from homeassistant.helpers.device import async_device_info_to_link_from_device_id
from homeassistant.helpers.entity import async_generate_entity_id
from homeassistant.helpers.entity_platform import AddEntitiesCallback
from homeassistant.helpers.event import async_call_later, async_track_point_in_utc_time
@ -86,6 +88,7 @@ BINARY_SENSOR_SCHEMA = vol.Schema(
vol.Optional(CONF_DEVICE_CLASS): DEVICE_CLASSES_SCHEMA,
vol.Required(CONF_STATE): cv.template,
vol.Optional(CONF_UNIT_OF_MEASUREMENT): cv.string,
vol.Optional(CONF_DEVICE_ID): selector.DeviceSelector(),
}
).extend(TEMPLATE_ENTITY_COMMON_SCHEMA.schema)
@ -244,6 +247,10 @@ class BinarySensorTemplate(TemplateEntity, BinarySensorEntity, RestoreEntity):
self._delay_on_raw = config.get(CONF_DELAY_ON)
self._delay_off = None
self._delay_off_raw = config.get(CONF_DELAY_OFF)
self._attr_device_info = async_device_info_to_link_from_device_id(
hass,
config.get(CONF_DEVICE_ID),
)
async def async_added_to_hass(self) -> None:
"""Restore state."""

View File

@ -18,6 +18,7 @@ from homeassistant.components.sensor import (
)
from homeassistant.const import (
CONF_DEVICE_CLASS,
CONF_DEVICE_ID,
CONF_NAME,
CONF_STATE,
CONF_UNIT_OF_MEASUREMENT,
@ -95,6 +96,8 @@ def generate_schema(domain: str, flow_type: str) -> dict[vol.Marker, Any]:
),
}
schema[vol.Optional(CONF_DEVICE_ID)] = selector.DeviceSelector()
return schema

View File

@ -25,6 +25,7 @@ from homeassistant.config_entries import ConfigEntry
from homeassistant.const import (
ATTR_ENTITY_ID,
CONF_DEVICE_CLASS,
CONF_DEVICE_ID,
CONF_ENTITY_PICTURE_TEMPLATE,
CONF_FRIENDLY_NAME,
CONF_FRIENDLY_NAME_TEMPLATE,
@ -40,7 +41,8 @@ from homeassistant.const import (
)
from homeassistant.core import HomeAssistant, callback
from homeassistant.exceptions import TemplateError
from homeassistant.helpers import config_validation as cv, template
from homeassistant.helpers import config_validation as cv, selector, template
from homeassistant.helpers.device import async_device_info_to_link_from_device_id
from homeassistant.helpers.entity import async_generate_entity_id
from homeassistant.helpers.entity_platform import AddEntitiesCallback
from homeassistant.helpers.trigger_template_entity import TEMPLATE_SENSOR_BASE_SCHEMA
@ -86,6 +88,7 @@ SENSOR_SCHEMA = vol.All(
{
vol.Required(CONF_STATE): cv.template,
vol.Optional(ATTR_LAST_RESET): cv.template,
vol.Optional(CONF_DEVICE_ID): selector.DeviceSelector(),
}
)
.extend(TEMPLATE_SENSOR_BASE_SCHEMA.schema)
@ -260,6 +263,10 @@ class SensorTemplate(TemplateEntity, SensorEntity):
self._attr_last_reset_template: template.Template | None = config.get(
ATTR_LAST_RESET
)
self._attr_device_info = async_device_info_to_link_from_device_id(
hass,
config.get(CONF_DEVICE_ID),
)
if (object_id := config.get(CONF_OBJECT_ID)) is not None:
self.entity_id = async_generate_entity_id(
ENTITY_ID_FORMAT, object_id, hass=hass

View File

@ -3,20 +3,28 @@
"step": {
"binary_sensor": {
"data": {
"device_id": "[%key:common::config_flow::data::device%]",
"device_class": "[%key:component::template::config::step::sensor::data::device_class%]",
"name": "[%key:common::config_flow::data::name%]",
"state": "[%key:component::template::config::step::sensor::data::state%]"
},
"data_description": {
"device_id": "[%key:component::template::config::step::sensor::data_description::device_id%]"
},
"title": "Template binary sensor"
},
"sensor": {
"data": {
"device_id": "[%key:common::config_flow::data::device%]",
"device_class": "Device class",
"name": "[%key:common::config_flow::data::name%]",
"state_class": "[%key:component::sensor::entity_component::_::state_attributes::state_class::name%]",
"state": "State template",
"unit_of_measurement": "Unit of measurement"
},
"data_description": {
"device_id": "Select a device to link to this entity."
},
"title": "Template sensor"
},
"user": {
@ -33,17 +41,25 @@
"step": {
"binary_sensor": {
"data": {
"device_id": "[%key:common::config_flow::data::device%]",
"state": "[%key:component::template::config::step::sensor::data::state%]"
},
"data_description": {
"device_id": "[%key:component::template::config::step::sensor::data_description::device_id%]"
},
"title": "[%key:component::template::config::step::binary_sensor::title%]"
},
"sensor": {
"data": {
"device_id": "[%key:common::config_flow::data::device%]",
"device_class": "[%key:component::template::config::step::sensor::data::device_class%]",
"state_class": "[%key:component::sensor::entity_component::_::state_attributes::state_class::name%]",
"state": "[%key:component::template::config::step::sensor::data::state%]",
"unit_of_measurement": "[%key:component::template::config::step::sensor::data::unit_of_measurement%]"
},
"data_description": {
"device_id": "[%key:component::template::config::step::sensor::data_description::device_id%]"
},
"title": "[%key:component::template::config::step::sensor::title%]"
}
}

View File

@ -28,11 +28,22 @@ def async_device_info_to_link_from_entity(
) -> 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."""
return async_device_info_to_link_from_device_id(
hass,
async_entity_id_to_device_id(hass, entity_id_or_uuid),
)
@callback
def async_device_info_to_link_from_device_id(
hass: HomeAssistant,
device_id: str | None,
) -> dr.DeviceInfo | None:
"""DeviceInfo with information to link a device to a configuration entry in the link category from a device id."""
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:
if device_id is None or (device := dev_reg.async_get(device_id=device_id)) is None:
return None
return dr.DeviceInfo(

View File

@ -19,7 +19,7 @@ from homeassistant.const import (
STATE_UNKNOWN,
)
from homeassistant.core import Context, CoreState, HomeAssistant, State
from homeassistant.helpers import entity_registry as er
from homeassistant.helpers import device_registry as dr, entity_registry as er
from homeassistant.helpers.entity_component import async_update_entity
from homeassistant.setup import async_setup_component
import homeassistant.util.dt as dt_util
@ -1403,3 +1403,40 @@ async def test_trigger_entity_restore_state_auto_off_expired(
state = hass.states.get("binary_sensor.test")
assert state.state == OFF
async def test_device_id(hass: HomeAssistant) -> None:
"""Test for device for Template."""
device_registry = dr.async_get(hass)
entity_registry = er.async_get(hass)
device_config_entry = MockConfigEntry()
device_config_entry.add_to_hass(hass)
device_entry = device_registry.async_get_or_create(
config_entry_id=device_config_entry.entry_id,
identifiers={("sensor", "identifier_test")},
connections={("mac", "30:31:32:33:34:35")},
)
await hass.async_block_till_done()
assert device_entry is not None
assert device_entry.id is not None
template_config_entry = MockConfigEntry(
data={},
domain=template.DOMAIN,
options={
"name": "My template",
"state": "{{10 > 8}}",
"template_type": "binary_sensor",
"device_id": device_entry.id,
},
title="My template",
)
template_config_entry.add_to_hass(hass)
assert await hass.config_entries.async_setup(template_config_entry.entry_id)
await hass.async_block_till_done()
template_entity = entity_registry.async_get("binary_sensor.my_template")
assert template_entity is not None
assert template_entity.device_id == device_entry.id

View File

@ -11,6 +11,7 @@ from homeassistant.components.template import DOMAIN, async_setup_entry
from homeassistant.const import STATE_UNAVAILABLE
from homeassistant.core import HomeAssistant
from homeassistant.data_entry_flow import FlowResultType
from homeassistant.helpers import device_registry as dr
from tests.common import MockConfigEntry
from tests.typing import WebSocketGenerator
@ -122,6 +123,91 @@ async def test_config_flow(
assert state.attributes[key] == extra_attrs[key]
@pytest.mark.parametrize(
(
"template_type",
"state_template",
),
[
(
"sensor",
"{{ 15 }}",
),
(
"binary_sensor",
"{{ false }}",
),
],
)
async def test_config_flow_device(
hass: HomeAssistant,
template_type: str,
state_template: str,
) -> None:
"""Test remove the device registry configuration entry when the device changes."""
device_registry = dr.async_get(hass)
# Configure a device registry
entry_device = MockConfigEntry()
entry_device.add_to_hass(hass)
device = device_registry.async_get_or_create(
config_entry_id=entry_device.entry_id,
identifiers={("test", "identifier_test1")},
connections={("mac", "20:31:32:33:34:01")},
)
await hass.async_block_till_done()
device_id = device.id
assert device_id is not None
result = await hass.config_entries.flow.async_init(
DOMAIN, context={"source": config_entries.SOURCE_USER}
)
assert result["type"] is FlowResultType.MENU
result = await hass.config_entries.flow.async_configure(
result["flow_id"],
{"next_step_id": template_type},
)
await hass.async_block_till_done()
assert result["type"] is FlowResultType.FORM
assert result["step_id"] == template_type
with patch(
"homeassistant.components.template.async_setup_entry", wraps=async_setup_entry
) as mock_setup_entry:
result = await hass.config_entries.flow.async_configure(
result["flow_id"],
{
"name": "My template",
"state": state_template,
"device_id": device_id,
},
)
await hass.async_block_till_done()
assert result["type"] is FlowResultType.CREATE_ENTRY
assert result["title"] == "My template"
assert result["data"] == {}
assert result["options"] == {
"name": "My template",
"state": state_template,
"template_type": template_type,
"device_id": device_id,
}
assert len(mock_setup_entry.mock_calls) == 1
config_entry = hass.config_entries.async_entries(DOMAIN)[0]
assert config_entry.data == {}
assert config_entry.options == {
"name": "My template",
"state": state_template,
"template_type": template_type,
"device_id": device_id,
}
def get_suggested(schema, key):
"""Get suggested value for key in voluptuous schema."""
for k in schema:
@ -852,3 +938,149 @@ async def test_option_flow_sensor_preview_config_entry_removed(
msg = await client.receive_json()
assert not msg["success"]
assert msg["error"] == {"code": "home_assistant_error", "message": "Unknown error"}
@pytest.mark.parametrize(
(
"template_type",
"state_template",
),
[
(
"sensor",
"{{ 15 }}",
),
(
"binary_sensor",
"{{ false }}",
),
],
)
async def test_options_flow_change_device(
hass: HomeAssistant,
template_type: str,
state_template: str,
) -> None:
"""Test remove the device registry configuration entry when the device changes."""
device_registry = dr.async_get(hass)
# Configure a device registry
entry_device1 = MockConfigEntry()
entry_device1.add_to_hass(hass)
device1 = device_registry.async_get_or_create(
config_entry_id=entry_device1.entry_id,
identifiers={("test", "identifier_test1")},
connections={("mac", "20:31:32:33:34:01")},
)
entry_device2 = MockConfigEntry()
entry_device2.add_to_hass(hass)
device2 = device_registry.async_get_or_create(
config_entry_id=entry_device1.entry_id,
identifiers={("test", "identifier_test2")},
connections={("mac", "20:31:32:33:34:02")},
)
await hass.async_block_till_done()
device_id1 = device1.id
assert device_id1 is not None
device_id2 = device2.id
assert device_id2 is not None
# Setup the config entry with device 1
template_config_entry = MockConfigEntry(
data={},
domain=DOMAIN,
options={
"template_type": template_type,
"name": "Test",
"state": state_template,
"device_id": device_id1,
},
title="Sensor template",
)
template_config_entry.add_to_hass(hass)
assert await hass.config_entries.async_setup(template_config_entry.entry_id)
await hass.async_block_till_done()
# Change to link to device 2
result = await hass.config_entries.options.async_init(
template_config_entry.entry_id
)
assert result["type"] is FlowResultType.FORM
result = await hass.config_entries.options.async_configure(
result["flow_id"],
user_input={
"state": state_template,
"device_id": device_id2,
},
)
assert result["type"] is FlowResultType.CREATE_ENTRY
assert result["data"] == {
"template_type": template_type,
"name": "Test",
"state": state_template,
"device_id": device_id2,
}
assert template_config_entry.data == {}
assert template_config_entry.options == {
"template_type": template_type,
"name": "Test",
"state": state_template,
"device_id": device_id2,
}
# Remove link with device
result = await hass.config_entries.options.async_init(
template_config_entry.entry_id
)
assert result["type"] is FlowResultType.FORM
result = await hass.config_entries.options.async_configure(
result["flow_id"],
user_input={
"state": state_template,
},
)
assert result["type"] is FlowResultType.CREATE_ENTRY
assert result["data"] == {
"template_type": template_type,
"name": "Test",
"state": state_template,
}
assert template_config_entry.data == {}
assert template_config_entry.options == {
"template_type": template_type,
"name": "Test",
"state": state_template,
}
# Change to link to device 1
result = await hass.config_entries.options.async_init(
template_config_entry.entry_id
)
assert result["type"] is FlowResultType.FORM
result = await hass.config_entries.options.async_configure(
result["flow_id"],
user_input={
"state": state_template,
"device_id": device_id1,
},
)
assert result["type"] is FlowResultType.CREATE_ENTRY
assert result["data"] == {
"template_type": template_type,
"name": "Test",
"state": state_template,
"device_id": device_id1,
}
assert template_config_entry.data == {}
assert template_config_entry.options == {
"template_type": template_type,
"name": "Test",
"state": state_template,
"device_id": device_id1,
}

View File

@ -8,11 +8,12 @@ import pytest
from homeassistant import config
from homeassistant.components.template import DOMAIN
from homeassistant.core import HomeAssistant
from homeassistant.helpers import device_registry as dr
from homeassistant.helpers.reload import SERVICE_RELOAD
from homeassistant.setup import async_setup_component
from homeassistant.util import dt as dt_util
from tests.common import async_fire_time_changed, get_fixture_path
from tests.common import MockConfigEntry, async_fire_time_changed, get_fixture_path
@pytest.mark.parametrize(("count", "domain"), [(1, "sensor")])
@ -268,3 +269,90 @@ async def async_yaml_patch_helper(hass, filename):
blocking=True,
)
await hass.async_block_till_done()
async def test_change_device(hass: HomeAssistant) -> None:
"""Test remove the device registry configuration entry when the device changes."""
device_registry = dr.async_get(hass)
# Configure a device registry
entry_device1 = MockConfigEntry()
entry_device1.add_to_hass(hass)
device1 = device_registry.async_get_or_create(
config_entry_id=entry_device1.entry_id,
identifiers={("test", "identifier_test1")},
connections={("mac", "20:31:32:33:34:01")},
)
entry_device2 = MockConfigEntry()
entry_device2.add_to_hass(hass)
device2 = device_registry.async_get_or_create(
config_entry_id=entry_device1.entry_id,
identifiers={("test", "identifier_test2")},
connections={("mac", "20:31:32:33:34:02")},
)
await hass.async_block_till_done()
device_id1 = device1.id
assert device_id1 is not None
device_id2 = device2.id
assert device_id2 is not None
# Setup the config entry (binary_sensor)
sensor_config_entry = MockConfigEntry(
data={},
domain=DOMAIN,
options={
"template_type": "binary_sensor",
"name": "Teste",
"state": "{{15}}",
"device_id": device_id1,
},
title="Binary sensor template",
)
sensor_config_entry.add_to_hass(hass)
assert await hass.config_entries.async_setup(sensor_config_entry.entry_id)
await hass.async_block_till_done()
# Confirm that the configuration entry has been added to the device 1 registry (current)
current_device = device_registry.async_get(device_id=device_id1)
assert sensor_config_entry.entry_id in current_device.config_entries
# Change configuration options to use device 2 and reload the integration
result = await hass.config_entries.options.async_init(sensor_config_entry.entry_id)
result = await hass.config_entries.options.async_configure(
result["flow_id"],
user_input={
"state": "{{15}}",
"device_id": device_id2,
},
)
await hass.async_block_till_done()
# Confirm that the configuration entry has been removed from the device 1 registry (previous)
previous_device = device_registry.async_get(device_id=device_id1)
assert sensor_config_entry.entry_id not in previous_device.config_entries
# Confirm that the configuration entry has been added to the device 2 registry (current)
current_device = device_registry.async_get(device_id=device_id2)
assert sensor_config_entry.entry_id in current_device.config_entries
result = await hass.config_entries.options.async_init(sensor_config_entry.entry_id)
result = await hass.config_entries.options.async_configure(
result["flow_id"],
user_input={
"state": "{{15}}",
},
)
await hass.async_block_till_done()
# Confirm that the configuration entry has been removed from the device 2 registry (previous)
previous_device = device_registry.async_get(device_id=device_id2)
assert sensor_config_entry.entry_id not in previous_device.config_entries
# Confirm that there is no device with the helper configuration entry
assert (
dr.async_entries_for_config_entry(device_registry, sensor_config_entry.entry_id)
== []
)

View File

@ -21,7 +21,7 @@ from homeassistant.const import (
STATE_UNKNOWN,
)
from homeassistant.core import Context, CoreState, HomeAssistant, State, callback
from homeassistant.helpers import entity_registry as er
from homeassistant.helpers import device_registry as dr, entity_registry as er
from homeassistant.helpers.entity_component import async_update_entity
from homeassistant.helpers.template import Template
from homeassistant.setup import ATTR_COMPONENT, async_setup_component
@ -1896,3 +1896,40 @@ async def test_trigger_action(
assert len(events) == 1
assert events[0].context.parent_id == context.id
async def test_device_id(hass: HomeAssistant) -> None:
"""Test for device for Template."""
device_registry = dr.async_get(hass)
device_config_entry = MockConfigEntry()
device_config_entry.add_to_hass(hass)
device_entry = device_registry.async_get_or_create(
config_entry_id=device_config_entry.entry_id,
identifiers={("sensor", "identifier_test")},
connections={("mac", "30:31:32:33:34:35")},
)
await hass.async_block_till_done()
assert device_entry is not None
assert device_entry.id is not None
template_config_entry = MockConfigEntry(
data={},
domain=template.DOMAIN,
options={
"name": "My template",
"state": "{{10}}",
"template_type": "sensor",
"device_id": device_entry.id,
},
title="My template",
)
template_config_entry.add_to_hass(hass)
assert await hass.config_entries.async_setup(template_config_entry.entry_id)
await hass.async_block_till_done()
entity_registry = er.async_get(hass)
template_entity = entity_registry.async_get("sensor.my_template")
assert template_entity is not None
assert template_entity.device_id == device_entry.id

View File

@ -6,6 +6,7 @@ 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_device_id,
async_device_info_to_link_from_entity,
async_entity_id_to_device_id,
async_remove_stale_devices_links_keep_current_device,
@ -90,12 +91,26 @@ async def test_device_info_to_link(
"connections": {("mac", "30:31:32:33:34:00")},
}
result = async_device_info_to_link_from_device_id(hass, device_id=device.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
# With a non-existent device id
result = async_device_info_to_link_from_device_id(hass, device_id="abcdefghi")
assert result is None
# With a None device id
result = async_device_info_to_link_from_device_id(hass, device_id=None)
assert result is None
async def test_remove_stale_device_links_keep_entity_device(
hass: HomeAssistant,