Create MQTT device referenced by via device (#131588)

homeassistant/components/mqtt/entity.py

@@ -1185,6 +1185,33 @@ def device_info_from_specifications(
     return info
 
 
+@callback
+def ensure_via_device_exists(
+    hass: HomeAssistant, device_info: DeviceInfo | 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(
+            identifiers={device_info["via_device"]}
+        )
+    ):
+        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_registry.async_get_or_create(
+        config_entry_id=config_entry.entry_id,
+        identifiers={device_info["via_device"]},
+    )
+
+
 class MqttEntityDeviceInfo(Entity):
     """Mixin used for mqtt platforms that support the device registry."""
 
@@ -1203,6 +1230,7 @@ class MqttEntityDeviceInfo(Entity):
         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
             )
@@ -1256,6 +1284,7 @@ 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)
 
     def _init_entity_id(self) -> None:
         """Set entity_id from object_id if defined in config."""
@@ -1490,6 +1519,8 @@ def update_device(
     config_entry_id = config_entry.entry_id
     device_info = device_info_from_specifications(config[CONF_DEVICE])
 
+    ensure_via_device_exists(hass, device_info, config_entry)
+
     if config_entry_id is not None and device_info is not None:
         update_device_info = cast(dict[str, Any], device_info)
         update_device_info["config_entry_id"] = config_entry_id

tests/components/mqtt/test_discovery.py

@@ -2987,3 +2987,139 @@ async def test_shared_state_topic(
     state = hass.states.get(entity_id)
     assert state is not None
     assert state.state == "New state3"
+
+
+@pytest.mark.parametrize("single_configs", [copy.deepcopy(TEST_SINGLE_CONFIGS)])
+async def test_discovery_with_late_via_device_discovery(
+    hass: HomeAssistant,
+    device_registry: dr.DeviceRegistry,
+    mqtt_mock_entry: MqttMockHAClientGenerator,
+    tag_mock: AsyncMock,
+    single_configs: list[tuple[str, dict[str, Any]]],
+) -> None:
+    """Test a via device is available and the discovery of the via device is late."""
+    await mqtt_mock_entry()
+
+    await hass.async_block_till_done()
+    await hass.async_block_till_done()
+
+    via_device_entry = device_registry.async_get_device(
+        {("mqtt", "id_via_very_unique")}
+    )
+    assert via_device_entry is None
+    # Discovery single config schema
+    for discovery_topic, config in single_configs:
+        config["device"]["via_device"] = "id_via_very_unique"
+        payload = json.dumps(config)
+        async_fire_mqtt_message(
+            hass,
+            discovery_topic,
+            payload,
+        )
+        via_device_entry = device_registry.async_get_device(
+            {("mqtt", "id_via_very_unique")}
+        )
+        assert via_device_entry is not None
+        assert via_device_entry.name is None
+
+    await hass.async_block_till_done()
+
+    # Now discover the via device (a switch)
+    via_device_config = {
+        "name": None,
+        "command_topic": "test-switch-topic",
+        "unique_id": "very_unique_switch",
+        "device": {"identifiers": ["id_via_very_unique"], "name": "My Switch"},
+    }
+    payload = json.dumps(via_device_config)
+    via_device_discovery_topic = "homeassistant/switch/very_unique/config"
+    async_fire_mqtt_message(
+        hass,
+        via_device_discovery_topic,
+        payload,
+    )
+    await hass.async_block_till_done()
+    await hass.async_block_till_done()
+    via_device_entry = device_registry.async_get_device(
+        {("mqtt", "id_via_very_unique")}
+    )
+    assert via_device_entry is not None
+    assert via_device_entry.name == "My Switch"
+
+    await help_check_discovered_items(hass, device_registry, tag_mock)
+
+
+@pytest.mark.parametrize("single_configs", [copy.deepcopy(TEST_SINGLE_CONFIGS)])
+async def test_discovery_with_late_via_device_update(
+    hass: HomeAssistant,
+    device_registry: dr.DeviceRegistry,
+    mqtt_mock_entry: MqttMockHAClientGenerator,
+    tag_mock: AsyncMock,
+    single_configs: list[tuple[str, dict[str, Any]]],
+) -> None:
+    """Test a via device is available and the discovery of the via device is is set via an update."""
+    await mqtt_mock_entry()
+
+    await hass.async_block_till_done()
+    await hass.async_block_till_done()
+
+    via_device_entry = device_registry.async_get_device(
+        {("mqtt", "id_via_very_unique")}
+    )
+    assert via_device_entry is None
+    # Discovery single config schema without via device
+    for discovery_topic, config in single_configs:
+        payload = json.dumps(config)
+        async_fire_mqtt_message(
+            hass,
+            discovery_topic,
+            payload,
+        )
+        via_device_entry = device_registry.async_get_device(
+            {("mqtt", "id_via_very_unique")}
+        )
+        await hass.async_block_till_done()
+        await hass.async_block_till_done()
+        assert via_device_entry is None
+
+    # Resend the discovery update to set the via device
+    for discovery_topic, config in single_configs:
+        config["device"]["via_device"] = "id_via_very_unique"
+        payload = json.dumps(config)
+        async_fire_mqtt_message(
+            hass,
+            discovery_topic,
+            payload,
+        )
+        via_device_entry = device_registry.async_get_device(
+            {("mqtt", "id_via_very_unique")}
+        )
+        assert via_device_entry is not None
+        assert via_device_entry.name is None
+
+    await hass.async_block_till_done()
+    await hass.async_block_till_done()
+
+    # Now discover the via device (a switch)
+    via_device_config = {
+        "name": None,
+        "command_topic": "test-switch-topic",
+        "unique_id": "very_unique_switch",
+        "device": {"identifiers": ["id_via_very_unique"], "name": "My Switch"},
+    }
+    payload = json.dumps(via_device_config)
+    via_device_discovery_topic = "homeassistant/switch/very_unique/config"
+    async_fire_mqtt_message(
+        hass,
+        via_device_discovery_topic,
+        payload,
+    )
+    await hass.async_block_till_done()
+    await hass.async_block_till_done()
+    via_device_entry = device_registry.async_get_device(
+        {("mqtt", "id_via_very_unique")}
+    )
+    assert via_device_entry is not None
+    assert via_device_entry.name == "My Switch"
+
+    await help_check_discovered_items(hass, device_registry, tag_mock)