Fix via_device race in toon (#177774)

homeassistant/components/toon/__init__.py

@@ -113,20 +113,47 @@ async def async_setup_entry(hass: HomeAssistant, entry: ToonConfigEntry) -> bool
 
     entry.runtime_data = coordinator
 
-    # Register device for the Meter Adapter, since it will have no entities.
+    agreement = coordinator.data.agreement
+    agreement_id = agreement.agreement_id
+
+    # Register the parent devices before forwarding the platforms, so that the
+    # child devices created by the entities can deterministically resolve their
+    # via_device_id.
     device_registry = dr.async_get(hass)
-    device_registry.async_get_or_create(
+    display_device = device_registry.async_get_or_create(
+        config_entry_id=entry.entry_id,
+        identifiers={(DOMAIN, agreement_id)},
+        manufacturer="Eneco",
+        model=agreement.display_hardware_version.rpartition("/")[0],
+        name="Toon Display",
+        sw_version=agreement.display_software_version.rpartition("/")[-1],
+    )
+    # The Meter Adapter has no entities of its own.
+    meter_adapter_device = device_registry.async_get_or_create(
         config_entry_id=entry.entry_id,
         identifiers={
-            (
-                DOMAIN,
-                coordinator.data.agreement.agreement_id,
-                "meter_adapter",
-            )  # type: ignore[arg-type]
+            (DOMAIN, agreement_id, "meter_adapter"),  # type: ignore[arg-type]
         },
         manufacturer="Eneco",
         name="Meter Adapter",
-        via_device=(DOMAIN, coordinator.data.agreement.agreement_id),
+        via_device_id=display_device.id,
+    )
+    device_registry.async_get_or_create(
+        config_entry_id=entry.entry_id,
+        identifiers={
+            (DOMAIN, agreement_id, "electricity"),  # type: ignore[arg-type]
+        },
+        name="Electricity Meter",
+        via_device_id=meter_adapter_device.id,
+    )
+    device_registry.async_get_or_create(
+        config_entry_id=entry.entry_id,
+        identifiers={
+            (DOMAIN, agreement_id, "boiler_module"),  # type: ignore[arg-type]
+        },
+        manufacturer="Eneco",
+        name="Boiler Module",
+        via_device_id=display_device.id,
     )
 
     # Spin up the platforms

homeassistant/components/toon/entity.py

@@ -3,6 +3,7 @@
 from dataclasses import dataclass
 from typing import override
 
+from homeassistant.helpers import device_registry as dr
 from homeassistant.helpers.device_registry import DeviceInfo
 from homeassistant.helpers.update_coordinator import CoordinatorEntity
 
@@ -44,10 +45,10 @@ class ToonElectricityMeterDeviceEntity(ToonEntity):
             identifiers={
                 (DOMAIN, agreement_id, "electricity"),  # type: ignore[arg-type]
             },
-            via_device=(
-                DOMAIN,
-                agreement_id,  # type: ignore[typeddict-item]
-                "meter_adapter",
+            via_device_id=dr.async_get_device_id_by_identifier(
+                self.coordinator.hass,
+                (DOMAIN, agreement_id, "meter_adapter"),  # type: ignore[arg-type]
+                config_entry_id=self.coordinator.config_entry.entry_id,
             ),
         )
 
@@ -65,10 +66,10 @@ class ToonGasMeterDeviceEntity(ToonEntity):
             identifiers={
                 (DOMAIN, agreement_id, "gas"),  # type: ignore[arg-type]
             },
-            via_device=(
-                DOMAIN,
-                agreement_id,  # type: ignore[typeddict-item]
-                "electricity",
+            via_device_id=dr.async_get_device_id_by_identifier(
+                self.coordinator.hass,
+                (DOMAIN, agreement_id, "electricity"),  # type: ignore[arg-type]
+                config_entry_id=self.coordinator.config_entry.entry_id,
             ),
         )
 
@@ -86,10 +87,10 @@ class ToonWaterMeterDeviceEntity(ToonEntity):
             identifiers={
                 (DOMAIN, agreement_id, "water"),  # type: ignore[arg-type]
             },
-            via_device=(
-                DOMAIN,
-                agreement_id,  # type: ignore[typeddict-item]
-                "electricity",
+            via_device_id=dr.async_get_device_id_by_identifier(
+                self.coordinator.hass,
+                (DOMAIN, agreement_id, "electricity"),  # type: ignore[arg-type]
+                config_entry_id=self.coordinator.config_entry.entry_id,
             ),
         )
 
@@ -107,10 +108,10 @@ class ToonSolarDeviceEntity(ToonEntity):
             identifiers={
                 (DOMAIN, agreement_id, "solar"),  # type: ignore[arg-type]
             },
-            via_device=(
-                DOMAIN,
-                agreement_id,  # type: ignore[typeddict-item]
-                "meter_adapter",
+            via_device_id=dr.async_get_device_id_by_identifier(
+                self.coordinator.hass,
+                (DOMAIN, agreement_id, "meter_adapter"),  # type: ignore[arg-type]
+                config_entry_id=self.coordinator.config_entry.entry_id,
             ),
         )
 
@@ -133,7 +134,11 @@ class ToonBoilerModuleDeviceEntity(ToonEntity):
                     "boiler_module",
                 )
             },
