Validate device id when adding or updating entity registry entry (#134982)

homeassistant/helpers/entity_registry.py

@@ -648,6 +648,7 @@ def _validate_item(
     domain: str,
     platform: str,
     *,
+    device_id: str | None | UndefinedType = None,
     disabled_by: RegistryEntryDisabler | None | UndefinedType = None,
     entity_category: EntityCategory | None | UndefinedType = None,
     hidden_by: RegistryEntryHider | None | UndefinedType = None,
@@ -671,6 +672,10 @@ def _validate_item(
             unique_id,
             report_issue,
         )
+    if device_id and device_id is not UNDEFINED:
+        device_registry = dr.async_get(hass)
+        if not device_registry.async_get(device_id):
+            raise ValueError(f"Device {device_id} does not exist")
     if (
         disabled_by
         and disabled_by is not UNDEFINED
@@ -859,6 +864,7 @@ class EntityRegistry(BaseRegistry):
             self.hass,
             domain,
             platform,
+            device_id=device_id,
             disabled_by=disabled_by,
             entity_category=entity_category,
             hidden_by=hidden_by,
@@ -1090,6 +1096,7 @@ class EntityRegistry(BaseRegistry):
                 self.hass,
                 old.domain,
                 old.platform,
+                device_id=device_id,
                 disabled_by=disabled_by,
                 entity_category=entity_category,
                 hidden_by=hidden_by,

tests/components/cloud/test_google_config.py

@@ -32,7 +32,7 @@ from homeassistant.helpers import device_registry as dr, entity_registry as er
 from homeassistant.setup import async_setup_component
 from homeassistant.util.dt import utcnow
 
-from tests.common import async_fire_time_changed
+from tests.common import MockConfigEntry, async_fire_time_changed
 
 
 @pytest.fixture
@@ -264,6 +264,7 @@ async def test_google_entity_registry_sync(
 @pytest.mark.usefixtures("mock_cloud_login")
 async def test_google_device_registry_sync(
     hass: HomeAssistant,
+    device_registry: dr.DeviceRegistry,
     entity_registry: er.EntityRegistry,
     cloud_prefs: CloudPreferences,
 ) -> None:
@@ -275,8 +276,14 @@ async def test_google_device_registry_sync(
     # Enable exposing new entities to Google
     expose_new(hass, True)
 
+    config_entry = MockConfigEntry(domain="test", data={})
+    config_entry.add_to_hass(hass)
+    device_entry = device_registry.async_get_or_create(
+        config_entry_id=config_entry.entry_id,
+        connections={(dr.CONNECTION_NETWORK_MAC, "12:34:56:AB:CD:EF")},
+    )
     entity_entry = entity_registry.async_get_or_create(
-        "light", "hue", "1234", device_id="1234"
+        "light", "hue", "1234", device_id=device_entry.id
     )
     entity_entry = entity_registry.async_update_entity(
         entity_entry.entity_id, area_id="ABCD"
@@ -294,7 +301,7 @@ async def test_google_device_registry_sync(
             dr.EVENT_DEVICE_REGISTRY_UPDATED,
             {
                 "action": "update",
-                "device_id": "1234",
+                "device_id": device_entry.id,
                 "changes": ["manufacturer"],
             },
         )
@@ -308,7 +315,7 @@ async def test_google_device_registry_sync(
             dr.EVENT_DEVICE_REGISTRY_UPDATED,
             {
                 "action": "update",
-                "device_id": "1234",
+                "device_id": device_entry.id,
                 "changes": ["area_id"],
             },
         )
@@ -324,7 +331,7 @@ async def test_google_device_registry_sync(
             dr.EVENT_DEVICE_REGISTRY_UPDATED,
             {
                 "action": "update",
-                "device_id": "1234",
+                "device_id": device_entry.id,
                 "changes": ["area_id"],
             },
         )

tests/components/homekit/test_config_flow.py

@@ -1406,7 +1406,6 @@ async def test_options_flow_exclude_mode_skips_category_entities(
         "switch",
         "sonos",
         "config",
-        device_id="1234",
         entity_category=EntityCategory.CONFIG,
     )
     hass.states.async_set(sonos_config_switch.entity_id, "off")
@@ -1415,7 +1414,6 @@ async def test_options_flow_exclude_mode_skips_category_entities(
         "switch",
         "sonos",
         "notconfig",
-        device_id="1234",
         entity_category=None,
     )
     hass.states.async_set(sonos_notconfig_switch.entity_id, "off")
@@ -1510,7 +1508,6 @@ async def test_options_flow_exclude_mode_skips_hidden_entities(
         "switch",
         "sonos",
         "config",
-        device_id="1234",
         hidden_by=er.RegistryEntryHider.INTEGRATION,
     )
     hass.states.async_set(sonos_hidden_switch.entity_id, "off")
@@ -1594,7 +1591,6 @@ async def test_options_flow_include_mode_allows_hidden_entities(
         "switch",
         "sonos",
         "config",
-        device_id="1234",
         hidden_by=er.RegistryEntryHider.INTEGRATION,
     )
     hass.states.async_set(sonos_hidden_switch.entity_id, "off")

tests/components/homekit/test_homekit.py

@@ -615,7 +615,6 @@ async def test_homekit_entity_glob_filter_with_config_entities(
         "select",
         "any",
         "any",
-        device_id="1234",
         entity_category=EntityCategory.CONFIG,
     )
     hass.states.async_set(select_config_entity.entity_id, "off")
@@ -624,7 +623,6 @@ async def test_homekit_entity_glob_filter_with_config_entities(
         "switch",
         "any",
         "any",
-        device_id="1234",
         entity_category=EntityCategory.CONFIG,
     )
     hass.states.async_set(switch_config_entity.entity_id, "off")
@@ -669,7 +667,6 @@ async def test_homekit_entity_glob_filter_with_hidden_entities(
         "select",
         "any",
         "any",
-        device_id="1234",
         hidden_by=er.RegistryEntryHider.INTEGRATION,
     )
     hass.states.async_set(select_config_entity.entity_id, "off")
@@ -678,7 +675,6 @@ async def test_homekit_entity_glob_filter_with_hidden_entities(
         "switch",
         "any",
         "any",
-        device_id="1234",
         hidden_by=er.RegistryEntryHider.INTEGRATION,
     )
     hass.states.async_set(switch_config_entity.entity_id, "off")
@@ -1867,7 +1863,11 @@ async def test_homekit_ignored_missing_devices(
     device_registry: dr.DeviceRegistry,
     entity_registry: er.EntityRegistry,
 ) -> None:
-    """Test HomeKit handles a device in the entity registry but missing from the device registry."""
+    """Test HomeKit handles a device in the entity registry but missing from the device registry.
+
+    If the entity registry is updated to remove entities linked to non-existent devices,
+    or set the link to None, this test can be removed.
+    """
 
     entry = await async_init_integration(hass)
     homekit = _mock_homekit(hass, entry, HOMEKIT_MODE_BRIDGE)
@@ -1885,47 +1885,37 @@ async def test_homekit_ignored_missing_devices(
         connections={(dr.CONNECTION_NETWORK_MAC, "12:34:56:AB:CD:EF")},
     )
 
-    entity_registry.async_get_or_create(
+    binary_sensor_entity = entity_registry.async_get_or_create(
         "binary_sensor",
         "powerwall",
         "battery_charging",
         device_id=device_entry.id,
         original_device_class=BinarySensorDeviceClass.BATTERY_CHARGING,
     )
-    entity_registry.async_get_or_create(
+    sensor_entity = entity_registry.async_get_or_create(
         "sensor",
         "powerwall",
         "battery",
         device_id=device_entry.id,
         original_device_class=SensorDeviceClass.BATTERY,
     )
-    light = entity_registry.async_get_or_create(
+    light_entity = light = entity_registry.async_get_or_create(
         "light", "powerwall", "demo", device_id=device_entry.id
     )
     # Delete the device to make sure we fallback
     # to using the platform
-    device_registry.async_remove_device(device_entry.id)
+    with patch(
-    # Wait for the entities to be removed
+        "homeassistant.helpers.entity_registry.async_entries_for_device",
-    await asyncio.sleep(0)
+        return_value=[],
-    await asyncio.sleep(0)
+    ):
-    # Restore the registry
+        device_registry.async_remove_device(device_entry.id)
-    entity_registry.async_get_or_create(
+        # Wait for the device registry event handlers to execute
-        "binary_sensor",
+        await asyncio.sleep(0)
-        "powerwall",
+        await asyncio.sleep(0)
-        "battery_charging",
+    # Check the entities were not removed
-        device_id=device_entry.id,
+    assert binary_sensor_entity.entity_id in entity_registry.entities
-        original_device_class=BinarySensorDeviceClass.BATTERY_CHARGING,
+    assert sensor_entity.entity_id in entity_registry.entities
-    )
+    assert light_entity.entity_id in entity_registry.entities
-    entity_registry.async_get_or_create(
-        "sensor",
-        "powerwall",
-        "battery",
-        device_id=device_entry.id,
-        original_device_class=SensorDeviceClass.BATTERY,
-    )
-    light = entity_registry.async_get_or_create(
-        "light", "powerwall", "demo", device_id=device_entry.id
-    )
 
     hass.states.async_set(light.entity_id, STATE_ON)
     hass.states.async_set("light.two", STATE_ON)

tests/components/number/test_device_action.py

@@ -243,18 +243,26 @@ async def test_action_legacy(
 
 
 async def test_capabilities(
-    hass: HomeAssistant, entity_registry: er.EntityRegistry
+    hass: HomeAssistant,
+    device_registry: dr.DeviceRegistry,
+    entity_registry: er.EntityRegistry,
 ) -> None:
     """Test getting capabilities."""
-    entry = entity_registry.async_get_or_create(
+    config_entry = MockConfigEntry(domain="test", data={})
-        DOMAIN, "test", "5678", device_id="abcdefgh"
+    config_entry.add_to_hass(hass)
+    device_entry = device_registry.async_get_or_create(
+        config_entry_id=config_entry.entry_id,
+        connections={(dr.CONNECTION_NETWORK_MAC, "12:34:56:AB:CD:EF")},
+    )
+    entity_entry = entity_registry.async_get_or_create(
+        DOMAIN, "test", "5678", device_id=device_entry.id
     )
     capabilities = await device_action.async_get_action_capabilities(
         hass,
         {
             "domain": DOMAIN,
             "device_id": "abcdefgh",
-            "entity_id": entry.id,
+            "entity_id": entity_entry.id,
             "type": "set_value",
         },
     )
@@ -267,18 +275,26 @@ async def test_capabilities(
 
 
 async def test_capabilities_legacy(
-    hass: HomeAssistant, entity_registry: er.EntityRegistry
+    hass: HomeAssistant,
+    device_registry: dr.DeviceRegistry,
+    entity_registry: er.EntityRegistry,
 ) -> None:
     """Test getting capabilities."""
-    entry = entity_registry.async_get_or_create(
+    config_entry = MockConfigEntry(domain="test", data={})
-        DOMAIN, "test", "5678", device_id="abcdefgh"
+    config_entry.add_to_hass(hass)
+    device_entry = device_registry.async_get_or_create(
+        config_entry_id=config_entry.entry_id,
+        connections={(dr.CONNECTION_NETWORK_MAC, "12:34:56:AB:CD:EF")},
+    )
+    entity_entry = entity_registry.async_get_or_create(
+        DOMAIN, "test", "5678", device_id=device_entry.id
     )
     capabilities = await device_action.async_get_action_capabilities(
         hass,
         {
             "domain": DOMAIN,
             "device_id": "abcdefgh",
-            "entity_id": entry.entity_id,
+            "entity_id": entity_entry.entity_id,
             "type": "set_value",
         },
     )

tests/components/ps4/test_init.py

@@ -29,7 +29,7 @@ from homeassistant.const import (
 from homeassistant.core import HomeAssistant
 from homeassistant.data_entry_flow import FlowResultType
 from homeassistant.exceptions import HomeAssistantError
-from homeassistant.helpers import entity_registry as er
+from homeassistant.helpers import device_registry as dr, entity_registry as er
 from homeassistant.setup import async_setup_component
 from homeassistant.util import location
 
@@ -80,8 +80,6 @@ MOCK_DEVICE_VERSION_1 = {
 
 MOCK_DATA_VERSION_1 = {CONF_TOKEN: MOCK_CREDS, "devices": [MOCK_DEVICE_VERSION_1]}
 
-MOCK_DEVICE_ID = "somedeviceid"
-
 MOCK_ENTRY_VERSION_1 = MockConfigEntry(
     domain=DOMAIN, data=MOCK_DATA_VERSION_1, entry_id=MOCK_ENTRY_ID, version=1
 )
@@ -141,20 +139,26 @@ async def test_creating_entry_sets_up_media_player(hass: HomeAssistant) -> None:
 
 
 async def test_config_flow_entry_migrate(
-    hass: HomeAssistant, entity_registry: er.EntityRegistry
+    hass: HomeAssistant,
+    device_registry: dr.DeviceRegistry,
+    entity_registry: er.EntityRegistry,
 ) -> None:
     """Test that config flow entry is migrated correctly."""
     # Start with the config entry at Version 1.
     manager = hass.config_entries
     mock_entry = MOCK_ENTRY_VERSION_1
     mock_entry.add_to_manager(manager)
+    mock_device_entry = device_registry.async_get_or_create(
+        config_entry_id=mock_entry.entry_id,
+        connections={(dr.CONNECTION_NETWORK_MAC, "12:34:56:AB:CD:EF")},
+    )
     mock_entity_id = f"media_player.ps4_{MOCK_UNIQUE_ID}"
     mock_e_entry = entity_registry.async_get_or_create(
         "media_player",
         "ps4",
         MOCK_UNIQUE_ID,
         config_entry=mock_entry,
-        device_id=MOCK_DEVICE_ID,
+        device_id=mock_device_entry.id,
     )
     assert len(entity_registry.entities) == 1
     assert mock_e_entry.entity_id == mock_entity_id
@@ -180,7 +184,7 @@ async def test_config_flow_entry_migrate(
 
     # Test that entity_id remains the same.
     assert mock_entity.entity_id == mock_entity_id
-    assert mock_entity.device_id == MOCK_DEVICE_ID
+    assert mock_entity.device_id == mock_device_entry.id
 
     # Test that last four of credentials is appended to the unique_id.
     assert mock_entity.unique_id == f"{MOCK_UNIQUE_ID}_{MOCK_CREDS[-4:]}"

tests/helpers/test_device_registry.py

@@ -2276,7 +2276,9 @@ async def test_cleanup_startup(hass: HomeAssistant) -> None:
 
 
 @pytest.mark.parametrize("load_registries", [False])
-async def test_cleanup_entity_registry_change(hass: HomeAssistant) -> None:
+async def test_cleanup_entity_registry_change(
+    hass: HomeAssistant, mock_config_entry: MockConfigEntry
+) -> None:
     """Test we run a cleanup when entity registry changes.
 
     Don't pre-load the registries as the debouncer will then not be waiting for
@@ -2284,8 +2286,14 @@ async def test_cleanup_entity_registry_change(hass: HomeAssistant) -> None:
     """
     await dr.async_load(hass)
     await er.async_load(hass)
+    dev_reg = dr.async_get(hass)
     ent_reg = er.async_get(hass)
 
+    entry = dev_reg.async_get_or_create(
+        config_entry_id=mock_config_entry.entry_id,
+        connections={(dr.CONNECTION_NETWORK_MAC, "12:34:56:AB:CD:EF")},
+    )
+
     with patch(
         "homeassistant.helpers.device_registry.Debouncer.async_schedule_call"
     ) as mock_call:
@@ -2299,7 +2307,7 @@ async def test_cleanup_entity_registry_change(hass: HomeAssistant) -> None:
         assert len(mock_call.mock_calls) == 0
 
         # Device ID update triggers
-        ent_reg.async_get_or_create("light", "hue", "e1", device_id="bla")
+        ent_reg.async_get_or_create("light", "hue", "e1", device_id=entry.id)
         await hass.async_block_till_done()
         assert len(mock_call.mock_calls) == 1
 

tests/helpers/test_entity_registry.py

@@ -72,11 +72,18 @@ def test_get_or_create_suggested_object_id(entity_registry: er.EntityRegistry) -
 
 
 def test_get_or_create_updates_data(
+    hass: HomeAssistant,
+    device_registry: dr.DeviceRegistry,
     entity_registry: er.EntityRegistry,
     freezer: FrozenDateTimeFactory,
 ) -> None:
     """Test that we update data in get_or_create."""
     orig_config_entry = MockConfigEntry(domain="light")
+    orig_config_entry.add_to_hass(hass)
+    orig_device_entry = device_registry.async_get_or_create(
+        config_entry_id=orig_config_entry.entry_id,
+        connections={(dr.CONNECTION_NETWORK_MAC, "12:34:56:AB:CD:EF")},
+    )
     created = datetime.fromisoformat("2024-02-14T12:00:00.0+00:00")
     freezer.move_to(created)
 
@@ -86,7 +93,7 @@ def test_get_or_create_updates_data(
         "5678",
         capabilities={"max": 100},
         config_entry=orig_config_entry,
-        device_id="mock-dev-id",
+        device_id=orig_device_entry.id,
         disabled_by=er.RegistryEntryDisabler.HASS,
         entity_category=EntityCategory.CONFIG,
         has_entity_name=True,
@@ -99,7 +106,7 @@ def test_get_or_create_updates_data(
         unit_of_measurement="initial-unit_of_measurement",
     )
 
-    assert set(entity_registry.async_device_ids()) == {"mock-dev-id"}
+    assert set(entity_registry.async_device_ids()) == {orig_device_entry.id}
 
     assert orig_entry == er.RegistryEntry(
         "light.hue_5678",
@@ -109,7 +116,7 @@ def test_get_or_create_updates_data(
         config_entry_id=orig_config_entry.entry_id,
         created_at=created,
         device_class=None,
-        device_id="mock-dev-id",
+        device_id=orig_device_entry.id,
         disabled_by=er.RegistryEntryDisabler.HASS,
         entity_category=EntityCategory.CONFIG,
         has_entity_name=True,
@@ -127,6 +134,11 @@ def test_get_or_create_updates_data(
     )
 
     new_config_entry = MockConfigEntry(domain="light")
+    new_config_entry.add_to_hass(hass)
+    new_device_entry = device_registry.async_get_or_create(
+        config_entry_id=new_config_entry.entry_id,
+        connections={(dr.CONNECTION_NETWORK_MAC, "34:56:AB:CD:EF:12")},
+    )
     modified = created + timedelta(minutes=5)
     freezer.move_to(modified)
 
@@ -136,7 +148,7 @@ def test_get_or_create_updates_data(
         "5678",
         capabilities={"new-max": 150},
         config_entry=new_config_entry,
-        device_id="new-mock-dev-id",
+        device_id=new_device_entry.id,
         disabled_by=er.RegistryEntryDisabler.USER,
         entity_category=EntityCategory.DIAGNOSTIC,
         has_entity_name=False,
@@ -159,7 +171,7 @@ def test_get_or_create_updates_data(
         config_entry_id=new_config_entry.entry_id,
         created_at=created,
         device_class=None,
-        device_id="new-mock-dev-id",
+        device_id=new_device_entry.id,
         disabled_by=er.RegistryEntryDisabler.HASS,  # Should not be updated
         entity_category=EntityCategory.DIAGNOSTIC,
         has_entity_name=False,
@@ -176,7 +188,7 @@ def test_get_or_create_updates_data(
         unit_of_measurement="updated-unit_of_measurement",
     )
 
-    assert set(entity_registry.async_device_ids()) == {"new-mock-dev-id"}
+    assert set(entity_registry.async_device_ids()) == {new_device_entry.id}
     modified = created + timedelta(minutes=5)
     freezer.move_to(modified)
 
@@ -262,10 +274,18 @@ def test_create_triggers_save(entity_registry: er.EntityRegistry) -> None:
 
 
 async def test_loading_saving_data(
-    hass: HomeAssistant, entity_registry: er.EntityRegistry
+    hass: HomeAssistant,
+    device_registry: dr.DeviceRegistry,
+    entity_registry: er.EntityRegistry,
 ) -> None:
     """Test that we load/save data correctly."""
     mock_config = MockConfigEntry(domain="light")
+    mock_config.add_to_hass(hass)
+
+    device_entry = device_registry.async_get_or_create(
+        config_entry_id=mock_config.entry_id,
+        connections={(dr.CONNECTION_NETWORK_MAC, "12:34:56:AB:CD:EF")},
+    )
 
     orig_entry1 = entity_registry.async_get_or_create("light", "hue", "1234")
     orig_entry2 = entity_registry.async_get_or_create(
@@ -274,7 +294,7 @@ async def test_loading_saving_data(
         "5678",
         capabilities={"max": 100},
         config_entry=mock_config,
-        device_id="mock-dev-id",
+        device_id=device_entry.id,
         disabled_by=er.RegistryEntryDisabler.HASS,
         entity_category=EntityCategory.CONFIG,
         hidden_by=er.RegistryEntryHider.INTEGRATION,
@@ -338,7 +358,7 @@ async def test_loading_saving_data(
     assert new_entry2.capabilities == {"max": 100}
     assert new_entry2.config_entry_id == mock_config.entry_id
     assert new_entry2.device_class == "user-class"
-    assert new_entry2.device_id == "mock-dev-id"
+    assert new_entry2.device_id == device_entry.id
     assert new_entry2.disabled_by is er.RegistryEntryDisabler.HASS
     assert new_entry2.entity_category == "config"
     assert new_entry2.icon == "hass:user-icon"
@@ -1741,6 +1761,16 @@ def test_entity_registry_items() -> None:
     assert entities.get_entry(entry2.id) is None
 
 
+async def test_device_does_not_exist(entity_registry: er.EntityRegistry) -> None:
+    """Test adding an entity linked to an unknown device."""
+    with pytest.raises(ValueError):
+        entity_registry.async_get_or_create("light", "hue", "1234", device_id="blah")
+
+    entity_id = entity_registry.async_get_or_create("light", "hue", "1234").entity_id
+    with pytest.raises(ValueError):
+        entity_registry.async_update_entity(entity_id, device_id="blah")
+
+
 async def test_disabled_by_str_not_allowed(entity_registry: er.EntityRegistry) -> None:
     """Test we need to pass disabled by type."""
     with pytest.raises(ValueError):