Fix via_device race in airzone_cloud (#177707)

homeassistant/components/airzone_cloud/__init__.py

@@ -2,11 +2,20 @@
 
 from aioairzone_cloud.cloudapi import AirzoneCloudApi
 from aioairzone_cloud.common import ConnectionOptions
+from aioairzone_cloud.const import (
+    AZD_FIRMWARE,
+    AZD_MODEL,
+    AZD_NAME,
+    AZD_SYSTEMS,
+    AZD_WEBSERVER,
+    AZD_WEBSERVERS,
+)
 
 from homeassistant.const import CONF_ID, CONF_PASSWORD, CONF_USERNAME, Platform
-from homeassistant.core import HomeAssistant
-from homeassistant.helpers import aiohttp_client
+from homeassistant.core import HomeAssistant, callback
+from homeassistant.helpers import aiohttp_client, device_registry as dr
 
+from .const import DOMAIN, MANUFACTURER
 from .coordinator import AirzoneCloudConfigEntry, AirzoneUpdateCoordinator
 
 PLATFORMS: list[Platform] = [
@@ -42,11 +51,53 @@ async def async_setup_entry(
 
     entry.runtime_data = coordinator
 
+    _async_register_devices(hass, entry, coordinator)
+
     await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
 
     return True
 
 
+@callback
+def _async_register_devices(
+    hass: HomeAssistant,
+    entry: AirzoneCloudConfigEntry,
+    coordinator: AirzoneUpdateCoordinator,
+) -> None:
+    """Register WebServer and System devices referenced as via_device parents.
+
+    Child devices resolve their via_device_id at add time, so the parents must
+    already exist regardless of which platform creates their own entities.
+    """
+    device_registry = dr.async_get(hass)
+
+    for ws_id, ws_data in coordinator.data.get(AZD_WEBSERVERS, {}).items():
+        device_registry.async_get_or_create(
+            config_entry_id=entry.entry_id,
+            connections={(dr.CONNECTION_NETWORK_MAC, ws_id)},
+            identifiers={(DOMAIN, ws_id)},
+            manufacturer=MANUFACTURER,
+            model="WebServer",
+            name=ws_data[AZD_NAME],
+            sw_version=ws_data[AZD_FIRMWARE],
+        )
+
+    for system_id, system_data in coordinator.data.get(AZD_SYSTEMS, {}).items():
+        device_registry.async_get_or_create(
+            config_entry_id=entry.entry_id,
+            identifiers={(DOMAIN, system_id)},
+            manufacturer=MANUFACTURER,
+            model=system_data.get(AZD_MODEL),
+            name=system_data[AZD_NAME],
+            sw_version=system_data.get(AZD_FIRMWARE),
+            via_device_id=dr.async_get_device_id_by_identifier(
+                hass,
+                (DOMAIN, system_data[AZD_WEBSERVER]),
+                config_entry_id=entry.entry_id,
+            ),
+        )
+
+
 async def async_unload_entry(
     hass: HomeAssistant, entry: AirzoneCloudConfigEntry
 ) -> bool:

homeassistant/components/airzone_cloud/entity.py

@@ -73,7 +73,11 @@ class AirzoneAidooEntity(AirzoneEntity):
             manufacturer=MANUFACTURER,
             model=aidoo_data[AZD_MODEL],
             name=aidoo_data[AZD_NAME],
-            via_device=(DOMAIN, aidoo_data[AZD_WEBSERVER]),
+            via_device_id=dr.async_get_device_id_by_identifier(
+                coordinator.hass,
+                (DOMAIN, aidoo_data[AZD_WEBSERVER]),
+                config_entry_id=coordinator.config_entry.entry_id,
+            ),
         )
 
     @override
@@ -164,7 +168,11 @@ class AirzoneHotWaterEntity(AirzoneEntity):
             manufacturer=MANUFACTURER,
             model="Hot Water",
             name=dhw_data[AZD_NAME],
-            via_device=(DOMAIN, dhw_data[AZD_WEBSERVER]),
+            via_device_id=dr.async_get_device_id_by_identifier(
+                coordinator.hass,
+                (DOMAIN, dhw_data[AZD_WEBSERVER]),
+                config_entry_id=coordinator.config_entry.entry_id,
+            ),
         )
 
     @override
@@ -257,7 +265,11 @@ class AirzoneSystemEntity(AirzoneEntity):
             model=system_data.get(AZD_MODEL),
             manufacturer=MANUFACTURER,
             name=system_data[AZD_NAME],
-            via_device=(DOMAIN, system_data[AZD_WEBSERVER]),
+            via_device_id=dr.async_get_device_id_by_identifier(
+                coordinator.hass,
+                (DOMAIN, system_data[AZD_WEBSERVER]),
+                config_entry_id=coordinator.config_entry.entry_id,
+            ),
             sw_version=system_data.get(AZD_FIRMWARE),
         )
 
