Adapt mqtt to set via_device_id in DeviceInfo (#178234)

homeassistant/components/mqtt/entity.py

@@ -19,7 +19,6 @@ from homeassistant.const import (
     ATTR_SERIAL_NUMBER,
     ATTR_SUGGESTED_AREA,
     ATTR_SW_VERSION,
-    ATTR_VIA_DEVICE,
     CONF_DEVICE,
     CONF_ENTITY_CATEGORY,
     CONF_ICON,
@@ -1301,9 +1300,6 @@ def device_info_from_specifications(
     if CONF_SW_VERSION in specifications:
         info[ATTR_SW_VERSION] = specifications[CONF_SW_VERSION]
 
-    if CONF_VIA_DEVICE in specifications:
-        info[ATTR_VIA_DEVICE] = (DOMAIN, specifications[CONF_VIA_DEVICE])
-
     if CONF_SUGGESTED_AREA in specifications:
         info[ATTR_SUGGESTED_AREA] = specifications[CONF_SUGGESTED_AREA]
 
@@ -1313,33 +1309,67 @@ def device_info_from_specifications(
     return info
 
 
+def _via_device_identifier(
+    specifications: dict[str, Any] | None,
+) -> tuple[str, str] | None:
+    """Return the device registry identifier of the referenced via device."""
+    if not specifications or CONF_VIA_DEVICE not in specifications:
+        return None
+    return (DOMAIN, specifications[CONF_VIA_DEVICE])
+
+
 @callback
 def ensure_via_device_exists(
-    hass: HomeAssistant, device_info: DeviceInfo | None, config_entry: ConfigEntry
+    hass: HomeAssistant,
+    specifications: dict[str, Any] | None,
+    config_entry: ConfigEntry,
 ) -> None:
-    """Ensure the via device is in the device registry."""
-    if (
-        device_info is None
-        or CONF_VIA_DEVICE not in device_info
-        or (device_registry := dr.async_get(hass)).async_get_device_by_identifier(
-            device_info["via_device"], config_entry.entry_id
-        )
+    """Ensure the via device is in the device registry.
+
+    MQTT discovery can announce a child device before its via device, so the
+    referenced parent is stub-created here when it does not yet exist.
+    """
+    if (identifier := _via_device_identifier(specifications)) is None:
+        return
+
+    device_registry = dr.async_get(hass)
+    if device_registry.async_get_device_by_identifier(
+        identifier, config_entry.entry_id
     ):
         return
 
-    # Ensure the via device exists in the device registry
     _LOGGER.debug(
-        "Device identifier %s via_device reference from device_info %s "
-        "not found in the Device Registry, creating new entry",
-        device_info["via_device"],
-        device_info,
+        "Device identifier %s referenced as via_device not found in the "
+        "Device Registry, creating new entry",
+        identifier,
     )
     device_registry.async_get_or_create(
         config_entry_id=config_entry.entry_id,
-        identifiers={device_info["via_device"]},
+        identifiers={identifier},
     )
 
 
+@callback
+def _resolve_via_device_id(
+    hass: HomeAssistant,
+    specifications: dict[str, Any] | None,
+    config_entry: ConfigEntry,
+) -> str | None:
+    """Resolve the referenced via device to its device registry id.
+
+    Best-effort read-only lookup: returns None when no via device is referenced
+    or the referenced device is not (yet) registered. It never raises, so a
+    missing parent cannot abort entity setup given MQTT's non-deterministic
+    discovery order.
+    """
+    if (identifier := _via_device_identifier(specifications)) is None:
+        return None
+    via_device = dr.async_get(hass).async_get_device_by_identifier(
+        identifier, config_entry.entry_id
+    )
+    return via_device.id if via_device else None
+
+
 class MqttEntityDeviceInfo(Entity):
     """Mixin used for mqtt platforms that support the device registry."""
 
@@ -1355,10 +1385,12 @@ class MqttEntityDeviceInfo(Entity):
         self._device_specifications = config.get(CONF_DEVICE)
         device_registry = dr.async_get(self.hass)
         config_entry_id = self._config_entry.entry_id
+        ensure_via_device_exists(
+            self.hass, self._device_specifications, self._config_entry
+        )
         device_info = self.device_info
 
         if device_info is not None:
-            ensure_via_device_exists(self.hass, device_info, self._config_entry)
             device_registry.async_get_or_create(
                 config_entry_id=config_entry_id, **device_info
             )
@@ -1367,7 +1399,14 @@ class MqttEntityDeviceInfo(Entity):
     @override
     def device_info(self) -> DeviceInfo | None:
         """Return a device description for device registry."""
-        return device_info_from_specifications(self._device_specifications)
+        info = device_info_from_specifications(self._device_specifications)
+        if info is not None and (
+            via_device_id := _resolve_via_device_id(
+                self.hass, self._device_specifications, self._config_entry
+            )
+        ):
+            info["via_device_id"] = via_device_id
+        return info
 
 
 class MqttEntity(
@@ -1414,7 +1453,9 @@ class MqttEntity(
             self, hass, discovery_data, self.discovery_update
         )
         MqttEntityDeviceInfo.__init__(self, config.get(CONF_DEVICE), config_entry)
-        ensure_via_device_exists(self.hass, self.device_info, self._config_entry)
+        ensure_via_device_exists(
+            self.hass, self._device_specifications, self._config_entry
+        )
 
     def _init_entity_registry(self, discovery_data: DiscoveryInfoType | None) -> None:
         """Set entity_id from default_entity_id if defined in config.
@@ -1733,11 +1774,14 @@ def update_device(
     device: DeviceEntry | None = None
     device_registry = dr.async_get(hass)
     config_entry_id = config_entry.entry_id
-    device_info = device_info_from_specifications(config[CONF_DEVICE])
+    specifications = config[CONF_DEVICE]
 
-    ensure_via_device_exists(hass, device_info, config_entry)
+    ensure_via_device_exists(hass, specifications, config_entry)
+    device_info = device_info_from_specifications(specifications)
 
     if config_entry_id is not None and device_info is not None:
+        if via_device_id := _resolve_via_device_id(hass, specifications, config_entry):
+            device_info["via_device_id"] = via_device_id
         update_device_info = cast(dict[str, Any], device_info)
         update_device_info["config_entry_id"] = config_entry_id
         device = device_registry.async_get_or_create(**update_device_info)

tests/components/mqtt/test_device_trigger.py

@@ -1271,6 +1271,38 @@ async def test_entity_device_info_with_identifier(
     assert device.sw_version == "0.1-beta"
 
 
+async def test_entity_device_info_with_via_device(
+    hass: HomeAssistant,
+    device_registry: dr.DeviceRegistry,
+    mqtt_mock_entry: MqttMockHAClientGenerator,
+) -> None:
+    """Test device_trigger device registry integration links via_device_id."""
+    await mqtt_mock_entry()
+    mqtt_config_entry = hass.config_entries.async_entries(DOMAIN)[0]
+    hub = device_registry.async_get_or_create(
+        config_entry_id=mqtt_config_entry.entry_id,
+        identifiers={("mqtt", "hub-id")},
+        manufacturer="manufacturer",
+        model="hub",
+    )
+
+    data = json.dumps(
+        {
+            "automation_type": "trigger",
+            "topic": "test-topic",
+            "type": "foo",
+            "subtype": "bar",
+            "device": {"identifiers": ["helloworld"], "via_device": "hub-id"},
+        }
+    )
+    async_fire_mqtt_message(hass, "homeassistant/device_automation/bla/config", data)
+    await hass.async_block_till_done()
+
+    device = device_registry.async_get_device(identifiers={("mqtt", "helloworld")})
+    assert device is not None
+    assert device.via_device_id == hub.id
+
+
 async def test_entity_device_info_update(
     hass: HomeAssistant,
     device_registry: dr.DeviceRegistry,

tests/components/mqtt/test_discovery.py

@@ -51,7 +51,11 @@ from homeassistant.helpers.service_info.mqtt import MqttServiceInfo
 from homeassistant.setup import async_setup_component
 from homeassistant.util.signal_type import SignalTypeFormat
 
-from .common import help_all_subscribe_calls, help_test_unload_config_entry
+from .common import (
+    MOCK_NOTIFY_SUBENTRY_DATA_MULTI,
+    help_all_subscribe_calls,
+    help_test_unload_config_entry,
+)
 from .conftest import ENTRY_DEFAULT_BIRTH_MESSAGE
 from .test_tag import DEFAULT_TAG_ID, DEFAULT_TAG_SCAN
 
@@ -3246,6 +3250,12 @@ async def test_discovery_with_late_via_device_discovery(
 
     await hass.async_block_till_done()
 
+    # The child device links to the stub via device by via_device_id
+    stub_id = via_device_entry.id
+    child_device_entry = device_registry.async_get_device({("mqtt", "0AFFD2")})
+    assert child_device_entry is not None
+    assert child_device_entry.via_device_id == stub_id
+
     # Now discover the via device (a switch)
     via_device_config = {
         "name": None,
@@ -3267,6 +3277,12 @@ async def test_discovery_with_late_via_device_discovery(
     )
     assert via_device_entry is not None
     assert via_device_entry.name == "My Switch"
+    # The stub merges into the announced device, keeping its id, so the link
+    # from the child device survives
+    assert via_device_entry.id == stub_id
+    child_device_entry = device_registry.async_get_device({("mqtt", "0AFFD2")})
+    assert child_device_entry is not None
+    assert child_device_entry.via_device_id == stub_id
 
     await help_check_discovered_items(hass, device_registry, tag_mock)
 
@@ -3322,6 +3338,12 @@ async def test_discovery_with_late_via_device_update(
     await hass.async_block_till_done()
     await hass.async_block_till_done()
 
+    # The discovery update established the via_device_id link on the child device
+    stub_id = via_device_entry.id
+    child_device_entry = device_registry.async_get_device({("mqtt", "0AFFD2")})
+    assert child_device_entry is not None
+    assert child_device_entry.via_device_id == stub_id
+
     # Now discover the via device (a switch)
     via_device_config = {
         "name": None,
@@ -3343,10 +3365,117 @@ async def test_discovery_with_late_via_device_update(
     )
     assert via_device_entry is not None
     assert via_device_entry.name == "My Switch"
+    assert via_device_entry.id == stub_id
+    child_device_entry = device_registry.async_get_device({("mqtt", "0AFFD2")})
+    assert child_device_entry is not None
+    assert child_device_entry.via_device_id == stub_id
 
     await help_check_discovered_items(hass, device_registry, tag_mock)
 
 
+async def test_via_device_relinks_after_parent_removed(
+    hass: HomeAssistant,
+    device_registry: dr.DeviceRegistry,
+    mqtt_mock_entry: MqttMockHAClientGenerator,
+) -> None:
+    """Test a child re-links to its via device after the parent is removed."""
+    await mqtt_mock_entry()
+
+    parent_config = {
+        "name": "Parent",
+        "command_topic": "test-parent-topic",
+        "unique_id": "parent_unique",
+        "device": {"identifiers": ["parent-id"], "name": "Parent"},
+    }
+    async_fire_mqtt_message(
+        hass, "homeassistant/switch/parent/config", json.dumps(parent_config)
+    )
+    child_config = {
+        "name": "Child",
+        "state_topic": "test-child-topic",
+        "event_types": ["press"],
+        "unique_id": "child_unique",
+        "device": {"identifiers": ["child-id"], "via_device": "parent-id"},
+    }
+    async_fire_mqtt_message(
+        hass, "homeassistant/event/child/config", json.dumps(child_config)
+    )
+    await hass.async_block_till_done()
+
+    parent = device_registry.async_get_device({("mqtt", "parent-id")})
+    child = device_registry.async_get_device({("mqtt", "child-id")})
+    assert parent is not None
+    assert child is not None
+    assert child.via_device_id == parent.id
+
+    # Removing the parent clears the child's via_device_id
+    device_registry.async_remove_device(parent.id)
+    await hass.async_block_till_done()
+    child = device_registry.async_get_device({("mqtt", "child-id")})
+    assert child is not None
+    assert child.via_device_id is None
+
+    # A child discovery update re-creates the via device stub and re-links
+    child_config["name"] = "Child updated"
+    async_fire_mqtt_message(
+        hass, "homeassistant/event/child/config", json.dumps(child_config)
+    )
+    await hass.async_block_till_done()
+
+    parent_stub = device_registry.async_get_device({("mqtt", "parent-id")})
+    child = device_registry.async_get_device({("mqtt", "child-id")})
+    assert parent_stub is not None
+    assert child is not None
+    assert child.via_device_id == parent_stub.id
+
+
+@pytest.mark.parametrize(
+    "mqtt_config_subentries_data",
+    [
+        (
+            config_entries.ConfigSubentryData(
+                data=MOCK_NOTIFY_SUBENTRY_DATA_MULTI,
+                subentry_type="device",
+                title="Parent subentry device",
+            ),
+        )
+    ],
+)
+async def test_via_device_across_subentries(
+    hass: HomeAssistant,
+    device_registry: dr.DeviceRegistry,
+    mqtt_mock_entry: MqttMockHAClientGenerator,
+) -> None:
+    """Test via_device resolves across subentries within one config entry."""
+    await mqtt_mock_entry()
+
+    config_entry = hass.config_entries.async_entries(DOMAIN)[0]
+    subentry_id = next(iter(config_entry.subentries))
+    parent = device_registry.async_get_device({(DOMAIN, subentry_id)})
+    assert parent is not None
+    assert parent.config_subentry_id == subentry_id
+
+    child_config = {
+        "name": "Child",
+        "state_topic": "test-child-topic",
+        "event_types": ["press"],
+        "unique_id": "child_unique",
+        "device": {"identifiers": ["child-id"], "via_device": subentry_id},
+    }
+    async_fire_mqtt_message(
+        hass, "homeassistant/event/child/config", json.dumps(child_config)
+    )
+    await hass.async_block_till_done()
+
+    child = device_registry.async_get_device({(DOMAIN, "child-id")})
+    assert child is not None
+    # The parent lives in a subentry and the discovered child does not, yet the
+    # link resolves because lookups are scoped to the config entry.
+    assert child.config_subentry_id is None
+    assert child.config_entry_id == parent.config_entry_id
+    assert child.via_device_id == parent.id
+
+
 async def test_shared_options_in_sync_with_device_schema() -> None:
     """Test shared options in device discovery schema are in sync.
 

tests/components/mqtt/test_tag.py

@@ -506,6 +506,35 @@ async def test_entity_device_info_with_identifier(
     assert device.sw_version == "0.1-beta"
 
 
+async def test_entity_device_info_with_via_device(
+    hass: HomeAssistant,
+    device_registry: dr.DeviceRegistry,
+    mqtt_mock_entry: MqttMockHAClientGenerator,
+) -> None:
+    """Test tag device registry integration links via_device_id."""
+    await mqtt_mock_entry()
+    mqtt_config_entry = hass.config_entries.async_entries(DOMAIN)[0]
+    hub = device_registry.async_get_or_create(
+        config_entry_id=mqtt_config_entry.entry_id,
+        identifiers={("mqtt", "hub-id")},
+        manufacturer="manufacturer",
+        model="hub",
+    )
+
+    data = json.dumps(
+        {
+            "topic": "test-topic",
+            "device": {"identifiers": ["helloworld"], "via_device": "hub-id"},
+        }
+    )
+    async_fire_mqtt_message(hass, "homeassistant/tag/bla/config", data)
+    await hass.async_block_till_done()
+
+    device = device_registry.async_get_device(identifiers={("mqtt", "helloworld")})
+    assert device is not None
+    assert device.via_device_id == hub.id
+
+
 async def test_entity_device_info_update(
     hass: HomeAssistant,
     device_registry: dr.DeviceRegistry,