Fix via_device race in hue (#177874)

homeassistant/components/hue/__init__.py

@@ -2,13 +2,12 @@
 
 from aiohue.util import normalize_bridge_id
 
-from homeassistant.components import persistent_notification
 from homeassistant.config_entries import SOURCE_IGNORE
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers import config_validation as cv, device_registry as dr
 from homeassistant.helpers.typing import ConfigType
 
-from .bridge import HueBridge, HueConfigEntry
+from .bridge import HueBridge, HueConfigEntry, _async_register_bridge_device
 from .const import DOMAIN
 from .migration import check_migration
 from .services import async_setup_services
@@ -70,47 +69,18 @@ async def async_setup_entry(hass: HomeAssistant, entry: HueConfigEntry) -> bool:
             hass.async_create_task(hass.config_entries.async_remove(entry.entry_id))
             return False
 
-    # add bridge device to device registry
-    device_registry = dr.async_get(hass)
-    if bridge.api_version == 1:
-        device_registry.async_get_or_create(
-            config_entry_id=entry.entry_id,
-            connections={(dr.CONNECTION_NETWORK_MAC, api.config.mac_address)},
-            identifiers={(DOMAIN, api.config.bridge_id)},
-            manufacturer="Signify",
-            name=api.config.name,
-            model_id=api.config.model_id,
-            sw_version=api.config.software_version,
-        )
-        # create persistent notification if we found a bridge version
-        # with security vulnerability
-        if (
-            api.config.model_id == "BSB002"
-            and api.config.software_version < "1935144040"
-        ):
-            persistent_notification.async_create(
-                hass,
-                (
-                    "Your Hue hub has a known security vulnerability ([CVE-2020-6007] "
-                    "(https://cve.circl.lu/cve/CVE-2020-6007)). "
-                    "Go to the Hue app and check for software updates."
-                ),
-                "Signify Hue",
-                "hue_hub_firmware",
-            )
-    else:
-        device_registry.async_get_or_create(
-            config_entry_id=entry.entry_id,
-            connections={(dr.CONNECTION_NETWORK_MAC, api.config.mac_address)},
-            identifiers={
-                (DOMAIN, api.config.bridge_id),
-                (DOMAIN, api.config.bridge_device.id),
-            },
-            manufacturer=api.config.bridge_device.product_data.manufacturer_name,
-            name=api.config.name,
-            model_id=api.config.model_id,
-            sw_version=api.config.software_version,
+    # v1 bridges already register their device before platform forwarding, so
+    # light/sensor entities can resolve it as their via_device parent; only
+    # register it here if that has not already happened. v2 bridges are always
+    # (re)registered here to merge the network MAC connection into the bridge
+    # device created by async_setup_devices with only its Zigbee MAC.
+    if bridge.api_version != 1 or (
+        dr.async_get(hass).async_get_device_by_identifier(
+            (DOMAIN, api.config.bridge_id), entry.entry_id
         )
+        is None
+    ):
+        _async_register_bridge_device(hass, entry, api, bridge.api_version)
 
     return True
 

homeassistant/components/hue/bridge.py

@@ -11,10 +11,11 @@ from aiohue import HueBridgeV1, HueBridgeV2, LinkButtonNotPressed, Unauthorized
 from aiohue.errors import AiohueException, BridgeBusy
 
 from homeassistant import core
+from homeassistant.components import persistent_notification
 from homeassistant.config_entries import SOURCE_IMPORT, ConfigEntry
 from homeassistant.const import CONF_API_KEY, CONF_API_VERSION, CONF_HOST, Platform
 from homeassistant.exceptions import ConfigEntryNotReady, HomeAssistantError
-from homeassistant.helpers import aiohttp_client
+from homeassistant.helpers import aiohttp_client, device_registry as dr
 
 from .const import DOMAIN
 from .v1.sensor_base import SensorManager
@@ -105,6 +106,13 @@ class HueBridge:
         if self.api_version == 1:
             if self.api.sensors is not None:
                 self.sensor_manager = SensorManager(self)
+            # Register the bridge device before forwarding the platforms so the
+            # light and sensor entities can resolve it as their via_device parent
+            # while they are being added.
+            if self.hass.config_entries.async_get_entry(self.config_entry.entry_id):
+                _async_register_bridge_device(
+                    self.hass, self.config_entry, self.api, self.api_version
+                )
             await self.hass.config_entries.async_forward_entry_setups(
                 self.config_entry, PLATFORMS_v1
             )
@@ -179,6 +187,56 @@ class HueBridge:
         create_config_flow(self.hass, self.host)
 
 
+@core.callback
+def _async_register_bridge_device(
+    hass: core.HomeAssistant,
+    config_entry: HueConfigEntry,
+    api: HueBridgeV1 | HueBridgeV2,
+    api_version: int,
+) -> None:
+    """Add the bridge device to the device registry."""
+    device_registry = dr.async_get(hass)
+    if api_version == 1:
+        device_registry.async_get_or_create(
+            config_entry_id=config_entry.entry_id,
+            connections={(dr.CONNECTION_NETWORK_MAC, api.config.mac_address)},
+            identifiers={(DOMAIN, api.config.bridge_id)},
+            manufacturer="Signify",
+            name=api.config.name,
+            model_id=api.config.model_id,
+            sw_version=api.config.software_version,
+        )
+        # create persistent notification if we found a bridge version
+        # with security vulnerability
+        if (
+            api.config.model_id == "BSB002"
+            and api.config.software_version < "1935144040"
+        ):
+            persistent_notification.async_create(
+                hass,
+                (
+                    "Your Hue hub has a known security vulnerability ([CVE-2020-6007] "
+                    "(https://cve.circl.lu/cve/CVE-2020-6007)). "
+                    "Go to the Hue app and check for software updates."
+                ),
+                "Signify Hue",
+                "hue_hub_firmware",
+            )
+    else:
+        device_registry.async_get_or_create(
+            config_entry_id=config_entry.entry_id,
+            connections={(dr.CONNECTION_NETWORK_MAC, api.config.mac_address)},
+            identifiers={
+                (DOMAIN, api.config.bridge_id),
+                (DOMAIN, api.config.bridge_device.id),
+            },
+            manufacturer=api.config.bridge_device.product_data.manufacturer_name,
+            name=api.config.name,
+            model_id=api.config.model_id,
+            sw_version=api.config.software_version,
+        )
+
+
 async def _update_listener(hass: core.HomeAssistant, entry: HueConfigEntry) -> None:
     """Handle ConfigEntry options update."""
     await hass.config_entries.async_reload(entry.entry_id)

homeassistant/components/hue/v1/light.py

@@ -29,6 +29,7 @@ from homeassistant.components.light import (
 )
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.exceptions import PlatformNotReady
+from homeassistant.helpers import device_registry as dr
 from homeassistant.helpers.debounce import Debouncer
 from homeassistant.helpers.device_registry import DeviceInfo
 from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
@@ -545,7 +546,11 @@ class HueLight(CoordinatorEntity, LightEntity):
             name=self.name,
             sw_version=self.light.swversion,
             suggested_area=suggested_area,
-            via_device=(DOMAIN, self.bridge.api.config.bridgeid),
+            via_device_id=dr.async_get_device_id_by_identifier(
+                self.bridge.hass,
+                (DOMAIN, self.bridge.api.config.bridgeid),
+                config_entry_id=self.bridge.config_entry.entry_id,
+            ),
         )
 
     @override

homeassistant/components/hue/v1/sensor_device.py

@@ -2,7 +2,7 @@
 
 from typing import override
 
-from homeassistant.helpers import entity
+from homeassistant.helpers import device_registry as dr, entity
 from homeassistant.helpers.device_registry import DeviceInfo
 
 from ..const import CONF_ALLOW_UNREACHABLE, DEFAULT_ALLOW_UNREACHABLE, DOMAIN
@@ -61,5 +61,9 @@ class GenericHueDevice(entity.Entity):  # pylint: disable=home-assistant-enforce
             model=(self.primary_sensor.productname or self.primary_sensor.modelid),
             name=self.primary_sensor.name,
             sw_version=self.primary_sensor.swversion,
-            via_device=(DOMAIN, self.bridge.api.config.bridgeid),
+            via_device_id=dr.async_get_device_id_by_identifier(
+                self.bridge.hass,
+                (DOMAIN, self.bridge.api.config.bridgeid),
+                config_entry_id=self.bridge.config_entry.entry_id,
+            ),
         )

