evlli/hass-core
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Add config subentry support to entity platform

Browse source
This commit is contained in:
Erik 2024-10-11 14:20:32 +02:00
parent 9ff35d5a5a
commit b14b52fb14
3 changed files with 174 additions and 27 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

49
homeassistant/helpers/entity_platform.py
View file

@ -80,6 +80,22 @@ class AddEntitiesCallback(Protocol):
"""Define add_entities type."""
class AddConfigEntryEntitiesCallback(Protocol):
"""Protocol type for EntityPlatform.add_entities callback."""
def __call__(
self,
new_entities: Iterable[Entity],
update_before_add: bool = False,
*,
subentry_id: str | None = None,
) -> None:
"""Define add_entities type.
:param subentry_id: subentry which the entities should be added to
"""
class EntityPlatformModule(Protocol):
"""Protocol type for entity platform modules."""
@ -105,7 +121,7 @@ class EntityPlatformModule(Protocol):
self,
hass: HomeAssistant,
entry: config_entries.ConfigEntry,
async_add_entities: AddEntitiesCallback,
async_add_entities: AddConfigEntryEntitiesCallback,
) -> None:
"""Set up an integration platform from a config entry."""
@ -507,13 +523,17 @@ class EntityPlatform:
@callback
def _async_schedule_add_entities_for_entry(
self, new_entities: Iterable[Entity], update_before_add: bool = False
self,
new_entities: Iterable[Entity],
update_before_add: bool = False,
*,
subentry_id: str | None = None,
) -> None:
"""Schedule adding entities for a single platform async and track the task."""
assert self.config_entry
task = self.config_entry.async_create_task(
self.hass,
self.async_add_entities(new_entities, update_before_add=update_before_add),
self.async_add_entities(new_entities, update_before_add=update_before_add, subentry_id=subentry_id),
f"EntityPlatform async_add_entities_for_entry {self.domain}.{self.platform_name}",
eager_start=True,
)
@ -615,12 +635,26 @@ class EntityPlatform:
)
async def async_add_entities(
self, new_entities: Iterable[Entity], update_before_add: bool = False
self,
new_entities: Iterable[Entity],
update_before_add: bool = False,
*,
subentry_id: str | None = None,
) -> None:
"""Add entities for a single platform async.
This method must be run in the event loop.
:param subentry_id: subentry which the entities should be added to
"""
if subentry_id and (
not self.config_entry or subentry_id not in self.config_entry.subentries
):
raise HomeAssistantError(
f"Can't add entities to unknown subentry {subentry_id} of config "
f"entry {self.config_entry.entry_id if self.config_entry else None}"
)
# handle empty list from component/platform
if not new_entities: # type: ignore[truthy-iterable]
return
@ -631,7 +665,9 @@ class EntityPlatform:
entities: list[Entity] = []
for entity in new_entities:
coros.append(
self._async_add_entity(entity, update_before_add, entity_registry)
self._async_add_entity(
entity, update_before_add, entity_registry, subentry_id
)
)
entities.append(entity)
@ -710,6 +746,7 @@ class EntityPlatform:
entity: Entity,
update_before_add: bool,
entity_registry: EntityRegistry,
subentry_id: str | None,
) -> None:
"""Add an entity to the platform."""
if entity is None:
@ -769,6 +806,7 @@ class EntityPlatform:
try:
device = dev_reg.async_get(self.hass).async_get_or_create(
config_entry_id=self.config_entry.entry_id,
config_subentry_id=subentry_id,
**device_info,
)
except dev_reg.DeviceInfoError as exc:
@ -815,6 +853,7 @@ class EntityPlatform:
entity.unique_id,
capabilities=entity.capability_attributes,
config_entry=self.config_entry,
config_subentry_id=subentry_id,
device_id=device.id if device else None,
disabled_by=disabled_by,
entity_category=entity.entity_category,

46
tests/helpers/snapshots/test_entity_platform.ambr
View file

