Fix via_device race in eheimdigital (#177713)

homeassistant/components/eheimdigital/__init__.py

@@ -1,11 +1,17 @@
 """The EHEIM Digital integration."""
 
+from typing import TYPE_CHECKING
+
+from eheimdigital.device import EheimDigitalDevice
+
 from homeassistant.const import Platform
 from homeassistant.core import HomeAssistant
+from homeassistant.helpers import device_registry as dr
 from homeassistant.helpers.device_registry import DeviceEntry
 
 from .const import DOMAIN
 from .coordinator import EheimDigitalConfigEntry, EheimDigitalUpdateCoordinator
+from .entity import async_device_info
 
 PLATFORMS = [
     Platform.BINARY_SENSOR,
@@ -28,6 +34,17 @@ async def async_setup_entry(
     await coordinator.async_config_entry_first_refresh()
     entry.runtime_data = coordinator
 
+    main = coordinator.hub.main
+    if TYPE_CHECKING:
+        # After the first refresh at least one device is found and so there is
+        # always a main device set.
+        assert isinstance(main, EheimDigitalDevice)
+    # Register the main device up front so child devices can resolve their
+    # via_device_id link during concurrent platform setup.
+    dr.async_get(hass).async_get_or_create(
+        config_entry_id=entry.entry_id, **async_device_info(coordinator, main)
+    )
+
     await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
 
     return True

homeassistant/components/eheimdigital/entity.py

@@ -10,6 +10,7 @@ from eheimdigital.types import EheimDigitalClientError
 from homeassistant.const import CONF_HOST
 from homeassistant.core import callback
 from homeassistant.exceptions import HomeAssistantError
+from homeassistant.helpers import device_registry as dr
 from homeassistant.helpers.device_registry import CONNECTION_NETWORK_MAC, DeviceInfo
 from homeassistant.helpers.update_coordinator import CoordinatorEntity
 
@@ -17,6 +18,22 @@ from .const import DOMAIN
 from .coordinator import EheimDigitalUpdateCoordinator
 
 
+def async_device_info(
+    coordinator: EheimDigitalUpdateCoordinator, device: EheimDigitalDevice
+) -> DeviceInfo:
+    """Return the base device info for an EHEIM Digital device."""
+    return DeviceInfo(
+        configuration_url=f"http://{coordinator.config_entry.data[CONF_HOST]}",
+        name=device.name,
+        connections={(CONNECTION_NETWORK_MAC, device.mac_address)},
+        manufacturer="EHEIM",
+        model=device.model_name,
+        identifiers={(DOMAIN, device.mac_address)},
+        suggested_area=device.aquarium_name,
+        sw_version=device.sw_version,
+    )
+
+
 class EheimDigitalEntity[_DeviceT: EheimDigitalDevice](
     CoordinatorEntity[EheimDigitalUpdateCoordinator], ABC
 ):
@@ -29,21 +46,22 @@ class EheimDigitalEntity[_DeviceT: EheimDigitalDevice](
     ) -> None:
         """Initialize a EHEIM Digital entity."""
         super().__init__(coordinator)
+        main = coordinator.hub.main
         if TYPE_CHECKING:
             # At this point at least one device is found
             # and so there is always a main device set
-            assert isinstance(coordinator.hub.main, EheimDigitalDevice)
-        self._attr_device_info = DeviceInfo(
-            configuration_url=f"http://{coordinator.config_entry.data[CONF_HOST]}",
-            name=device.name,
-            connections={(CONNECTION_NETWORK_MAC, device.mac_address)},
-            manufacturer="EHEIM",
-            model=device.model_name,
-            identifiers={(DOMAIN, device.mac_address)},
-            suggested_area=device.aquarium_name,
-            sw_version=device.sw_version,
-            via_device=(DOMAIN, coordinator.hub.main.mac_address),
-        )
+            assert isinstance(main, EheimDigitalDevice)
+        self._attr_device_info = async_device_info(coordinator, device)
+        if device.mac_address != main.mac_address:
+            # The main device is registered during setup, before the platforms
+            # are forwarded, so this link always resolves deterministically.
+            self._attr_device_info["via_device_id"] = (
+                dr.async_get_device_id_by_identifier(
+                    coordinator.hass,
+                    (DOMAIN, main.mac_address),
+                    config_entry_id=coordinator.config_entry.entry_id,
+                )
+            )
         self._device = device
         self._device_address = device.mac_address
 

tests/components/eheimdigital/test_init.py

@@ -117,3 +117,32 @@ async def test_entry_setup_error(
     eheimdigital_hub_mock.return_value.connect.side_effect = EheimDigitalClientError()
     await init_integration(hass, mock_config_entry)
     assert mock_config_entry.state is ConfigEntryState.SETUP_RETRY
+
+
+async def test_child_device_via_device(
+    hass: HomeAssistant,
+    eheimdigital_hub_mock: MagicMock,
+    mock_config_entry: MockConfigEntry,
+    device_registry: dr.DeviceRegistry,
+) -> None:
+    """Test that child devices are linked to the main device."""
+    await init_integration(hass, mock_config_entry)
+
+    for device_address in eheimdigital_hub_mock.return_value.devices:
+        await eheimdigital_hub_mock.call_args.kwargs["device_found_callback"](
+            device_address,
+            eheimdigital_hub_mock.return_value.devices[device_address].device_type,
+        )
+    await hass.async_block_till_done()
+
+    main_device = device_registry.async_get_device_by_identifier(
+        (DOMAIN, "00:00:00:00:00:01"), mock_config_entry.entry_id
+    )
+    assert main_device is not None
+    assert main_device.via_device_id is None
+
+    child_device = device_registry.async_get_device_by_identifier(
+        (DOMAIN, "00:00:00:00:00:02"), mock_config_entry.entry_id
+    )
+    assert child_device is not None
+    assert child_device.via_device_id == main_device.id