homeassistant/components/hue/v2/device.py

@@ -18,7 +18,6 @@ from homeassistant.const import (
     ATTR_NAME,
     ATTR_SUGGESTED_AREA,
     ATTR_SW_VERSION,
-    ATTR_VIA_DEVICE,
 )
 from homeassistant.core import callback
 from homeassistant.helpers import device_registry as dr
@@ -49,7 +48,11 @@ async def async_setup_devices(bridge: HueBridge):
                 name=hue_resource.metadata.name,
                 model=hue_resource.type.value.replace("_", " ").title(),
                 manufacturer=api.config.bridge_device.product_data.manufacturer_name,
-                via_device=(DOMAIN, api.config.bridge_device.id),
+                via_device_id=dr.async_get_device_id_by_identifier(
+                    hass,
+                    (DOMAIN, api.config.bridge_device.id),
+                    config_entry_id=entry.entry_id,
+                ),
                 suggested_area=hue_resource.metadata.name
                 if hue_resource.type == ResourceTypes.ROOM
                 else None,
@@ -68,7 +71,13 @@ async def async_setup_devices(bridge: HueBridge):
         if hue_resource.id == api.config.bridge_device.id:
             params[ATTR_IDENTIFIERS].add((DOMAIN, api.config.bridge_id))
         else:
-            params[ATTR_VIA_DEVICE] = (DOMAIN, api.config.bridge_device.id)
+            # The bridge device is always registered first (see sort below), so
+            # its id can be resolved here for the via_device link.
+            params["via_device_id"] = dr.async_get_device_id_by_identifier(
+                hass,
+                (DOMAIN, api.config.bridge_device.id),
+                config_entry_id=entry.entry_id,
+            )
         zigbee = dev_controller.get_zigbee_connectivity(hue_resource.id)
         if zigbee and zigbee.mac_address:
             params[ATTR_CONNECTIONS] = {(dr.CONNECTION_NETWORK_MAC, zigbee.mac_address)}

tests/components/hue/conftest.py

@@ -128,6 +128,8 @@ def create_mock_api_v1() -> Mock:
         software_version="1935144040",
     )
     api.config.name = "Home"
+    # aiohue exposes the bridge id as both `bridge_id` and `bridgeid`
+    api.config.bridgeid = api.config.bridge_id
 
     api.lights = aiohue_v1.Lights(logger, {}, mock_request)
     api.groups = aiohue_v1.Groups(logger, {}, mock_request)

tests/components/hue/test_bridge.py

@@ -9,7 +9,7 @@ from aiohue.v1 import HueBridgeV1
 from aiohue.v2 import HueBridgeV2
 import pytest
 
