Adapt qbus to set via_device_id in DeviceInfo (#177750)

homeassistant/components/qbus/entity.py

@@ -3,13 +3,14 @@
 from abc import ABC, abstractmethod
 from collections.abc import Callable
 import re
-from typing import cast, override
+from typing import TYPE_CHECKING, cast, override
 
 from qbusmqttapi.discovery import QbusMqttDevice, QbusMqttOutput
 from qbusmqttapi.factory import QbusMqttMessageFactory, QbusMqttTopicFactory
 from qbusmqttapi.state import QbusMqttState
 
 from homeassistant.components.mqtt import ReceiveMessage, client as mqtt
+from homeassistant.helpers import device_registry as dr
 from homeassistant.helpers.device_registry import DeviceInfo, format_mac
 from homeassistant.helpers.entity import Entity
 
@@ -92,6 +93,7 @@ class QbusEntity[StateT: QbusMqttState](Entity, ABC):
         """Initialize the Qbus entity."""
 
         self._mqtt_output = mqtt_output
+        self._link_to_main_device = link_to_main_device
 
         self._topic_factory = QbusMqttTopicFactory()
         self._message_factory = QbusMqttMessageFactory()
@@ -99,8 +101,8 @@ class QbusEntity[StateT: QbusMqttState](Entity, ABC):
             mqtt_output.device.id, mqtt_output.id
         )
 
-        ref_id = format_ref_id(mqtt_output.ref_id)
-        suffix = ref_id or ""
+        self._ref_id = format_ref_id(mqtt_output.ref_id)
+        suffix = self._ref_id or ""
 
         if id_suffix:
             suffix += f"_{id_suffix}"
@@ -109,19 +111,32 @@ class QbusEntity[StateT: QbusMqttState](Entity, ABC):
             mqtt_output.device.serial_number, suffix
         )
 
-        if link_to_main_device:
-            self._attr_device_info = DeviceInfo(
-                identifiers={create_device_identifier(mqtt_output.device)}
-            )
-        else:
-            self._attr_device_info = DeviceInfo(
-                name=mqtt_output.name.title(),
-                manufacturer=MANUFACTURER,
-                identifiers={(DOMAIN, f"{mqtt_output.device.serial_number}_{ref_id}")},
-                suggested_area=mqtt_output.location.title(),
-                via_device=create_device_identifier(mqtt_output.device),
+    @property
+    @override
+    def device_info(self) -> DeviceInfo:
+        """Return the device info."""
+        if self._link_to_main_device:
+            return DeviceInfo(
+                identifiers={create_device_identifier(self._mqtt_output.device)}
             )
 
+        config_entry = self.platform.config_entry
+        if TYPE_CHECKING:
+            assert config_entry is not None
+        return DeviceInfo(
+            name=self._mqtt_output.name.title(),
+            manufacturer=MANUFACTURER,
+            identifiers={
+                (DOMAIN, f"{self._mqtt_output.device.serial_number}_{self._ref_id}")
+            },
+            suggested_area=self._mqtt_output.location.title(),
+            via_device_id=dr.async_get_device_id_by_identifier(
+                self.hass,
+                create_device_identifier(self._mqtt_output.device),
+                config_entry_id=config_entry.entry_id,
+            ),
+        )
+
     @override
     async def async_added_to_hass(self) -> None:
         """Run when entity about to be added to hass."""

tests/components/qbus/test_light.py

@@ -3,6 +3,7 @@
 from collections.abc import Awaitable, Callable
 from unittest.mock import patch
 
+import pytest
 from syrupy.assertion import SnapshotAssertion
 
 from homeassistant.components.light import (
@@ -14,9 +15,10 @@ from homeassistant.components.light import (
     SERVICE_TURN_OFF,
     SERVICE_TURN_ON,
 )
+from homeassistant.components.qbus.const import DOMAIN
 from homeassistant.const import ATTR_ENTITY_ID, STATE_OFF, STATE_ON, Platform
 from homeassistant.core import HomeAssistant
-from homeassistant.helpers import entity_registry as er
+from homeassistant.helpers import device_registry as dr, entity_registry as er
 
 from tests.common import MockConfigEntry, async_fire_mqtt_message, snapshot_platform
 from tests.typing import MqttMockHAClient
@@ -95,6 +97,8 @@ _TOPIC_COLOR_SET_STATE = "cloudapp/QBUSMQTTGW/UL1/UL100/setState"
 _DIMMER_ENTITY_ID = "light.media_room_media_room"
 _COLOR_ENTITY_ID = "light.media_room_tv"
 
+_CONTROLLER_MAC = "001122334455"
+
 
 async def test_light(
     hass: HomeAssistant,
@@ -111,6 +115,29 @@ async def test_light(
     await snapshot_platform(hass, entity_registry, snapshot, mock_config_entry.entry_id)
 
 
+@pytest.mark.usefixtures("setup_integration")
+async def test_light_via_device_id(
+    hass: HomeAssistant,
+    entity_registry: er.EntityRegistry,
+    device_registry: dr.DeviceRegistry,
+    mock_config_entry: MockConfigEntry,
+) -> None:
+    """Test the light's device links to the controller device via via_device_id."""
+    entity_entry = entity_registry.async_get(_DIMMER_ENTITY_ID)
+    assert entity_entry
+    assert entity_entry.device_id
+
+    device = device_registry.async_get(entity_entry.device_id)
+    assert device
+
+    controller_device = device_registry.async_get_device_by_identifier(
+        (DOMAIN, dr.format_mac(_CONTROLLER_MAC)), mock_config_entry.entry_id
+    )
+    assert controller_device
+
+    assert device.via_device_id == controller_device.id
+
+
 async def test_dimmer(
     hass: HomeAssistant,
     mqtt_mock: MqttMockHAClient,