Fix via_device race in tplink

homeassistant/components/tplink/__init__.py

@@ -242,6 +242,15 @@ async def async_setup_entry(hass: HomeAssistant, entry: TPLinkConfigEntry) -> bo
     live_view = entry.data.get(CONF_LIVE_VIEW)
 
     entry.runtime_data = TPLinkData(parent_coordinator, camera_creds, live_view)
+
+    # Register the parent device before forwarding platforms so child device
+    # entities can resolve their via_device_id from it.
+    device_registry = dr.async_get(hass)
+    device_registry.async_get_or_create(
+        config_entry_id=entry.entry_id,
+        identifiers={(DOMAIN, device.device_id)},
+    )
+
     await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
 
     return True

homeassistant/components/tplink/entity.py

@@ -236,7 +236,13 @@ class CoordinatedTPLinkEntity(CoordinatorEntity[TPLinkDataUpdateCoordinator], AB
             and parent != registry_device
             and parent.device_type is not Device.Type.WallSwitch
         ):
-            self._attr_device_info["via_device"] = (DOMAIN, parent.device_id)
+            self._attr_device_info["via_device_id"] = (
+                dr.async_get_device_id_by_identifier(
+                    self.coordinator.hass,
+                    (DOMAIN, parent.device_id),
+                    config_entry_id=self.coordinator.config_entry.entry_id,
+                )
+            )
         else:
             self._attr_device_info["connections"] = {
                 (dr.CONNECTION_NETWORK_MAC, device.mac)
@@ -642,7 +648,6 @@ class CoordinatedTPLinkModuleEntity(CoordinatedTPLinkEntity, ABC):
                     platform_domain=platform_domain,
                 )
             )
-            has_parent_entities = bool(entities)
 
         children = _get_new_children(
             device, coordinator, known_child_device_ids, entity_class.__name__
@@ -678,16 +683,6 @@ class CoordinatedTPLinkModuleEntity(CoordinatedTPLinkEntity, ABC):
             )
             entities.extend(child_entities)
 
-        if first_check and entities and not has_parent_entities:
-            # Get or create the parent device for via_device.
-            # This is a timing factor in case this platform is loaded before
-            # other platforms that will have entities on the parent. Eventually
-            # those other platforms will update the parent with full DeviceInfo
-            device_registry = dr.async_get(hass)
-            device_registry.async_get_or_create(
-                config_entry_id=coordinator.config_entry.entry_id,
-                identifiers={(DOMAIN, device.device_id)},
-            )
         return entities
 
 

tests/components/tplink/test_init.py

@@ -46,6 +46,7 @@ from homeassistant.util import dt as dt_util
 from . import (
     _mocked_device,
     _mocked_feature,
+    _mocked_strip_children,
     _patch_connect,
     _patch_discovery,
     _patch_single_discovery,
@@ -1245,3 +1246,84 @@ async def test_automatic_device_addition_does_not_remove_disabled_default(
     check_entities("hub")
     for child_id in (1, 2, 3):
         check_entities(f"child_{child_id}")
+
+
+async def test_device_via_device_links(
+    hass: HomeAssistant,
+    device_registry: dr.DeviceRegistry,
+) -> None:
+    """Test that strip child devices link to the parent via via_device_id."""
+    config_entry = MockConfigEntry(
+        domain=DOMAIN, data={CONF_HOST: "127.0.0.1"}, unique_id=MAC_ADDRESS
+    )
+    config_entry.add_to_hass(hass)
+    feature = _mocked_feature(
+        "consumption_this_month",
+        value=5.2,
+        type_=Feature.Type.Sensor,
+        category=Feature.Category.Primary,
+    )
+    parent = _mocked_device(
+        alias="my_plug",
+        features=[feature],
+        children=_mocked_strip_children(features=[feature]),
+        device_type=DeviceType.Strip,
+    )
+    with _patch_discovery(device=parent), _patch_connect(device=parent):
+        await hass.config_entries.async_setup(config_entry.entry_id)
+        await hass.async_block_till_done()
+
+    parent_device = device_registry.async_get_device_by_identifier(
+        (DOMAIN, parent.device_id), config_entry.entry_id
+    )
+    assert parent_device is not None
+
+    assert parent.children
+    for child in parent.children:
+        child_device = device_registry.async_get_device_by_identifier(
+            (DOMAIN, child.device_id), config_entry.entry_id
+        )
+        assert child_device is not None
+        assert child_device.id != parent_device.id
+        assert child_device.via_device_id == parent_device.id
+
+
+async def test_wall_switch_child_uses_connections(
+    hass: HomeAssistant,
+    device_registry: dr.DeviceRegistry,
+) -> None:
+    """Test that WallSwitch child devices merge with the parent via connections."""
+    config_entry = MockConfigEntry(
+        domain=DOMAIN, data={CONF_HOST: "127.0.0.1"}, unique_id=MAC_ADDRESS
+    )
+    config_entry.add_to_hass(hass)
+    feature = _mocked_feature(
+        "consumption_this_month",
+        value=5.2,
+        type_=Feature.Type.Sensor,
+        category=Feature.Category.Primary,
+    )
+    parent = _mocked_device(
+        alias="my_plug",
+        features=[feature],
+        children=_mocked_strip_children(features=[feature]),
+        device_type=DeviceType.WallSwitch,
+    )
+    with _patch_discovery(device=parent), _patch_connect(device=parent):
+        await hass.config_entries.async_setup(config_entry.entry_id)
+        await hass.async_block_till_done()
+
+    parent_device = device_registry.async_get_device_by_identifier(
+        (DOMAIN, parent.device_id), config_entry.entry_id
+    )
+    assert parent_device is not None
+    assert parent_device.via_device_id is None
+
+    assert parent.children
+    for child in parent.children:
+        # WallSwitch children merge into the parent device via the mac connection
+        child_device = device_registry.async_get_device_by_identifier(
+            (DOMAIN, child.device_id), config_entry.entry_id
+        )
+        assert child_device is not None
+        assert child_device.id == parent_device.id