Fix via_device race in homewizard (#177934)

homeassistant/components/homewizard/__init__.py

@@ -17,6 +17,7 @@ from homeassistant.helpers.issue_registry import IssueSeverity, async_create_iss
 
 from .const import DOMAIN, PLATFORMS
 from .coordinator import HomeWizardConfigEntry, HWEnergyDeviceUpdateCoordinator
+from .entity import create_main_device_info
 
 
 async def async_setup_entry(hass: HomeAssistant, entry: HomeWizardConfigEntry) -> bool:
@@ -60,6 +61,14 @@ async def async_setup_entry(hass: HomeAssistant, entry: HomeWizardConfigEntry) -
         ):
             hass.config_entries.flow.async_abort(progress_flow["flow_id"])
 
+    # Register the main device up front so external sub-device sensors can
+    # resolve it as their via_device parent regardless of entity add order.
+    if coordinator.data.device.serial is not None:
+        dr.async_get(hass).async_get_or_create(
+            config_entry_id=entry.entry_id,
+            **create_main_device_info(coordinator),
+        )
+
     # Finalize
     entry.async_on_unload(coordinator.api.close)
     await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)

homeassistant/components/homewizard/entity.py

@@ -8,6 +8,23 @@ from .const import DOMAIN
 from .coordinator import HWEnergyDeviceUpdateCoordinator
 
 
+def create_main_device_info(
+    coordinator: HWEnergyDeviceUpdateCoordinator,
+) -> DeviceInfo:
+    """Return the device info for the main HomeWizard device."""
+    device_info = DeviceInfo(
+        manufacturer="HomeWizard",
+        sw_version=coordinator.data.device.firmware_version,
+        model_id=coordinator.data.device.product_type,
+        model=coordinator.data.device.model_name,
+    )
+    if (serial_number := coordinator.data.device.serial) is not None:
+        device_info[ATTR_CONNECTIONS] = {(CONNECTION_NETWORK_MAC, serial_number)}
+        device_info[ATTR_IDENTIFIERS] = {(DOMAIN, serial_number)}
+        device_info[ATTR_SERIAL_NUMBER] = serial_number
+    return device_info
+
+
 class HomeWizardEntity(CoordinatorEntity[HWEnergyDeviceUpdateCoordinator]):
     """Defines a HomeWizard entity."""
 
@@ -16,16 +33,4 @@ class HomeWizardEntity(CoordinatorEntity[HWEnergyDeviceUpdateCoordinator]):
     def __init__(self, coordinator: HWEnergyDeviceUpdateCoordinator) -> None:
         """Initialize the HomeWizard entity."""
         super().__init__(coordinator)
-        self._attr_device_info = DeviceInfo(
-            manufacturer="HomeWizard",
-            sw_version=coordinator.data.device.firmware_version,
-            model_id=coordinator.data.device.product_type,
-            model=coordinator.data.device.model_name,
-        )
-
-        if (serial_number := coordinator.data.device.serial) is not None:
-            self._attr_device_info[ATTR_CONNECTIONS] = {
-                (CONNECTION_NETWORK_MAC, serial_number)
-            }
-            self._attr_device_info[ATTR_IDENTIFIERS] = {(DOMAIN, serial_number)}
-            self._attr_device_info[ATTR_SERIAL_NUMBER] = serial_number
+        self._attr_device_info = create_main_device_info(coordinator)

homeassistant/components/homewizard/sensor.py

@@ -16,7 +16,6 @@ from homeassistant.components.sensor import (
     SensorStateClass,
 )
 from homeassistant.const import (
-    ATTR_VIA_DEVICE,
     PERCENTAGE,
     SIGNAL_STRENGTH_DECIBELS_MILLIWATT,
     EntityCategory,
@@ -31,7 +30,10 @@ from homeassistant.const import (
     UnitOfVolumeFlowRate,
 )
 from homeassistant.core import HomeAssistant
-from homeassistant.helpers.device_registry import DeviceInfo
+from homeassistant.helpers.device_registry import (
+    DeviceInfo,
+    async_get_device_id_by_identifier,
+)
 from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
 from homeassistant.helpers.typing import StateType
 from homeassistant.util.dt import utcnow
@@ -839,9 +841,10 @@ class HomeWizardExternalSensorEntity(HomeWizardEntity, SensorEntity):
             serial_number=device_unique_id,
         )
         if coordinator.data.device.serial is not None:
-            self._attr_device_info[ATTR_VIA_DEVICE] = (
-                DOMAIN,
-                coordinator.data.device.serial,
+            self._attr_device_info["via_device_id"] = async_get_device_id_by_identifier(
+                coordinator.hass,
+                (DOMAIN, coordinator.data.device.serial),
+                config_entry_id=coordinator.config_entry.entry_id,
             )
 
     @property

tests/components/homewizard/test_init.py

@@ -14,6 +14,7 @@ from homeassistant.components.homewizard.const import (
     battery_mode_cloud_issue_id,
 )
 from homeassistant.config_entries import SOURCE_REAUTH, ConfigEntryState
+from homeassistant.const import Platform
 from homeassistant.core import Event, HomeAssistant, callback
 from homeassistant.helpers import device_registry as dr, issue_registry as ir
 
@@ -324,3 +325,35 @@ async def test_battery_cloud_issue_stale_issue_cleared_on_reload(
     await hass.config_entries.async_reload(mock_config_entry.entry_id)
     await hass.async_block_till_done()
     assert ir.async_get(hass).async_get_issue(DOMAIN, issue_id) is None  # pylint: disable=home-assistant-tests-registry-fixtures
+
+
+async def test_main_device_registered_before_platform_forwarding(
+    hass: HomeAssistant,
+    device_registry: dr.DeviceRegistry,
+    mock_config_entry: MockConfigEntry,
+    mock_homewizardenergy: MagicMock,
+) -> None:
+    """Test the main device is registered before platforms are set up.
+
+    Restricting setup to the sensor platform removes the button platform,
+    whose main-device entity would otherwise register the main device
+    before the external sensor needs it, masking a missing up-front
+    registration.
+    """
+    mock_config_entry.add_to_hass(hass)
+    with patch("homeassistant.components.homewizard.PLATFORMS", [Platform.SENSOR]):
+        await hass.config_entries.async_setup(mock_config_entry.entry_id)
+        await hass.async_block_till_done()
+
+    assert mock_config_entry.state is ConfigEntryState.LOADED
+
+    main_device = device_registry.async_get_device_by_identifier(
+        (DOMAIN, "5c2fafabcdef"), mock_config_entry.entry_id
+    )
+    assert main_device is not None
+
+    gas_meter_device = device_registry.async_get_device_by_identifier(
+        (DOMAIN, "gas_meter_G001"), mock_config_entry.entry_id
+    )
+    assert gas_meter_device is not None
+    assert gas_meter_device.via_device_id == main_device.id

tests/components/homewizard/test_sensor.py

@@ -7,13 +7,13 @@ from homewizard_energy.models import CombinedModels, Measurement, State, System
 import pytest
 from syrupy.assertion import SnapshotAssertion
 
-from homeassistant.components.homewizard.const import UPDATE_INTERVAL
+from homeassistant.components.homewizard.const import DOMAIN, UPDATE_INTERVAL
 from homeassistant.const import STATE_UNAVAILABLE
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers import device_registry as dr, entity_registry as er
 from homeassistant.util import dt as dt_util
 
-from tests.common import async_fire_time_changed
+from tests.common import MockConfigEntry, async_fire_time_changed
 
 pytestmark = [
     pytest.mark.usefixtures("init_integration"),
@@ -393,6 +393,24 @@ async def test_sensors(
         assert snapshot(name=f"{entity_id}:device-registry") == device_entry
 
 
+async def test_external_device_via_device_id(
+    hass: HomeAssistant,
+    device_registry: dr.DeviceRegistry,
+    init_integration: MockConfigEntry,
+) -> None:
+    """Test an external device is linked to the main device via via_device_id."""
+    main_device = device_registry.async_get_device_by_identifier(
+        (DOMAIN, "5c2fafabcdef"), init_integration.entry_id
+    )
+    assert main_device is not None
+
+    gas_meter_device = device_registry.async_get_device_by_identifier(
+        (DOMAIN, "gas_meter_G001"), init_integration.entry_id
+    )
+    assert gas_meter_device is not None
+    assert gas_meter_device.via_device_id == main_device.id
+
+
 @pytest.mark.parametrize(
     ("device_fixture", "entity_ids"),
     [