-            via_device=(DOMAIN, agreement_id),
+            via_device_id=dr.async_get_device_id_by_identifier(
+                self.coordinator.hass,
+                (DOMAIN, agreement_id),
+                config_entry_id=self.coordinator.config_entry.entry_id,
+            ),
         )
 
 
@@ -150,10 +155,10 @@ class ToonBoilerDeviceEntity(ToonEntity):
             identifiers={
                 (DOMAIN, agreement_id, "boiler"),  # type: ignore[arg-type]
             },
-            via_device=(
-                DOMAIN,
-                agreement_id,  # type: ignore[typeddict-item]
-                "boiler_module",
+            via_device_id=dr.async_get_device_id_by_identifier(
+                self.coordinator.hass,
+                (DOMAIN, agreement_id, "boiler_module"),  # type: ignore[arg-type]
+                config_entry_id=self.coordinator.config_entry.entry_id,
             ),
         )
 

tests/components/toon/test_init.py

@@ -1,10 +1,15 @@
 """Tests for the Toon component."""
 
+import time
 from unittest.mock import patch
 
+from toonapi import Agreement, Status
+from toonapi.models import ThermostatInfo
+
 from homeassistant.components.toon import DOMAIN
 from homeassistant.config_entries import ConfigEntryState
 from homeassistant.core import HomeAssistant
+from homeassistant.helpers import config_entry_oauth2_flow, device_registry as dr
 from homeassistant.helpers.config_entry_oauth2_flow import (
     ImplementationUnavailableError,
 )
@@ -66,3 +71,90 @@ async def test_migrate_entry_minor_version_2_2(hass: HomeAssistant) -> None:
         assert entry.version == 2
         assert entry.minor_version == 2
         assert entry.unique_id == "123"
+
+
+async def test_device_registry_via_devices(
+    hass: HomeAssistant,
+    device_registry: dr.DeviceRegistry,
+) -> None:
+    """Test that child devices are linked to their parent via via_device_id."""
+    config_entry = MockConfigEntry(
+        domain=DOMAIN,
+        version=2,
+        minor_version=2,
+        unique_id="test-agreement-id",
+        data={
+            "auth_implementation": DOMAIN,
+            "token": {
+                "refresh_token": "mock-refresh-token",
+                "access_token": "mock-access-token",
+                "type": "Bearer",
+                "expires_in": 60,
+                "expires_at": time.time() + 3600,
+            },
+            "agreement_id": "test-agreement-id",
+        },
+    )
+    config_entry.add_to_hass(hass)
+
+    config_entry_oauth2_flow.async_register_implementation(
+        hass,
+        DOMAIN,
+        config_entry_oauth2_flow.LocalOAuth2Implementation(
+            hass,
+            DOMAIN,
+            "client-id",
+            "client-secret",
+            "https://api.toon.eu/authorize",
+            "https://api.toon.eu/token",
+        ),
+    )
+
+    agreement = Agreement(
+        agreement_id="test-agreement-id",
+        display_common_name="display-common-name",
+        display_hardware_version="qb2/ICY/v0.8",
+        display_software_version="qb2/v1.2",
+        heating_type="gas",
+        is_toon_solar=True,
+    )
+    status = Status(agreement=agreement)
+    status.thermostat = ThermostatInfo(have_opentherm_boiler=True)
+
+    with (
+        patch("toonapi.Toon.activate_agreement"),
+        patch("toonapi.Toon.update", return_value=status),
+        patch(
+            "homeassistant.components.toon.coordinator."
+            "ToonDataUpdateCoordinator.register_webhook"
+        ),
+    ):
+        assert await hass.config_entries.async_setup(config_entry.entry_id)
+        await hass.async_block_till_done()
+
+    assert config_entry.state is ConfigEntryState.LOADED
+
+    def get_device(*identifier: str) -> dr.DeviceEntry:
+        device = device_registry.async_get_device_by_identifier(
+            (DOMAIN, *identifier),  # type: ignore[arg-type]
+            config_entry.entry_id,
+        )
+        assert device is not None
+        return device
+
+    display = get_device("test-agreement-id")
+    assert display.via_device_id is None
+
+    meter_adapter = get_device("test-agreement-id", "meter_adapter")
+    assert meter_adapter.via_device_id == display.id
+
+    electricity = get_device("test-agreement-id", "electricity")
+    assert electricity.via_device_id == meter_adapter.id
+
+    boiler_module = get_device("test-agreement-id", "boiler_module")
+    assert boiler_module.via_device_id == display.id
+
+    assert get_device("test-agreement-id", "gas").via_device_id == electricity.id
+    assert get_device("test-agreement-id", "water").via_device_id == electricity.id
+    assert get_device("test-agreement-id", "solar").via_device_id == meter_adapter.id
+    assert get_device("test-agreement-id", "boiler").via_device_id == boiler_module.id