@ -3,6 +3,11 @@
DeviceRegistryEntrySnapshot({
'area_id': 'heliport',
'config_entries': <ANY>,
'config_subentries': dict({
'super-mock-id': set({
None,
}),
}),
'configuration_url': 'http://192.168.0.100/config',
'connections': set({
tuple(
@ -35,3 +40,44 @@
'via_device_id': <ANY>,
})
# ---
# name: test_device_info_called.1
DeviceRegistryEntrySnapshot({
'area_id': 'heliport',
'config_entries': <ANY>,
'config_subentries': dict({
'super-mock-id': set({
'mock-subentry-id-1',
}),
}),
'configuration_url': 'http://192.168.0.100/config',
'connections': set({
tuple(
'mac',
'efgh',
),
}),
'disabled_by': None,
'entry_type': <DeviceEntryType.SERVICE: 'service'>,
'hw_version': 'test-hw',
'id': <ANY>,
'identifiers': set({
tuple(
'hue',
'efgh',
),
}),
'is_new': False,
'labels': set({
}),
'manufacturer': 'test-manuf',
'model': 'test-model',
'model_id': None,
'name': 'test-name',
'name_by_user': None,
'primary_config_entry': <ANY>,
'serial_number': None,
'suggested_area': 'Heliport',
'sw_version': 'test-sw',
'via_device_id': <ANY>,
})
# ---

106
tests/helpers/test_entity_platform.py
View file

@ -11,7 +11,7 @@ import pytest
from syrupy.assertion import SnapshotAssertion
import voluptuous as vol
from homeassistant.config_entries import ConfigEntry
from homeassistant.config_entries import ConfigEntry, ConfigSubentryData
from homeassistant.const import EVENT_HOMEASSISTANT_STARTED, PERCENTAGE, EntityCategory
from homeassistant.core import (
CoreState,
@ -36,7 +36,10 @@ from homeassistant.helpers.entity_component import (
DEFAULT_SCAN_INTERVAL,
EntityComponent,
)
from homeassistant.helpers.entity_platform import AddEntitiesCallback
from homeassistant.helpers.entity_platform import (
AddConfigEntryEntitiesCallback,
AddEntitiesCallback,
)
from homeassistant.helpers.typing import ConfigType, DiscoveryInfoType
import homeassistant.util.dt as dt_util
@ -223,7 +226,7 @@ async def test_set_scan_interval_via_platform(hass: HomeAssistant) -> None:
def platform_setup(
hass: HomeAssistant,
config: ConfigType,
add_entities: entity_platform.AddEntitiesCallback,
add_entities: AddEntitiesCallback,
discovery_info: DiscoveryInfoType | None = None,
) -> None:
"""Test the platform setup."""
@ -862,13 +865,28 @@ async def test_setup_entry(
async def async_setup_entry(
hass: HomeAssistant,
config_entry: ConfigEntry,
async_add_entities: AddEntitiesCallback,
async_add_entities: AddConfigEntryEntitiesCallback,
) -> None:
"""Mock setup entry method."""
async_add_entities([MockEntity(name="test1", unique_id="unique")])
async_add_entities([MockEntity(name="test1", unique_id="unique1")])
async_add_entities(
[MockEntity(name="test2", unique_id="unique2")],
subentry_id="mock-subentry-id-1",
)
platform = MockPlatform(async_setup_entry=async_setup_entry)
config_entry = MockConfigEntry(entry_id="super-mock-id")
config_entry = MockConfigEntry(
entry_id="super-mock-id",
subentries_data=(
ConfigSubentryData(
data={},
subentry_id="mock-subentry-id-1",
title="Mock title",
unique_id="test",
),
),
)
config_entry.add_to_hass(hass)
entity_platform = MockEntityPlatform(
hass, platform_name=config_entry.domain, platform=platform
)
@ -877,11 +895,16 @@ async def test_setup_entry(
await hass.async_block_till_done()
full_name = f"{config_entry.domain}.{entity_platform.domain}"
assert full_name in hass.config.components
assert len(hass.states.async_entity_ids()) == 1
assert len(entity_registry.entities) == 1
assert len(hass.states.async_entity_ids()) == 2
assert len(entity_registry.entities) == 2
entity_registry_entry = entity_registry.entities["test_domain.test1"]
assert entity_registry_entry.config_entry_id == "super-mock-id"
assert entity_registry_entry.config_subentry_id is None
entity_registry_entry = entity_registry.entities["test_domain.test2"]
assert entity_registry_entry.config_entry_id == "super-mock-id"
assert entity_registry_entry.config_subentry_id == "mock-subentry-id-1"
async def test_setup_entry_platform_not_ready(
@ -1137,7 +1160,17 @@ async def test_device_info_called(
snapshot: SnapshotAssertion,
) -> None:
"""Test device info is forwarded correctly."""
config_entry = MockConfigEntry(entry_id="super-mock-id")
config_entry = MockConfigEntry(
entry_id="super-mock-id",
subentries_data=(
ConfigSubentryData(
data={},
subentry_id="mock-subentry-id-1",
title="Mock title",
unique_id="test",
),
),
)
config_entry.add_to_hass(hass)
via = device_registry.async_get_or_create(
config_entry_id=config_entry.entry_id,
@ -1150,7 +1183,7 @@ async def test_device_info_called(
async def async_setup_entry(
hass: HomeAssistant,
config_entry: ConfigEntry,
async_add_entities: AddEntitiesCallback,
async_add_entities: AddConfigEntryEntitiesCallback,
) -> None:
"""Mock setup entry method."""
async_add_entities(
@ -1176,6 +1209,28 @@ async def test_device_info_called(
),
]
)
async_add_entities(
[
# Valid device info
MockEntity(
unique_id="efgh",
device_info={
"identifiers": {("hue", "efgh")},
"configuration_url": "http://192.168.0.100/config",
"connections": {(dr.CONNECTION_NETWORK_MAC, "efgh")},
"manufacturer": "test-manuf",
"model": "test-model",
"name": "test-name",
"sw_version": "test-sw",
"hw_version": "test-hw",
"suggested_area": "Heliport",
"entry_type": dr.DeviceEntryType.SERVICE,
"via_device": ("hue", "via-id"),
},
),
],
subentry_id="mock-subentry-id-1",
)
platform = MockPlatform(async_setup_entry=async_setup_entry)
entity_platform = MockEntityPlatform(
@ -1185,11 +1240,18 @@ async def test_device_info_called(
assert await entity_platform.async_setup_entry(config_entry)
await hass.async_block_till_done()
assert len(hass.states.async_entity_ids()) == 2
assert len(hass.states.async_entity_ids()) == 3
device = device_registry.async_get_device(identifiers={("hue", "1234")})
assert device == snapshot
assert device.config_entries == {config_entry.entry_id}
assert device.config_subentries == {config_entry.entry_id: {None}}
assert device.primary_config_entry == config_entry.entry_id
assert device.via_device_id == via.id
device = device_registry.async_get_device(identifiers={("hue", "efgh")})
assert device == snapshot
assert device.config_entries == {config_entry.entry_id}
assert device.config_subentries == {config_entry.entry_id: {"mock-subentry-id-1"}}
assert device.primary_config_entry == config_entry.entry_id
assert device.via_device_id == via.id
@ -1213,7 +1275,7 @@ async def test_device_info_not_overrides(
async def async_setup_entry(
hass: HomeAssistant,
config_entry: ConfigEntry,
async_add_entities: AddEntitiesCallback,
async_add_entities: AddConfigEntryEntitiesCallback,
) -> None:
"""Mock setup entry method."""
async_add_entities(
@ -1266,7 +1328,7 @@ async def test_device_info_homeassistant_url(
async def async_setup_entry(
hass: HomeAssistant,
config_entry: ConfigEntry,
async_add_entities: AddEntitiesCallback,
async_add_entities: AddConfigEntryEntitiesCallback,
) -> None:
"""Mock setup entry method."""
async_add_entities(
@ -1318,7 +1380,7 @@ async def test_device_info_change_to_no_url(
async def async_setup_entry(
hass: HomeAssistant,
config_entry: ConfigEntry,
async_add_entities: AddEntitiesCallback,
async_add_entities: AddConfigEntryEntitiesCallback,
) -> None:
"""Mock setup entry method."""
async_add_entities(
@ -1390,7 +1452,7 @@ async def test_entity_disabled_by_device(
async def async_setup_entry(
hass: HomeAssistant,
config_entry: ConfigEntry,
async_add_entities: AddEntitiesCallback,
async_add_entities: AddConfigEntryEntitiesCallback,
) -> None:
"""Mock setup entry method."""
async_add_entities([entity_disabled])
@ -1876,7 +1938,7 @@ async def test_setup_entry_with_entities_that_block_forever(
async def async_setup_entry(
hass: HomeAssistant,
config_entry: ConfigEntry,
async_add_entities: AddEntitiesCallback,
async_add_entities: AddConfigEntryEntitiesCallback,
) -> None:
"""Mock setup entry method."""
async_add_entities(
@ -1924,7 +1986,7 @@ async def test_cancellation_is_not_blocked(
async def async_setup_entry(
hass: HomeAssistant,
config_entry: ConfigEntry,
async_add_entities: AddEntitiesCallback,
async_add_entities: AddConfigEntryEntitiesCallback,
) -> None:
"""Mock setup entry method."""
async_add_entities(
@ -2021,7 +2083,7 @@ async def test_entity_name_influences_entity_id(
async def async_setup_entry(
hass: HomeAssistant,
config_entry: ConfigEntry,
async_add_entities: AddEntitiesCallback,
async_add_entities: AddConfigEntryEntitiesCallback,
) -> None:
"""Mock setup entry method."""
async_add_entities(
@ -2109,7 +2171,7 @@ async def test_translated_entity_name_influences_entity_id(
async def async_setup_entry(
hass: HomeAssistant,
config_entry: ConfigEntry,
async_add_entities: AddEntitiesCallback,
async_add_entities: AddConfigEntryEntitiesCallback,
) -> None:
"""Mock setup entry method."""
async_add_entities(
@ -2197,7 +2259,7 @@ async def test_translated_device_class_name_influences_entity_id(
async def async_setup_entry(
hass: HomeAssistant,
config_entry: ConfigEntry,
async_add_entities: AddEntitiesCallback,
async_add_entities: AddConfigEntryEntitiesCallback,
) -> None:
"""Mock setup entry method."""
async_add_entities([TranslatedDeviceClassEntity(device_class, has_entity_name)])
@ -2259,7 +2321,7 @@ async def test_device_name_defaulting_config_entry(
async def async_setup_entry(
hass: HomeAssistant,
config_entry: ConfigEntry,
async_add_entities: AddEntitiesCallback,
async_add_entities: AddConfigEntryEntitiesCallback,
) -> None:
"""Mock setup entry method."""
async_add_entities([DeviceNameEntity()])
@ -2315,7 +2377,7 @@ async def test_device_type_error_checking(
async def async_setup_entry(
hass: HomeAssistant,
config_entry: ConfigEntry,
async_add_entities: AddEntitiesCallback,
async_add_entities: AddConfigEntryEntitiesCallback,
) -> None:
"""Mock setup entry method."""
async_add_entities([DeviceNameEntity()])