"""Tests for the Entity Registry.""" from unittest.mock import patch import pytest import voluptuous as vol from homeassistant import config_entries from homeassistant.const import ( EVENT_HOMEASSISTANT_START, STATE_UNAVAILABLE, EntityCategory, ) from homeassistant.core import CoreState, HomeAssistant, callback from homeassistant.exceptions import MaxLengthExceeded from homeassistant.helpers import device_registry as dr, entity_registry as er from tests.common import MockConfigEntry, flush_store YAML__OPEN_PATH = "homeassistant.util.yaml.loader.open" @pytest.fixture def update_events(hass): """Capture update events.""" events = [] @callback def async_capture(event): events.append(event.data) hass.bus.async_listen(er.EVENT_ENTITY_REGISTRY_UPDATED, async_capture) return events async def test_get_or_create_returns_same_entry(hass, entity_registry, update_events): """Make sure we do not duplicate entries.""" entry = entity_registry.async_get_or_create("light", "hue", "1234") entry2 = entity_registry.async_get_or_create("light", "hue", "1234") await hass.async_block_till_done() assert len(entity_registry.entities) == 1 assert entry is entry2 assert entry.entity_id == "light.hue_1234" assert len(update_events) == 1 assert update_events[0]["action"] == "create" assert update_events[0]["entity_id"] == entry.entity_id def test_get_or_create_suggested_object_id(entity_registry): """Test that suggested_object_id works.""" entry = entity_registry.async_get_or_create( "light", "hue", "1234", suggested_object_id="beer" ) assert entry.entity_id == "light.beer" def test_get_or_create_updates_data(entity_registry): """Test that we update data in get_or_create.""" orig_config_entry = MockConfigEntry(domain="light") orig_entry = entity_registry.async_get_or_create( "light", "hue", "5678", capabilities={"max": 100}, config_entry=orig_config_entry, device_id="mock-dev-id", disabled_by=er.RegistryEntryDisabler.HASS, entity_category=EntityCategory.CONFIG, has_entity_name=True, hidden_by=er.RegistryEntryHider.INTEGRATION, original_device_class="mock-device-class", original_icon="initial-original_icon", original_name="initial-original_name", supported_features=5, translation_key="initial-translation_key", unit_of_measurement="initial-unit_of_measurement", ) assert orig_entry == er.RegistryEntry( "light.hue_5678", "5678", "hue", capabilities={"max": 100}, config_entry_id=orig_config_entry.entry_id, device_class=None, device_id="mock-dev-id", disabled_by=er.RegistryEntryDisabler.HASS, entity_category=EntityCategory.CONFIG, has_entity_name=True, hidden_by=er.RegistryEntryHider.INTEGRATION, icon=None, id=orig_entry.id, name=None, original_device_class="mock-device-class", original_icon="initial-original_icon", original_name="initial-original_name", supported_features=5, translation_key="initial-translation_key", unit_of_measurement="initial-unit_of_measurement", ) new_config_entry = MockConfigEntry(domain="light") new_entry = entity_registry.async_get_or_create( "light", "hue", "5678", capabilities={"new-max": 150}, config_entry=new_config_entry, device_id="new-mock-dev-id", disabled_by=er.RegistryEntryDisabler.USER, entity_category=EntityCategory.DIAGNOSTIC, has_entity_name=False, hidden_by=er.RegistryEntryHider.USER, original_device_class="new-mock-device-class", original_icon="updated-original_icon", original_name="updated-original_name", supported_features=10, translation_key="updated-translation_key", unit_of_measurement="updated-unit_of_measurement", ) assert new_entry == er.RegistryEntry( "light.hue_5678", "5678", "hue", aliases=set(), area_id=None, capabilities={"new-max": 150}, config_entry_id=new_config_entry.entry_id, device_class=None, device_id="new-mock-dev-id", disabled_by=er.RegistryEntryDisabler.HASS, # Should not be updated entity_category=EntityCategory.DIAGNOSTIC, has_entity_name=False, hidden_by=er.RegistryEntryHider.INTEGRATION, # Should not be updated icon=None, id=orig_entry.id, name=None, original_device_class="new-mock-device-class", original_icon="updated-original_icon", original_name="updated-original_name", supported_features=10, translation_key="updated-translation_key", unit_of_measurement="updated-unit_of_measurement", ) new_entry = entity_registry.async_get_or_create( "light", "hue", "5678", capabilities=None, config_entry=None, device_id=None, disabled_by=None, entity_category=None, has_entity_name=None, hidden_by=None, original_device_class=None, original_icon=None, original_name=None, supported_features=None, translation_key=None, unit_of_measurement=None, ) assert new_entry == er.RegistryEntry( "light.hue_5678", "5678", "hue", aliases=set(), area_id=None, capabilities=None, config_entry_id=None, device_class=None, device_id=None, disabled_by=er.RegistryEntryDisabler.HASS, # Should not be updated entity_category=None, has_entity_name=None, hidden_by=er.RegistryEntryHider.INTEGRATION, # Should not be updated icon=None, id=orig_entry.id, name=None, original_device_class=None, original_icon=None, original_name=None, supported_features=0, # supported_features is stored as an int translation_key=None, unit_of_measurement=None, ) def test_get_or_create_suggested_object_id_conflict_register(entity_registry): """Test that we don't generate an entity id that is already registered.""" entry = entity_registry.async_get_or_create( "light", "hue", "1234", suggested_object_id="beer" ) entry2 = entity_registry.async_get_or_create( "light", "hue", "5678", suggested_object_id="beer" ) assert entry.entity_id == "light.beer" assert entry2.entity_id == "light.beer_2" def test_get_or_create_suggested_object_id_conflict_existing(hass, entity_registry): """Test that we don't generate an entity id that currently exists.""" hass.states.async_set("light.hue_1234", "on") entry = entity_registry.async_get_or_create("light", "hue", "1234") assert entry.entity_id == "light.hue_1234_2" def test_create_triggers_save(entity_registry): """Test that registering entry triggers a save.""" with patch.object(entity_registry, "async_schedule_save") as mock_schedule_save: entity_registry.async_get_or_create("light", "hue", "1234") assert len(mock_schedule_save.mock_calls) == 1 async def test_loading_saving_data(hass, entity_registry): """Test that we load/save data correctly.""" mock_config = MockConfigEntry(domain="light") orig_entry1 = entity_registry.async_get_or_create("light", "hue", "1234") orig_entry2 = entity_registry.async_get_or_create( "light", "hue", "5678", capabilities={"max": 100}, config_entry=mock_config, device_id="mock-dev-id", disabled_by=er.RegistryEntryDisabler.HASS, entity_category=EntityCategory.CONFIG, hidden_by=er.RegistryEntryHider.INTEGRATION, has_entity_name=True, original_device_class="mock-device-class", original_icon="hass:original-icon", original_name="Original Name", supported_features=5, translation_key="initial-translation_key", unit_of_measurement="initial-unit_of_measurement", ) entity_registry.async_update_entity( orig_entry2.entity_id, aliases={"initial_alias_1", "initial_alias_2"}, area_id="mock-area-id", device_class="user-class", name="User Name", icon="hass:user-icon", ) entity_registry.async_update_entity_options( orig_entry2.entity_id, "light", {"minimum_brightness": 20} ) orig_entry2 = entity_registry.async_get(orig_entry2.entity_id) assert len(entity_registry.entities) == 2 # Now load written data in new registry registry2 = er.EntityRegistry(hass) await flush_store(entity_registry._store) await registry2.async_load() # Ensure same order assert list(entity_registry.entities) == list(registry2.entities) new_entry1 = entity_registry.async_get_or_create("light", "hue", "1234") new_entry2 = entity_registry.async_get_or_create("light", "hue", "5678") assert orig_entry1 == new_entry1 assert orig_entry2 == new_entry2 assert new_entry2.area_id == "mock-area-id" 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.disabled_by is er.RegistryEntryDisabler.HASS assert new_entry2.entity_category == "config" assert new_entry2.icon == "hass:user-icon" assert new_entry2.hidden_by == er.RegistryEntryHider.INTEGRATION assert new_entry2.has_entity_name is True assert new_entry2.name == "User Name" assert new_entry2.options == {"light": {"minimum_brightness": 20}} assert new_entry2.original_device_class == "mock-device-class" assert new_entry2.original_icon == "hass:original-icon" assert new_entry2.original_name == "Original Name" assert new_entry2.supported_features == 5 assert new_entry2.translation_key == "initial-translation_key" assert new_entry2.unit_of_measurement == "initial-unit_of_measurement" def test_generate_entity_considers_registered_entities(entity_registry): """Test that we don't create entity id that are already registered.""" entry = entity_registry.async_get_or_create("light", "hue", "1234") assert entry.entity_id == "light.hue_1234" assert ( entity_registry.async_generate_entity_id("light", "hue_1234") == "light.hue_1234_2" ) def test_generate_entity_considers_existing_entities(hass, entity_registry): """Test that we don't create entity id that currently exists.""" hass.states.async_set("light.kitchen", "on") assert ( entity_registry.async_generate_entity_id("light", "kitchen") == "light.kitchen_2" ) def test_is_registered(entity_registry): """Test that is_registered works.""" entry = entity_registry.async_get_or_create("light", "hue", "1234") assert entity_registry.async_is_registered(entry.entity_id) assert not entity_registry.async_is_registered("light.non_existing") @pytest.mark.parametrize("load_registries", [False]) async def test_filter_on_load(hass, hass_storage): """Test we transform some data when loading from storage.""" hass_storage[er.STORAGE_KEY] = { "version": er.STORAGE_VERSION_MAJOR, "minor_version": 1, "data": { "entities": [ { "entity_id": "test.named", "platform": "super_platform", "unique_id": "with-name", "name": "registry override", }, # This entity's name should be None { "entity_id": "test.no_name", "platform": "super_platform", "unique_id": "without-name", }, { "entity_id": "test.disabled_user", "platform": "super_platform", "unique_id": "disabled-user", "disabled_by": "user", # We store the string representation }, { "entity_id": "test.disabled_hass", "platform": "super_platform", "unique_id": "disabled-hass", "disabled_by": "hass", # We store the string representation }, # This entry should have the entity_category reset to None { "entity_id": "test.system_entity", "platform": "super_platform", "unique_id": "system-entity", "entity_category": "system", }, ] }, } await er.async_load(hass) registry = er.async_get(hass) assert len(registry.entities) == 5 assert set(registry.entities.keys()) == { "test.disabled_hass", "test.disabled_user", "test.named", "test.no_name", "test.system_entity", } entry_with_name = registry.async_get_or_create( "test", "super_platform", "with-name" ) entry_without_name = registry.async_get_or_create( "test", "super_platform", "without-name" ) assert entry_with_name.name == "registry override" assert entry_without_name.name is None assert not entry_with_name.disabled entry_disabled_hass = registry.async_get_or_create( "test", "super_platform", "disabled-hass" ) entry_disabled_user = registry.async_get_or_create( "test", "super_platform", "disabled-user" ) assert entry_disabled_hass.disabled assert entry_disabled_hass.disabled_by is er.RegistryEntryDisabler.HASS assert entry_disabled_user.disabled assert entry_disabled_user.disabled_by is er.RegistryEntryDisabler.USER entry_system_category = registry.async_get_or_create( "test", "system_entity", "system-entity" ) assert entry_system_category.entity_category is None def test_async_get_entity_id(entity_registry): """Test that entity_id is returned.""" entry = entity_registry.async_get_or_create("light", "hue", "1234") assert entry.entity_id == "light.hue_1234" assert ( entity_registry.async_get_entity_id("light", "hue", "1234") == "light.hue_1234" ) assert entity_registry.async_get_entity_id("light", "hue", "123") is None async def test_updating_config_entry_id(hass, entity_registry, update_events): """Test that we update config entry id in registry.""" mock_config_1 = MockConfigEntry(domain="light", entry_id="mock-id-1") entry = entity_registry.async_get_or_create( "light", "hue", "5678", config_entry=mock_config_1 ) mock_config_2 = MockConfigEntry(domain="light", entry_id="mock-id-2") entry2 = entity_registry.async_get_or_create( "light", "hue", "5678", config_entry=mock_config_2 ) assert entry.entity_id == entry2.entity_id assert entry2.config_entry_id == "mock-id-2" await hass.async_block_till_done() assert len(update_events) == 2 assert update_events[0]["action"] == "create" assert update_events[0]["entity_id"] == entry.entity_id assert update_events[1]["action"] == "update" assert update_events[1]["entity_id"] == entry.entity_id assert update_events[1]["changes"] == {"config_entry_id": "mock-id-1"} async def test_removing_config_entry_id(hass, entity_registry, update_events): """Test that we update config entry id in registry.""" mock_config = MockConfigEntry(domain="light", entry_id="mock-id-1") entry = entity_registry.async_get_or_create( "light", "hue", "5678", config_entry=mock_config ) assert entry.config_entry_id == "mock-id-1" entity_registry.async_clear_config_entry("mock-id-1") assert not entity_registry.entities await hass.async_block_till_done() assert len(update_events) == 2 assert update_events[0]["action"] == "create" assert update_events[0]["entity_id"] == entry.entity_id assert update_events[1]["action"] == "remove" assert update_events[1]["entity_id"] == entry.entity_id async def test_removing_area_id(entity_registry): """Make sure we can clear area id.""" entry = entity_registry.async_get_or_create("light", "hue", "5678") entry_w_area = entity_registry.async_update_entity( entry.entity_id, area_id="12345A" ) entity_registry.async_clear_area_id("12345A") entry_wo_area = entity_registry.async_get(entry.entity_id) assert not entry_wo_area.area_id assert entry_w_area != entry_wo_area @pytest.mark.parametrize("load_registries", [False]) async def test_migration_1_1(hass, hass_storage): """Test migration from version 1.1.""" hass_storage[er.STORAGE_KEY] = { "version": 1, "minor_version": 1, "data": { "entities": [ { "device_class": "best_class", "entity_id": "test.entity", "platform": "super_platform", "unique_id": "very_unique", }, ] }, } await er.async_load(hass) registry = er.async_get(hass) entry = registry.async_get_or_create("test", "super_platform", "very_unique") assert entry.device_class is None assert entry.original_device_class == "best_class" @pytest.mark.parametrize("load_registries", [False]) async def test_migration_1_7(hass, hass_storage): """Test migration from version 1.7. This tests cleanup after frontend bug which incorrectly updated device_class """ entity_dict = { "area_id": None, "capabilities": {}, "config_entry_id": None, "device_id": None, "disabled_by": None, "entity_category": None, "has_entity_name": False, "hidden_by": None, "icon": None, "id": "12345", "name": None, "options": None, "original_icon": None, "original_name": None, "platform": "super_platform", "supported_features": 0, "unique_id": "very_unique", "unit_of_measurement": None, } hass_storage[er.STORAGE_KEY] = { "version": 1, "minor_version": 7, "data": { "entities": [ { **entity_dict, "device_class": "original_class_by_integration", "entity_id": "test.entity", "original_device_class": "new_class_by_integration", }, { **entity_dict, "device_class": "class_by_user", "entity_id": "binary_sensor.entity", "original_device_class": "class_by_integration", }, { **entity_dict, "device_class": "class_by_user", "entity_id": "cover.entity", "original_device_class": "class_by_integration", }, ] }, } await er.async_load(hass) registry = er.async_get(hass) entry = registry.async_get_or_create("test", "super_platform", "very_unique") assert entry.device_class is None assert entry.original_device_class == "new_class_by_integration" entry = registry.async_get_or_create( "binary_sensor", "super_platform", "very_unique" ) assert entry.device_class == "class_by_user" assert entry.original_device_class == "class_by_integration" entry = registry.async_get_or_create("cover", "super_platform", "very_unique") assert entry.device_class == "class_by_user" assert entry.original_device_class == "class_by_integration" async def test_update_entity_unique_id(entity_registry): """Test entity's unique_id is updated.""" mock_config = MockConfigEntry(domain="light", entry_id="mock-id-1") entry = entity_registry.async_get_or_create( "light", "hue", "5678", config_entry=mock_config ) assert ( entity_registry.async_get_entity_id("light", "hue", "5678") == entry.entity_id ) new_unique_id = "1234" with patch.object(entity_registry, "async_schedule_save") as mock_schedule_save: updated_entry = entity_registry.async_update_entity( entry.entity_id, new_unique_id=new_unique_id ) assert updated_entry != entry assert updated_entry.unique_id == new_unique_id assert mock_schedule_save.call_count == 1 assert entity_registry.async_get_entity_id("light", "hue", "5678") is None assert ( entity_registry.async_get_entity_id("light", "hue", "1234") == entry.entity_id ) async def test_update_entity_unique_id_conflict(entity_registry): """Test migration raises when unique_id already in use.""" mock_config = MockConfigEntry(domain="light", entry_id="mock-id-1") entry = entity_registry.async_get_or_create( "light", "hue", "5678", config_entry=mock_config ) entry2 = entity_registry.async_get_or_create( "light", "hue", "1234", config_entry=mock_config ) with patch.object( entity_registry, "async_schedule_save" ) as mock_schedule_save, pytest.raises(ValueError): entity_registry.async_update_entity( entry.entity_id, new_unique_id=entry2.unique_id ) assert mock_schedule_save.call_count == 0 assert ( entity_registry.async_get_entity_id("light", "hue", "5678") == entry.entity_id ) assert ( entity_registry.async_get_entity_id("light", "hue", "1234") == entry2.entity_id ) async def test_update_entity_entity_id(entity_registry): """Test entity's entity_id is updated.""" entry = entity_registry.async_get_or_create("light", "hue", "5678") assert ( entity_registry.async_get_entity_id("light", "hue", "5678") == entry.entity_id ) new_entity_id = "light.blah" assert new_entity_id != entry.entity_id with patch.object(entity_registry, "async_schedule_save") as mock_schedule_save: updated_entry = entity_registry.async_update_entity( entry.entity_id, new_entity_id=new_entity_id ) assert updated_entry != entry assert updated_entry.entity_id == new_entity_id assert mock_schedule_save.call_count == 1 assert entity_registry.async_get(entry.entity_id) is None assert entity_registry.async_get(new_entity_id) is not None async def test_update_entity_entity_id_entity_id(hass: HomeAssistant, entity_registry): """Test update raises when entity_id already in use.""" entry = entity_registry.async_get_or_create("light", "hue", "5678") entry2 = entity_registry.async_get_or_create("light", "hue", "1234") state_entity_id = "light.blah" hass.states.async_set(state_entity_id, "on") assert entry.entity_id != state_entity_id assert entry2.entity_id != state_entity_id # Try updating to a registered entity_id with patch.object( entity_registry, "async_schedule_save" ) as mock_schedule_save, pytest.raises(ValueError): entity_registry.async_update_entity( entry.entity_id, new_entity_id=entry2.entity_id ) assert mock_schedule_save.call_count == 0 assert ( entity_registry.async_get_entity_id("light", "hue", "5678") == entry.entity_id ) assert entity_registry.async_get(entry.entity_id) is entry assert ( entity_registry.async_get_entity_id("light", "hue", "1234") == entry2.entity_id ) assert entity_registry.async_get(entry2.entity_id) is entry2 # Try updating to an entity_id which is in the state machine with patch.object( entity_registry, "async_schedule_save" ) as mock_schedule_save, pytest.raises(ValueError): entity_registry.async_update_entity( entry.entity_id, new_entity_id=state_entity_id ) assert mock_schedule_save.call_count == 0 assert ( entity_registry.async_get_entity_id("light", "hue", "5678") == entry.entity_id ) assert entity_registry.async_get(entry.entity_id) is entry assert entity_registry.async_get(state_entity_id) is None async def test_update_entity(entity_registry): """Test updating entity.""" mock_config = MockConfigEntry(domain="light", entry_id="mock-id-1") entry = entity_registry.async_get_or_create( "light", "hue", "5678", config_entry=mock_config ) for attr_name, new_value in ( ("aliases", {"alias_1", "alias_2"}), ("disabled_by", er.RegistryEntryDisabler.USER), ("icon", "new icon"), ("name", "new name"), ): changes = {attr_name: new_value} updated_entry = entity_registry.async_update_entity(entry.entity_id, **changes) assert updated_entry != entry assert getattr(updated_entry, attr_name) == new_value assert getattr(updated_entry, attr_name) != getattr(entry, attr_name) assert ( entity_registry.async_get_entity_id("light", "hue", "5678") == updated_entry.entity_id ) entry = updated_entry async def test_update_entity_options(entity_registry): """Test updating entity.""" mock_config = MockConfigEntry(domain="light", entry_id="mock-id-1") entry = entity_registry.async_get_or_create( "light", "hue", "5678", config_entry=mock_config ) entity_registry.async_update_entity_options( entry.entity_id, "light", {"minimum_brightness": 20} ) new_entry_1 = entity_registry.async_get(entry.entity_id) assert entry.options == {} assert new_entry_1.options == {"light": {"minimum_brightness": 20}} entity_registry.async_update_entity_options( entry.entity_id, "light", {"minimum_brightness": 30} ) new_entry_2 = entity_registry.async_get(entry.entity_id) assert entry.options == {} assert new_entry_1.options == {"light": {"minimum_brightness": 20}} assert new_entry_2.options == {"light": {"minimum_brightness": 30}} async def test_disabled_by(entity_registry): """Test that we can disable an entry when we create it.""" entry = entity_registry.async_get_or_create( "light", "hue", "5678", disabled_by=er.RegistryEntryDisabler.HASS ) assert entry.disabled_by is er.RegistryEntryDisabler.HASS entry = entity_registry.async_get_or_create( "light", "hue", "5678", disabled_by=er.RegistryEntryDisabler.INTEGRATION ) assert entry.disabled_by is er.RegistryEntryDisabler.HASS entry2 = entity_registry.async_get_or_create("light", "hue", "1234") assert entry2.disabled_by is None async def test_disabled_by_config_entry_pref(entity_registry): """Test config entry preference setting disabled_by.""" mock_config = MockConfigEntry( domain="light", entry_id="mock-id-1", pref_disable_new_entities=True, ) entry = entity_registry.async_get_or_create( "light", "hue", "AAAA", config_entry=mock_config ) assert entry.disabled_by is er.RegistryEntryDisabler.INTEGRATION entry2 = entity_registry.async_get_or_create( "light", "hue", "BBBB", config_entry=mock_config, disabled_by=er.RegistryEntryDisabler.USER, ) assert entry2.disabled_by is er.RegistryEntryDisabler.USER async def test_restore_states(hass: HomeAssistant) -> None: """Test restoring states.""" hass.state = CoreState.not_running registry = er.async_get(hass) registry.async_get_or_create( "light", "hue", "1234", suggested_object_id="simple", ) # Should not be created registry.async_get_or_create( "light", "hue", "5678", suggested_object_id="disabled", disabled_by=er.RegistryEntryDisabler.HASS, ) registry.async_get_or_create( "light", "hue", "9012", suggested_object_id="all_info_set", capabilities={"max": 100}, supported_features=5, original_device_class="mock-device-class", original_name="Mock Original Name", original_icon="hass:original-icon", ) hass.bus.async_fire(EVENT_HOMEASSISTANT_START, {}) await hass.async_block_till_done() simple = hass.states.get("light.simple") assert simple is not None assert simple.state == STATE_UNAVAILABLE assert simple.attributes == {"restored": True, "supported_features": 0} disabled = hass.states.get("light.disabled") assert disabled is None all_info_set = hass.states.get("light.all_info_set") assert all_info_set is not None assert all_info_set.state == STATE_UNAVAILABLE assert all_info_set.attributes == { "max": 100, "supported_features": 5, "device_class": "mock-device-class", "restored": True, "friendly_name": "Mock Original Name", "icon": "hass:original-icon", } registry.async_remove("light.disabled") registry.async_remove("light.simple") registry.async_remove("light.all_info_set") await hass.async_block_till_done() assert hass.states.get("light.simple") is None assert hass.states.get("light.disabled") is None assert hass.states.get("light.all_info_set") is None async def test_async_get_device_class_lookup(hass: HomeAssistant) -> None: """Test registry device class lookup.""" hass.state = CoreState.not_running ent_reg = er.async_get(hass) ent_reg.async_get_or_create( "binary_sensor", "light", "battery_charging", device_id="light_device_entry_id", original_device_class="battery_charging", ) ent_reg.async_get_or_create( "sensor", "light", "battery", device_id="light_device_entry_id", original_device_class="battery", ) ent_reg.async_get_or_create( "light", "light", "demo", device_id="light_device_entry_id" ) ent_reg.async_get_or_create( "binary_sensor", "vacuum", "battery_charging", device_id="vacuum_device_entry_id", original_device_class="battery_charging", ) ent_reg.async_get_or_create( "sensor", "vacuum", "battery", device_id="vacuum_device_entry_id", original_device_class="battery", ) ent_reg.async_get_or_create( "vacuum", "vacuum", "demo", device_id="vacuum_device_entry_id" ) ent_reg.async_get_or_create( "binary_sensor", "remote", "battery_charging", device_id="remote_device_entry_id", original_device_class="battery_charging", ) ent_reg.async_get_or_create( "remote", "remote", "demo", device_id="remote_device_entry_id" ) device_lookup = ent_reg.async_get_device_class_lookup( {("binary_sensor", "battery_charging"), ("sensor", "battery")} ) assert device_lookup == { "remote_device_entry_id": { ( "binary_sensor", "battery_charging", ): "binary_sensor.remote_battery_charging" }, "light_device_entry_id": { ( "binary_sensor", "battery_charging", ): "binary_sensor.light_battery_charging", ("sensor", "battery"): "sensor.light_battery", }, "vacuum_device_entry_id": { ( "binary_sensor", "battery_charging", ): "binary_sensor.vacuum_battery_charging", ("sensor", "battery"): "sensor.vacuum_battery", }, } async def test_remove_device_removes_entities(hass, entity_registry, device_registry): """Test that we remove entities tied to a device.""" config_entry = MockConfigEntry(domain="light") 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")}, ) entry = entity_registry.async_get_or_create( "light", "hue", "5678", config_entry=config_entry, device_id=device_entry.id, ) assert entity_registry.async_is_registered(entry.entity_id) device_registry.async_remove_device(device_entry.id) await hass.async_block_till_done() assert not entity_registry.async_is_registered(entry.entity_id) async def test_remove_config_entry_from_device_removes_entities( hass, device_registry, entity_registry ): """Test that we remove entities tied to a device when config entry is removed.""" config_entry_1 = MockConfigEntry(domain="hue") config_entry_2 = MockConfigEntry(domain="device_tracker") # Create device with two config entries device_registry.async_get_or_create( config_entry_id=config_entry_1.entry_id, connections={(dr.CONNECTION_NETWORK_MAC, "12:34:56:AB:CD:EF")}, ) device_entry = device_registry.async_get_or_create( config_entry_id=config_entry_2.entry_id, connections={(dr.CONNECTION_NETWORK_MAC, "12:34:56:AB:CD:EF")}, ) assert device_entry.config_entries == { config_entry_1.entry_id, config_entry_2.entry_id, } # Create one entity for each config entry entry_1 = entity_registry.async_get_or_create( "light", "hue", "5678", config_entry=config_entry_1, device_id=device_entry.id, ) entry_2 = entity_registry.async_get_or_create( "sensor", "device_tracker", "6789", config_entry=config_entry_2, device_id=device_entry.id, ) assert entity_registry.async_is_registered(entry_1.entity_id) assert entity_registry.async_is_registered(entry_2.entity_id) # Remove the first config entry from the device, the entity associated with it # should be removed device_registry.async_update_device( device_entry.id, remove_config_entry_id=config_entry_1.entry_id ) await hass.async_block_till_done() assert device_registry.async_get(device_entry.id) assert not entity_registry.async_is_registered(entry_1.entity_id) assert entity_registry.async_is_registered(entry_2.entity_id) # Remove the second config entry from the device, the entity associated with it # (and the device itself) should be removed device_registry.async_update_device( device_entry.id, remove_config_entry_id=config_entry_2.entry_id ) await hass.async_block_till_done() assert not device_registry.async_get(device_entry.id) assert not entity_registry.async_is_registered(entry_1.entity_id) assert not entity_registry.async_is_registered(entry_2.entity_id) async def test_remove_config_entry_from_device_removes_entities_2( hass, device_registry, entity_registry ): """Test that we don't remove entities with no config entry when device is modified.""" config_entry_1 = MockConfigEntry(domain="hue") config_entry_2 = MockConfigEntry(domain="device_tracker") # Create device with two config entries device_registry.async_get_or_create( config_entry_id=config_entry_1.entry_id, connections={(dr.CONNECTION_NETWORK_MAC, "12:34:56:AB:CD:EF")}, ) device_entry = device_registry.async_get_or_create( config_entry_id=config_entry_2.entry_id, connections={(dr.CONNECTION_NETWORK_MAC, "12:34:56:AB:CD:EF")}, ) assert device_entry.config_entries == { config_entry_1.entry_id, config_entry_2.entry_id, } # Create one entity for each config entry entry_1 = entity_registry.async_get_or_create( "light", "hue", "5678", device_id=device_entry.id, ) assert entity_registry.async_is_registered(entry_1.entity_id) # Remove the first config entry from the device device_registry.async_update_device( device_entry.id, remove_config_entry_id=config_entry_1.entry_id ) await hass.async_block_till_done() assert device_registry.async_get(device_entry.id) assert entity_registry.async_is_registered(entry_1.entity_id) async def test_update_device_race(hass, device_registry, entity_registry): """Test race when a device is created, updated and removed.""" config_entry = MockConfigEntry(domain="light") # Create device 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")}, ) # Update it device_registry.async_get_or_create( config_entry_id=config_entry.entry_id, identifiers={("bridgeid", "0123")}, connections={(dr.CONNECTION_NETWORK_MAC, "12:34:56:AB:CD:EF")}, ) # Add entity to the device entry = entity_registry.async_get_or_create( "light", "hue", "5678", config_entry=config_entry, device_id=device_entry.id, ) assert entity_registry.async_is_registered(entry.entity_id) device_registry.async_remove_device(device_entry.id) await hass.async_block_till_done() assert not entity_registry.async_is_registered(entry.entity_id) async def test_disable_device_disables_entities(hass, device_registry, entity_registry): """Test that we disable entities tied to a device.""" config_entry = MockConfigEntry(domain="light") 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")}, ) entry1 = entity_registry.async_get_or_create( "light", "hue", "5678", config_entry=config_entry, device_id=device_entry.id, ) entry2 = entity_registry.async_get_or_create( "light", "hue", "ABCD", config_entry=config_entry, device_id=device_entry.id, disabled_by=er.RegistryEntryDisabler.USER, ) entry3 = entity_registry.async_get_or_create( "light", "hue", "EFGH", config_entry=config_entry, device_id=device_entry.id, disabled_by=er.RegistryEntryDisabler.CONFIG_ENTRY, ) assert not entry1.disabled assert entry2.disabled assert entry3.disabled device_registry.async_update_device( device_entry.id, disabled_by=dr.DeviceEntryDisabler.USER ) await hass.async_block_till_done() entry1 = entity_registry.async_get(entry1.entity_id) assert entry1.disabled assert entry1.disabled_by is er.RegistryEntryDisabler.DEVICE entry2 = entity_registry.async_get(entry2.entity_id) assert entry2.disabled assert entry2.disabled_by is er.RegistryEntryDisabler.USER entry3 = entity_registry.async_get(entry3.entity_id) assert entry3.disabled assert entry3.disabled_by is er.RegistryEntryDisabler.CONFIG_ENTRY device_registry.async_update_device(device_entry.id, disabled_by=None) await hass.async_block_till_done() entry1 = entity_registry.async_get(entry1.entity_id) assert not entry1.disabled entry2 = entity_registry.async_get(entry2.entity_id) assert entry2.disabled assert entry2.disabled_by is er.RegistryEntryDisabler.USER entry3 = entity_registry.async_get(entry3.entity_id) assert entry3.disabled assert entry3.disabled_by is er.RegistryEntryDisabler.CONFIG_ENTRY async def test_disable_config_entry_disables_entities( hass, device_registry, entity_registry ): """Test that we disable entities tied to a config entry.""" config_entry = MockConfigEntry(domain="light") 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")}, ) entry1 = entity_registry.async_get_or_create( "light", "hue", "5678", config_entry=config_entry, device_id=device_entry.id, ) entry2 = entity_registry.async_get_or_create( "light", "hue", "ABCD", config_entry=config_entry, device_id=device_entry.id, disabled_by=er.RegistryEntryDisabler.USER, ) entry3 = entity_registry.async_get_or_create( "light", "hue", "EFGH", config_entry=config_entry, device_id=device_entry.id, disabled_by=er.RegistryEntryDisabler.DEVICE, ) assert not entry1.disabled assert entry2.disabled assert entry3.disabled await hass.config_entries.async_set_disabled_by( config_entry.entry_id, config_entries.ConfigEntryDisabler.USER ) await hass.async_block_till_done() entry1 = entity_registry.async_get(entry1.entity_id) assert entry1.disabled assert entry1.disabled_by is er.RegistryEntryDisabler.CONFIG_ENTRY entry2 = entity_registry.async_get(entry2.entity_id) assert entry2.disabled assert entry2.disabled_by is er.RegistryEntryDisabler.USER entry3 = entity_registry.async_get(entry3.entity_id) assert entry3.disabled assert entry3.disabled_by is er.RegistryEntryDisabler.DEVICE await hass.config_entries.async_set_disabled_by(config_entry.entry_id, None) await hass.async_block_till_done() entry1 = entity_registry.async_get(entry1.entity_id) assert not entry1.disabled entry2 = entity_registry.async_get(entry2.entity_id) assert entry2.disabled assert entry2.disabled_by is er.RegistryEntryDisabler.USER # The device was re-enabled, so entity disabled by the device will be re-enabled too entry3 = entity_registry.async_get(entry3.entity_id) assert not entry3.disabled_by async def test_disabled_entities_excluded_from_entity_list( hass, device_registry, entity_registry ): """Test that disabled entities are excluded from async_entries_for_device.""" config_entry = MockConfigEntry(domain="light") 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")}, ) entry1 = entity_registry.async_get_or_create( "light", "hue", "5678", config_entry=config_entry, device_id=device_entry.id, ) entry2 = entity_registry.async_get_or_create( "light", "hue", "ABCD", config_entry=config_entry, device_id=device_entry.id, disabled_by=er.RegistryEntryDisabler.USER, ) entries = er.async_entries_for_device(entity_registry, device_entry.id) assert entries == [entry1] entries = er.async_entries_for_device( entity_registry, device_entry.id, include_disabled_entities=True ) assert entries == [entry1, entry2] async def test_entity_max_length_exceeded(entity_registry): """Test that an exception is raised when the max character length is exceeded.""" long_domain_name = ( "1234567890123456789012345678901234567890123456789012345678901234567890" "1234567890123456789012345678901234567890123456789012345678901234567890" "1234567890123456789012345678901234567890123456789012345678901234567890" "1234567890123456789012345678901234567890123456789012345678901234567890" ) with pytest.raises(MaxLengthExceeded) as exc_info: entity_registry.async_generate_entity_id(long_domain_name, "sensor") assert exc_info.value.property_name == "domain" assert exc_info.value.max_length == 64 assert exc_info.value.value == long_domain_name # Try again but force a number to get added to the entity ID long_entity_id_name = ( "1234567890123456789012345678901234567890123456789012345678901234567890" "1234567890123456789012345678901234567890123456789012345678901234567890" "1234567890123456789012345678901234567890123456789012345678901234567890" "1234567890123456789012345678901234567" ) known = [] new_id = entity_registry.async_generate_entity_id( "sensor", long_entity_id_name, known ) assert new_id == "sensor." + long_entity_id_name[: 255 - 7] known.append(new_id) new_id = entity_registry.async_generate_entity_id( "sensor", long_entity_id_name, known ) assert new_id == "sensor." + long_entity_id_name[: 255 - 7 - 2] + "_2" known.append(new_id) new_id = entity_registry.async_generate_entity_id( "sensor", long_entity_id_name, known ) assert new_id == "sensor." + long_entity_id_name[: 255 - 7 - 2] + "_3" async def test_resolve_entity_ids(entity_registry): """Test resolving entity IDs.""" entry1 = entity_registry.async_get_or_create( "light", "hue", "1234", suggested_object_id="beer" ) assert entry1.entity_id == "light.beer" entry2 = entity_registry.async_get_or_create( "light", "hue", "2345", suggested_object_id="milk" ) assert entry2.entity_id == "light.milk" expected = ["light.beer", "light.milk"] assert ( er.async_validate_entity_ids(entity_registry, [entry1.id, entry2.id]) == expected ) expected = ["light.beer", "light.milk"] assert ( er.async_validate_entity_ids(entity_registry, ["light.beer", entry2.id]) == expected ) with pytest.raises(vol.Invalid): er.async_validate_entity_ids(entity_registry, ["light.beer", "bad_uuid"]) expected = ["light.unknown"] assert er.async_validate_entity_ids(entity_registry, ["light.unknown"]) == expected with pytest.raises(vol.Invalid): er.async_validate_entity_ids(entity_registry, ["unknown_uuid"]) def test_entity_registry_items() -> None: """Test the EntityRegistryItems container.""" entities = er.EntityRegistryItems() assert entities.get_entity_id(("a", "b", "c")) is None assert entities.get_entry("abc") is None entry1 = er.RegistryEntry("test.entity1", "1234", "hue") entry2 = er.RegistryEntry("test.entity2", "2345", "hue") entities["test.entity1"] = entry1 entities["test.entity2"] = entry2 assert entities["test.entity1"] is entry1 assert entities["test.entity2"] is entry2 assert entities.get_entity_id(("test", "hue", "1234")) is entry1.entity_id assert entities.get_entry(entry1.id) is entry1 assert entities.get_entity_id(("test", "hue", "2345")) is entry2.entity_id assert entities.get_entry(entry2.id) is entry2 entities.pop("test.entity1") del entities["test.entity2"] assert entities.get_entity_id(("test", "hue", "1234")) is None assert entities.get_entry(entry1.id) is None assert entities.get_entity_id(("test", "hue", "2345")) is None assert entities.get_entry(entry2.id) is None async def test_disabled_by_str_not_allowed(hass: HomeAssistant) -> None: """Test we need to pass disabled by type.""" reg = er.async_get(hass) with pytest.raises(ValueError): reg.async_get_or_create( "light", "hue", "1234", disabled_by=er.RegistryEntryDisabler.USER.value ) entity_id = reg.async_get_or_create("light", "hue", "1234").entity_id with pytest.raises(ValueError): reg.async_update_entity( entity_id, disabled_by=er.RegistryEntryDisabler.USER.value ) async def test_entity_category_str_not_allowed(hass: HomeAssistant) -> None: """Test we need to pass entity category type.""" reg = er.async_get(hass) with pytest.raises(ValueError): reg.async_get_or_create( "light", "hue", "1234", entity_category=EntityCategory.DIAGNOSTIC.value ) entity_id = reg.async_get_or_create("light", "hue", "1234").entity_id with pytest.raises(ValueError): reg.async_update_entity( entity_id, entity_category=EntityCategory.DIAGNOSTIC.value ) async def test_hidden_by_str_not_allowed(hass: HomeAssistant) -> None: """Test we need to pass hidden by type.""" reg = er.async_get(hass) with pytest.raises(ValueError): reg.async_get_or_create( "light", "hue", "1234", hidden_by=er.RegistryEntryHider.USER.value ) entity_id = reg.async_get_or_create("light", "hue", "1234").entity_id with pytest.raises(ValueError): reg.async_update_entity(entity_id, hidden_by=er.RegistryEntryHider.USER.value) def test_migrate_entity_to_new_platform(hass, entity_registry): """Test migrate_entity_to_new_platform.""" orig_config_entry = MockConfigEntry(domain="light") orig_unique_id = "5678" orig_entry = entity_registry.async_get_or_create( "light", "hue", orig_unique_id, suggested_object_id="light", config_entry=orig_config_entry, disabled_by=er.RegistryEntryDisabler.USER, entity_category=EntityCategory.CONFIG, original_device_class="mock-device-class", original_icon="initial-original_icon", original_name="initial-original_name", ) assert entity_registry.async_get("light.light") is orig_entry entity_registry.async_update_entity( "light.light", name="new_name", icon="new_icon", ) new_config_entry = MockConfigEntry(domain="light") new_unique_id = "1234" assert entity_registry.async_update_entity_platform( "light.light", "hue2", new_unique_id=new_unique_id, new_config_entry_id=new_config_entry.entry_id, ) assert not entity_registry.async_get_entity_id("light", "hue", orig_unique_id) assert (new_entry := entity_registry.async_get("light.light")) is not orig_entry assert new_entry.config_entry_id == new_config_entry.entry_id assert new_entry.unique_id == new_unique_id assert new_entry.name == "new_name" assert new_entry.icon == "new_icon" assert new_entry.platform == "hue2" # Test nonexisting entity with pytest.raises(KeyError): entity_registry.async_update_entity_platform( "light.not_a_real_light", "hue2", new_unique_id=new_unique_id, new_config_entry_id=new_config_entry.entry_id, ) # Test migrate entity without new config entry ID with pytest.raises(ValueError): entity_registry.async_update_entity_platform( "light.light", "hue3", ) # Test entity with a state hass.states.async_set("light.light", "on") with pytest.raises(ValueError): entity_registry.async_update_entity_platform( "light.light", "hue2", new_unique_id=new_unique_id, new_config_entry_id=new_config_entry.entry_id, )