Add config subentry support to entity platform

homeassistant/helpers/entity_platform.py

@@ -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."""
 
@@ -516,13 +532,21 @@ 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,
         )
@@ -624,12 +648,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
@@ -640,7 +678,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)
 
@@ -719,6 +759,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:
@@ -778,6 +819,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:
@@ -824,6 +866,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,

tests/helpers/snapshots/test_entity_platform.ambr

@@ -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>,
+  })
+# ---

tests/helpers/test_entity_platform.py

@@ -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(hass.states.async_entity_ids()) == 2
-    assert len(entity_registry.entities) == 1
+    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()])