@@ -322,7 +334,11 @@ class AirzoneZoneEntity(AirzoneEntity):
             model=zone_data.get(AZD_THERMOSTAT_MODEL),
             manufacturer=MANUFACTURER,
             name=zone_data[AZD_NAME],
-            via_device=(DOMAIN, self.system_id),
+            via_device_id=dr.async_get_device_id_by_identifier(
+                coordinator.hass,
+                (DOMAIN, self.system_id),
+                config_entry_id=coordinator.config_entry.entry_id,
+            ),
             sw_version=zone_data.get(AZD_THERMOSTAT_FW),
         )
 

tests/components/airzone_cloud/test_init.py

@@ -7,8 +7,9 @@ from aioairzone_cloud.exceptions import AirzoneTimeout
 from homeassistant.components.airzone_cloud.const import DOMAIN
 from homeassistant.config_entries import ConfigEntryState
 from homeassistant.core import HomeAssistant
+from homeassistant.helpers import device_registry as dr
 
-from .util import CONFIG
+from .util import CONFIG, WS_ID, WS_ID_AIDOO, async_init_integration
 
 from tests.common import MockConfigEntry
 
@@ -54,6 +55,51 @@ async def test_unload_entry(hass: HomeAssistant) -> None:
         assert config_entry.state is ConfigEntryState.NOT_LOADED
 
 
+async def test_device_via_device(
+    hass: HomeAssistant,
+    device_registry: dr.DeviceRegistry,
+) -> None:
+    """Test that child devices are linked to their via_device parents."""
+    await async_init_integration(hass)
+
+    config_entry = hass.config_entries.async_entries(DOMAIN)[0]
+
+    ws_device = device_registry.async_get_device_by_identifier(
+        (DOMAIN, WS_ID), config_entry.entry_id
+    )
+    assert ws_device is not None
+    assert ws_device.via_device_id is None
+
+    ws_aidoo_device = device_registry.async_get_device_by_identifier(
+        (DOMAIN, WS_ID_AIDOO), config_entry.entry_id
+    )
+    assert ws_aidoo_device is not None
+
+    system_device = device_registry.async_get_device_by_identifier(
+        (DOMAIN, "system1"), config_entry.entry_id
+    )
+    assert system_device is not None
+    assert system_device.via_device_id == ws_device.id
+
+    zone_device = device_registry.async_get_device_by_identifier(
+        (DOMAIN, "zone1"), config_entry.entry_id
+    )
+    assert zone_device is not None
+    assert zone_device.via_device_id == system_device.id
+
+    dhw_device = device_registry.async_get_device_by_identifier(
+        (DOMAIN, "dhw1"), config_entry.entry_id
+    )
+    assert dhw_device is not None
+    assert dhw_device.via_device_id == ws_device.id
+
+    aidoo_device = device_registry.async_get_device_by_identifier(
+        (DOMAIN, "aidoo1"), config_entry.entry_id
+    )
+    assert aidoo_device is not None
+    assert aidoo_device.via_device_id == ws_aidoo_device.id
+
+
 async def test_init_api_timeout(hass: HomeAssistant) -> None:
     """Test API timeouts when loading the Airzone Cloud integration."""