Fix via_device race in zinvolt (#178169)

homeassistant/components/zinvolt/__init__.py

@@ -8,8 +8,10 @@ from zinvolt.exceptions import ZinvoltError
 from homeassistant.const import CONF_ACCESS_TOKEN, Platform
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import ConfigEntryNotReady
+from homeassistant.helpers import device_registry as dr
 from homeassistant.helpers.aiohttp_client import async_get_clientsession
 
+from .const import DOMAIN
 from .coordinator import ZinvoltConfigEntry, ZinvoltDeviceCoordinator
 
 _PLATFORMS: list[Platform] = [
@@ -40,6 +42,15 @@ async def async_setup_entry(hass: HomeAssistant, entry: ZinvoltConfigEntry) -> b
 
     entry.runtime_data = coordinators
 
+    # Register the main battery devices before forwarding platforms so unit
+    # sub-devices can resolve their via_device_id parent.
+    device_registry = dr.async_get(hass)
+    for coordinator in coordinators.values():
+        device_registry.async_get_or_create(
+            config_entry_id=entry.entry_id,
+            identifiers={(DOMAIN, coordinator.data.battery.serial_number)},
+        )
+
     await hass.config_entries.async_forward_entry_setups(entry, _PLATFORMS)
 
     return True

homeassistant/components/zinvolt/coordinator.py

@@ -40,6 +40,7 @@ class BatteryData:
 class ZinvoltDeviceCoordinator(DataUpdateCoordinator[ZinvoltData]):
     """Class for Zinvolt devices."""
 
+    config_entry: ZinvoltConfigEntry
     battery_units: dict[str, Unit]
 
     def __init__(

homeassistant/components/zinvolt/entity.py

@@ -4,7 +4,7 @@ from typing import override
 
 from zinvolt.models import OnlineStatus, Unit
 
-from homeassistant.const import ATTR_VIA_DEVICE
+from homeassistant.helpers import device_registry as dr
 from homeassistant.helpers.device_registry import DeviceInfo
 from homeassistant.helpers.update_coordinator import CoordinatorEntity
 
@@ -65,9 +65,14 @@ class ZinvoltUnitEntity(ZinvoltEntity):
             model_id=self.battery.model,
         )
         if not is_main_device:
-            self._attr_device_info[ATTR_VIA_DEVICE] = (
-                DOMAIN,
-                coordinator.data.battery.serial_number,
+            # The main battery device is pre-registered in async_setup_entry, so it
+            # is guaranteed to exist here.
+            self._attr_device_info["via_device_id"] = (
+                dr.async_get_device_id_by_identifier(
+                    coordinator.hass,
+                    (DOMAIN, coordinator.data.battery.serial_number),
+                    config_entry_id=coordinator.config_entry.entry_id,
+                )
             )
 
     @property

tests/components/zinvolt/test_init.py

@@ -4,6 +4,7 @@ from unittest.mock import AsyncMock
 
 from syrupy.assertion import SnapshotAssertion
 
+from homeassistant.components.zinvolt.const import DOMAIN
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers import device_registry as dr
 
@@ -24,3 +25,24 @@ async def test_device(
     devices = device_registry.devices
     for device in devices.values():
         assert device == snapshot(name=list(device.identifiers)[0][1])
+
+
+async def test_device_via_device_links(
+    hass: HomeAssistant,
+    device_registry: dr.DeviceRegistry,
+    mock_config_entry: MockConfigEntry,
+    mock_zinvolt_client: AsyncMock,
+) -> None:
+    """Test that a unit sub-device links to its main battery device via via_device_id."""
+    await setup_integration(hass, mock_config_entry)
+
+    battery_device = device_registry.async_get_device_by_identifier(
+        (DOMAIN, "ZVG011025120088"), mock_config_entry.entry_id
+    )
+    assert battery_device is not None
+
+    unit_device = device_registry.async_get_device_by_identifier(
+        (DOMAIN, "BAT002"), mock_config_entry.entry_id
+    )
+    assert unit_device is not None
+    assert unit_device.via_device_id == battery_device.id