Fix via_device race in litejet (#177724)

homeassistant/components/litejet/__init__.py

@@ -8,8 +8,9 @@ from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_PORT, EVENT_HOMEASSISTANT_STOP
 from homeassistant.core import Event, HomeAssistant
 from homeassistant.exceptions import ConfigEntryNotReady
+from homeassistant.helpers import device_registry as dr
 
-from .const import PLATFORMS
+from .const import DOMAIN, PLATFORMS
 
 type LiteJetConfigEntry = ConfigEntry[pylitejet.LiteJet]
 
@@ -42,6 +43,14 @@ async def async_setup_entry(hass: HomeAssistant, entry: LiteJetConfigEntry) -> b
 
     entry.runtime_data = system
 
+    dr.async_get(hass).async_get_or_create(
+        config_entry_id=entry.entry_id,
+        identifiers={(DOMAIN, f"{entry.entry_id}_mcp")},
+        name="LiteJet",
+        manufacturer="Centralite",
+        model=system.model_name,
+    )
+
     await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
 
     return True

homeassistant/components/litejet/light.py

@@ -13,6 +13,7 @@ from homeassistant.components.light import (
 )
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import HomeAssistantError
+from homeassistant.helpers import device_registry as dr
 from homeassistant.helpers.device_registry import DeviceInfo
 from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
 
@@ -34,7 +35,7 @@ async def async_setup_entry(
     entities = []
     for index in system.loads():
         name = await system.get_load_name(index)
-        entities.append(LiteJetLight(config_entry, system, index, name))
+        entities.append(LiteJetLight(hass, config_entry, system, index, name))
 
     async_add_entities(entities, True)
 
@@ -50,7 +51,12 @@ class LiteJetLight(LightEntity):
     _attr_name = None
 
     def __init__(
-        self, config_entry: LiteJetConfigEntry, system: LiteJet, index: int, name: str
+        self,
+        hass: HomeAssistant,
+        config_entry: LiteJetConfigEntry,
+        system: LiteJet,
+        index: int,
+        name: str,
     ) -> None:
         """Initialize a LiteJet light."""
         self._config_entry = config_entry
@@ -63,7 +69,11 @@ class LiteJetLight(LightEntity):
         self._attr_device_info = DeviceInfo(
             identifiers={(DOMAIN, f"{config_entry.entry_id}_light_{index}")},
             name=name,
-            via_device=(DOMAIN, f"{config_entry.entry_id}_mcp"),
+            via_device_id=dr.async_get_device_id_by_identifier(
+                hass,
+                (DOMAIN, f"{config_entry.entry_id}_mcp"),
+                config_entry_id=config_entry.entry_id,
+            ),
         )
 
     @override

homeassistant/components/litejet/switch.py

@@ -7,6 +7,7 @@ from pylitejet import LiteJet, LiteJetError
 from homeassistant.components.switch import SwitchDeviceClass, SwitchEntity
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import HomeAssistantError
+from homeassistant.helpers import device_registry as dr
 from homeassistant.helpers.device_registry import DeviceInfo
 from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
 
@@ -28,7 +29,7 @@ async def async_setup_entry(
     entities = []
     for i in system.button_switches():
         name = await system.get_switch_name(i)
-        entities.append(LiteJetSwitch(config_entry.entry_id, system, i, name))
+        entities.append(LiteJetSwitch(hass, config_entry.entry_id, system, i, name))
 
     async_add_entities(entities, True)
 
@@ -41,7 +42,9 @@ class LiteJetSwitch(SwitchEntity):
     _attr_entity_registry_enabled_default = False
     _attr_device_class = SwitchDeviceClass.SWITCH
 
-    def __init__(self, entry_id: str, system: LiteJet, i: int, name: str) -> None:
+    def __init__(
+        self, hass: HomeAssistant, entry_id: str, system: LiteJet, i: int, name: str
+    ) -> None:
         """Initialize a LiteJet switch."""
         self._lj = system
         self._index = i
@@ -55,7 +58,9 @@ class LiteJetSwitch(SwitchEntity):
             identifiers={(DOMAIN, f"{entry_id}_keypad_{keypad_number}")},
             name=system.get_switch_keypad_name(i),
             manufacturer="Centralite",
-            via_device=(DOMAIN, f"{entry_id}_mcp"),
+            via_device_id=dr.async_get_device_id_by_identifier(
+                hass, (DOMAIN, f"{entry_id}_mcp"), config_entry_id=entry_id
+            ),
         )
 
     @override

tests/components/litejet/test_init.py

@@ -1,7 +1,10 @@
 """The tests for the litejet component."""
 
+import pytest
+
 from homeassistant.components.litejet.const import DOMAIN
 from homeassistant.core import HomeAssistant
+from homeassistant.helpers import device_registry as dr
 from homeassistant.setup import async_setup_component
 
 from . import async_init_integration
@@ -13,6 +16,31 @@ async def test_setup_with_no_config(hass: HomeAssistant) -> None:
     assert DOMAIN not in hass.data
 
 
+@pytest.mark.usefixtures("mock_litejet")
+async def test_child_devices_link_to_mcp(
+    hass: HomeAssistant, device_registry: dr.DeviceRegistry
+) -> None:
+    """Test that light and switch devices are linked to the MCP parent device."""
+    entry = await async_init_integration(hass, use_switch=True)
+
+    parent = device_registry.async_get_device_by_identifier(
+        (DOMAIN, f"{entry.entry_id}_mcp"), entry.entry_id
+    )
+    assert parent is not None
+
+    light_device = device_registry.async_get_device_by_identifier(
+        (DOMAIN, f"{entry.entry_id}_light_1"), entry.entry_id
+    )
+    assert light_device is not None
+    assert light_device.via_device_id == parent.id
+
+    switch_device = device_registry.async_get_device_by_identifier(
+        (DOMAIN, f"{entry.entry_id}_keypad_101"), entry.entry_id
+    )
+    assert switch_device is not None
+    assert switch_device.via_device_id == parent.id
+
+
 async def test_unload_entry(hass: HomeAssistant, mock_litejet) -> None:
     """Test being able to unload an entry."""
     entry = await async_init_integration(hass, use_switch=True, use_scene=True)