-from homeassistant.components.hue import bridge
+from homeassistant.components.hue import bridge, migration
 from homeassistant.components.hue.const import (
     CONF_ALLOW_HUE_GROUPS,
     CONF_ALLOW_UNREACHABLE,
@@ -17,21 +17,41 @@ from homeassistant.components.hue.const import (
 )
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import ConfigEntryNotReady
+from homeassistant.helpers import device_registry as dr
+from homeassistant.util.json import JsonArrayType
 
-from tests.common import MockConfigEntry
+from .conftest import setup_platform
+from .test_light_v1 import LIGHT_RESPONSE
+
+from tests.common import MockConfigEntry, async_capture_events
 
 
-async def test_bridge_setup_v1(hass: HomeAssistant, mock_api_v1: Mock) -> None:
+async def test_bridge_setup_v1(
+    hass: HomeAssistant, mock_api_v1: Mock, device_registry: dr.DeviceRegistry
+) -> None:
     """Test a successful setup for V1 bridge."""
     config_entry = MockConfigEntry(
         domain=DOMAIN,
         data={"host": "1.2.3.4", "api_key": "mock-api-key", "api_version": 1},
         options={CONF_ALLOW_HUE_GROUPS: False, CONF_ALLOW_UNREACHABLE: False},
     )
+    config_entry.add_to_hass(hass)
+
+    def assert_bridge_device_registered(*args: object, **kwargs: object) -> None:
+        # The bridge device must already be registered by the time platforms
+        # are forwarded, so light/sensor entities can resolve it as their
+        # via_device parent while they are being added.
+        assert device_registry.async_get_device_by_identifier(
+            (DOMAIN, mock_api_v1.config.bridge_id), config_entry.entry_id
+        )
 
     with (
         patch.object(bridge, "HueBridgeV1", return_value=mock_api_v1),
-        patch.object(hass.config_entries, "async_forward_entry_setups") as mock_forward,
+        patch.object(
+            hass.config_entries,
+            "async_forward_entry_setups",
+            side_effect=assert_bridge_device_registered,
+        ) as mock_forward,
     ):
         hue_bridge = bridge.HueBridge(hass, config_entry)
         async with config_entry.setup_lock:
@@ -45,6 +65,88 @@ async def test_bridge_setup_v1(hass: HomeAssistant, mock_api_v1: Mock) -> None:
     assert forward_entries == {"light", "binary_sensor", "sensor"}
 
 
+async def test_bridge_device_v1(
+    hass: HomeAssistant, mock_api_v1: Mock, device_registry: dr.DeviceRegistry
+) -> None:
+    """Test the bridge device after a full v1 setup."""
+    config_entry = MockConfigEntry(
+        domain=DOMAIN,
+        data={"host": "1.2.3.4", "api_key": "mock-api-key", "api_version": 1},
+        options={CONF_ALLOW_HUE_GROUPS: False, CONF_ALLOW_UNREACHABLE: False},
+    )
+    config_entry.add_to_hass(hass)
+    mock_api_v1.mock_light_responses.append(LIGHT_RESPONSE)
+    mock_api_v1.mock_group_responses.append({})
+    mock_api_v1.mock_sensor_responses.append({})
+    events = async_capture_events(hass, dr.EVENT_DEVICE_REGISTRY_UPDATED)
+
+    with (
+        patch.object(bridge, "HueBridgeV1", return_value=mock_api_v1),
+        patch.object(migration, "is_v2_bridge", return_value=False),
+    ):
+        assert await hass.config_entries.async_setup(config_entry.entry_id)
+        await hass.async_block_till_done()
+
+    bridge_device = device_registry.async_get_device_by_identifier(
+        (DOMAIN, mock_api_v1.config.bridge_id), config_entry.entry_id
+    )
+    assert bridge_device is not None
+    assert bridge_device.connections == {
+        (dr.CONNECTION_NETWORK_MAC, mock_api_v1.config.mac_address)
+    }
+    # The bridge device is registered exactly once
+    create_events = [
+        event
+        for event in events
+        if event.data["action"] == "create"
+        and event.data["device_id"] == bridge_device.id
+    ]
+    assert len(create_events) == 1
+    # The light devices resolve the bridge device as their via_device parent
+    light_device = device_registry.async_get_device_by_identifier(
+        (DOMAIN, "456"), config_entry.entry_id
+    )
+    assert light_device is not None
+    assert light_device.via_device_id == bridge_device.id
+
+
+async def test_bridge_device_v2(
+    hass: HomeAssistant,
+    mock_bridge_v2: Mock,
+    v2_resources_test_data: JsonArrayType,
+    device_registry: dr.DeviceRegistry,
+) -> None:
+    """Test the bridge device after a fresh v2 setup."""
+    await mock_bridge_v2.api.load_test_data(v2_resources_test_data)
+    events = async_capture_events(hass, dr.EVENT_DEVICE_REGISTRY_UPDATED)
+    await setup_platform(hass, mock_bridge_v2, [])
+
+    bridge_device = device_registry.async_get_device_by_identifier(
+        (DOMAIN, mock_bridge_v2.api.config.bridge_id),
+        mock_bridge_v2.config_entry.entry_id,
+    )
+    assert bridge_device is not None
+    assert bridge_device.identifiers == {
+        (DOMAIN, mock_bridge_v2.api.config.bridge_id),
+        (DOMAIN, mock_bridge_v2.api.config.bridge_device.id),
+    }
+    # The bridge device has both the Zigbee MAC connection (set by
+    # async_setup_devices) and the network MAC connection (merged in by
+    # _async_register_bridge_device)
+    assert bridge_device.connections == {
+        (dr.CONNECTION_NETWORK_MAC, "00:17:88:01:aa:bb:fd:c7"),
+        (dr.CONNECTION_NETWORK_MAC, mock_bridge_v2.api.config.mac_address),
+    }
+    # The bridge device is registered exactly once
+    create_events = [
+        event
+        for event in events
+        if event.data["action"] == "create"
+        and event.data["device_id"] == bridge_device.id
+    ]
+    assert len(create_events) == 1
+
+
 async def test_bridge_setup_v2(hass: HomeAssistant, mock_api_v2: Mock) -> None:
     """Test a successful setup for V2 bridge."""
     config_entry = MockConfigEntry(

tests/components/hue/test_device_trigger_v2.py

@@ -73,6 +73,12 @@ async def test_get_triggers(
     hue_wall_switch_device = device_registry.async_get_device(
         identifiers={(hue.DOMAIN, "3ff06175-29e8-44a8-8fe7-af591b0025da")}
     )
+    # The device is linked to the bridge device as its via_device.
+    bridge_device = device_registry.async_get_device_by_identifier(
+        (hue.DOMAIN, mock_bridge_v2.api.config.bridge_id),
+        mock_bridge_v2.config_entry.entry_id,
+    )
+    assert hue_wall_switch_device.via_device_id == bridge_device.id
     hue_bat_sensor = entity_registry.async_get(
         "sensor.wall_switch_with_2_controls_battery"
     )

tests/components/hue/test_light_v1.py

@@ -191,6 +191,12 @@ async def setup_bridge(hass: HomeAssistant, mock_bridge_v1: Mock) -> None:
     config_entry.mock_state(hass, ConfigEntryState.LOADED)
     mock_bridge_v1.config_entry = config_entry
     config_entry.runtime_data = mock_bridge_v1
+    # Register the bridge device so the light entities can resolve it as their
+    # via_device parent while they are being added.
+    dr.async_get(hass).async_get_or_create(
+        config_entry_id=config_entry.entry_id,
+        identifiers={(hue.DOMAIN, mock_bridge_v1.api.config.bridgeid)},
+    )
     await hass.config_entries.async_forward_entry_setups(config_entry, [Platform.LIGHT])
     # To flush out the service call to update the group
     await hass.async_block_till_done()

tests/components/hue/test_sensor_v1.py

@@ -488,6 +488,12 @@ async def test_hue_events(
     hue_tap_device = device_registry.async_get_device(
         identifiers={(hue.DOMAIN, "00:00:00:00:00:44:23:08")}
     )
+    # The sensor device is linked to the bridge device as its via_device.
+    bridge_device = device_registry.async_get_device_by_identifier(
+        (hue.DOMAIN, mock_bridge_v1.api.config.bridgeid),
+        mock_bridge_v1.config_entry.entry_id,
+    )
+    assert hue_tap_device.via_device_id == bridge_device.id
 
     mock_bridge_v1.api.sensors["7"].last_event = {"type": "button"}
     mock_bridge_v1.api.sensors["8"].last_event = {"type": "button"}