From 4435c641decd0269e03ba752c35e0aca468c1ab3 Mon Sep 17 00:00:00 2001 From: Erik Montnemery Date: Wed, 8 Jun 2022 21:36:43 +0200 Subject: [PATCH] Enforce RegistryEntryHider in entity registry (#73219) --- homeassistant/helpers/entity_registry.py | 12 +++++- tests/components/group/test_config_flow.py | 5 ++- tests/components/group/test_init.py | 10 ++--- .../switch_as_x/test_config_flow.py | 4 +- tests/components/switch_as_x/test_init.py | 4 +- tests/helpers/test_entity_registry.py | 42 ++++++++++++++++--- 6 files changed, 60 insertions(+), 17 deletions(-) diff --git a/homeassistant/helpers/entity_registry.py b/homeassistant/helpers/entity_registry.py index d03d272b1ac..eb5590b7fdf 100644 --- a/homeassistant/helpers/entity_registry.py +++ b/homeassistant/helpers/entity_registry.py @@ -369,6 +369,8 @@ class EntityRegistry: if disabled_by and not isinstance(disabled_by, RegistryEntryDisabler): raise ValueError("disabled_by must be a RegistryEntryDisabler value") + if hidden_by and not isinstance(hidden_by, RegistryEntryHider): + raise ValueError("hidden_by must be a RegistryEntryHider value") if ( disabled_by is None @@ -520,6 +522,12 @@ class EntityRegistry: and not isinstance(disabled_by, RegistryEntryDisabler) ): raise ValueError("disabled_by must be a RegistryEntryDisabler value") + if ( + hidden_by + and hidden_by is not UNDEFINED + and not isinstance(hidden_by, RegistryEntryHider) + ): + raise ValueError("hidden_by must be a RegistryEntryHider value") from .entity import EntityCategory # pylint: disable=import-outside-toplevel @@ -729,7 +737,9 @@ class EntityRegistry: if entity["entity_category"] else None, entity_id=entity["entity_id"], - hidden_by=entity["hidden_by"], + hidden_by=RegistryEntryHider(entity["hidden_by"]) + if entity["hidden_by"] + else None, icon=entity["icon"], id=entity["id"], name=entity["name"], diff --git a/tests/components/group/test_config_flow.py b/tests/components/group/test_config_flow.py index 83741a2e851..9d6b099557d 100644 --- a/tests/components/group/test_config_flow.py +++ b/tests/components/group/test_config_flow.py @@ -348,7 +348,10 @@ async def test_all_options( @pytest.mark.parametrize( "hide_members,hidden_by_initial,hidden_by", - ((False, "integration", None), (True, None, "integration")), + ( + (False, er.RegistryEntryHider.INTEGRATION, None), + (True, None, er.RegistryEntryHider.INTEGRATION), + ), ) @pytest.mark.parametrize( "group_type,extra_input", diff --git a/tests/components/group/test_init.py b/tests/components/group/test_init.py index 56553ff263c..945f6555789 100644 --- a/tests/components/group/test_init.py +++ b/tests/components/group/test_init.py @@ -1421,12 +1421,12 @@ async def test_setup_and_remove_config_entry( @pytest.mark.parametrize( "hide_members,hidden_by_initial,hidden_by", ( - (False, "integration", "integration"), + (False, er.RegistryEntryHider.INTEGRATION, er.RegistryEntryHider.INTEGRATION), (False, None, None), - (False, "user", "user"), - (True, "integration", None), + (False, er.RegistryEntryHider.USER, er.RegistryEntryHider.USER), + (True, er.RegistryEntryHider.INTEGRATION, None), (True, None, None), - (True, "user", "user"), + (True, er.RegistryEntryHider.USER, er.RegistryEntryHider.USER), ), ) @pytest.mark.parametrize( @@ -1444,7 +1444,7 @@ async def test_unhide_members_on_remove( group_type: str, extra_options: dict[str, Any], hide_members: bool, - hidden_by_initial: str, + hidden_by_initial: er.RegistryEntryHider, hidden_by: str, ) -> None: """Test removing a config entry.""" diff --git a/tests/components/switch_as_x/test_config_flow.py b/tests/components/switch_as_x/test_config_flow.py index dc4ca96aa97..d80f7e24bb1 100644 --- a/tests/components/switch_as_x/test_config_flow.py +++ b/tests/components/switch_as_x/test_config_flow.py @@ -65,8 +65,8 @@ async def test_config_flow( @pytest.mark.parametrize( "hidden_by_before,hidden_by_after", ( - (er.RegistryEntryHider.USER.value, er.RegistryEntryHider.USER.value), - (None, er.RegistryEntryHider.INTEGRATION.value), + (er.RegistryEntryHider.USER, er.RegistryEntryHider.USER), + (None, er.RegistryEntryHider.INTEGRATION), ), ) @pytest.mark.parametrize("target_domain", PLATFORMS_TO_TEST) diff --git a/tests/components/switch_as_x/test_init.py b/tests/components/switch_as_x/test_init.py index e2b875b813b..9c3eec1884c 100644 --- a/tests/components/switch_as_x/test_init.py +++ b/tests/components/switch_as_x/test_init.py @@ -365,8 +365,8 @@ async def test_setup_and_remove_config_entry( @pytest.mark.parametrize( "hidden_by_before,hidden_by_after", ( - (er.RegistryEntryHider.USER.value, er.RegistryEntryHider.USER.value), - (er.RegistryEntryHider.INTEGRATION.value, None), + (er.RegistryEntryHider.USER, er.RegistryEntryHider.USER), + (er.RegistryEntryHider.INTEGRATION, None), ), ) @pytest.mark.parametrize("target_domain", PLATFORMS_TO_TEST) diff --git a/tests/helpers/test_entity_registry.py b/tests/helpers/test_entity_registry.py index 21d29736bd0..8f5b4a7d333 100644 --- a/tests/helpers/test_entity_registry.py +++ b/tests/helpers/test_entity_registry.py @@ -79,6 +79,7 @@ def test_get_or_create_updates_data(registry): device_id="mock-dev-id", disabled_by=er.RegistryEntryDisabler.HASS, entity_category=EntityCategory.CONFIG, + hidden_by=er.RegistryEntryHider.INTEGRATION, original_device_class="mock-device-class", original_icon="initial-original_icon", original_name="initial-original_name", @@ -97,6 +98,7 @@ def test_get_or_create_updates_data(registry): device_id="mock-dev-id", disabled_by=er.RegistryEntryDisabler.HASS, entity_category=EntityCategory.CONFIG, + hidden_by=er.RegistryEntryHider.INTEGRATION, icon=None, id=orig_entry.id, name=None, @@ -119,6 +121,7 @@ def test_get_or_create_updates_data(registry): device_id="new-mock-dev-id", disabled_by=er.RegistryEntryDisabler.USER, entity_category=None, + hidden_by=er.RegistryEntryHider.USER, original_device_class="new-mock-device-class", original_icon="updated-original_icon", original_name="updated-original_name", @@ -137,6 +140,7 @@ def test_get_or_create_updates_data(registry): device_id="new-mock-dev-id", disabled_by=er.RegistryEntryDisabler.HASS, # Should not be updated entity_category=EntityCategory.CONFIG, + hidden_by=er.RegistryEntryHider.INTEGRATION, # Should not be updated icon=None, id=orig_entry.id, name=None, @@ -191,6 +195,7 @@ async def test_loading_saving_data(hass, registry): device_id="mock-dev-id", disabled_by=er.RegistryEntryDisabler.HASS, entity_category=EntityCategory.CONFIG, + hidden_by=er.RegistryEntryHider.INTEGRATION, original_device_class="mock-device-class", original_icon="hass:original-icon", original_name="Original Name", @@ -231,6 +236,7 @@ async def test_loading_saving_data(hass, registry): 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.name == "User Name" assert new_entry2.options == {"light": {"minimum_brightness": 20}} assert new_entry2.original_device_class == "mock-device-class" @@ -261,8 +267,8 @@ def test_is_registered(registry): @pytest.mark.parametrize("load_registries", [False]) -async def test_loading_extra_values(hass, hass_storage): - """Test we load extra data from the registry.""" +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, @@ -274,6 +280,7 @@ async def test_loading_extra_values(hass, hass_storage): "unique_id": "with-name", "name": "registry override", }, + # This entity's name should be None { "entity_id": "test.no_name", "platform": "super_platform", @@ -283,20 +290,22 @@ async def test_loading_extra_values(hass, hass_storage): "entity_id": "test.disabled_user", "platform": "super_platform", "unique_id": "disabled-user", - "disabled_by": er.RegistryEntryDisabler.USER, + "disabled_by": "user", # We store the string representation }, { "entity_id": "test.disabled_hass", "platform": "super_platform", "unique_id": "disabled-hass", - "disabled_by": er.RegistryEntryDisabler.HASS, + "disabled_by": "hass", # We store the string representation }, + # This entry should not be loaded because the entity_id is invalid { "entity_id": "test.invalid__entity", "platform": "super_platform", "unique_id": "invalid-hass", - "disabled_by": er.RegistryEntryDisabler.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", @@ -311,6 +320,13 @@ async def test_loading_extra_values(hass, hass_storage): 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" @@ -1221,7 +1237,7 @@ def test_entity_registry_items(): async def test_disabled_by_str_not_allowed(hass): - """Test we need to pass entity category type.""" + """Test we need to pass disabled by type.""" reg = er.async_get(hass) with pytest.raises(ValueError): @@ -1252,6 +1268,20 @@ async def test_entity_category_str_not_allowed(hass): ) +async def test_hidden_by_str_not_allowed(hass): + """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, registry): """Test migrate_entity_to_new_platform.""" orig_config_entry = MockConfigEntry(domain="light")