From 8a78b65f0d99f1383cfdbcea13e95fbe7e3f2288 Mon Sep 17 00:00:00 2001 From: Alexei Chetroi Date: Fri, 17 Jan 2020 11:41:46 -0500 Subject: [PATCH] Extract collection entity registry cleaner into a helper (#30844) * Extract entity_registry_keeper into separate helper * Refactor input_*, timer and person to use new helper. * Make Mypy happy. * Better name. --- .../components/input_boolean/__init__.py | 16 +++--------- .../components/input_datetime/__init__.py | 16 +++--------- .../components/input_number/__init__.py | 16 +++--------- .../components/input_select/__init__.py | 16 +++--------- .../components/input_text/__init__.py | 16 +++--------- homeassistant/components/person/__init__.py | 13 ++-------- homeassistant/components/timer/__init__.py | 16 +++--------- homeassistant/helpers/collection.py | 26 +++++++++++++++++++ tests/components/input_datetime/test_init.py | 21 +++++++++++++-- 9 files changed, 65 insertions(+), 91 deletions(-) diff --git a/homeassistant/components/input_boolean/__init__.py b/homeassistant/components/input_boolean/__init__.py index c805af0a758..daadfac3705 100644 --- a/homeassistant/components/input_boolean/__init__.py +++ b/homeassistant/components/input_boolean/__init__.py @@ -16,7 +16,7 @@ from homeassistant.const import ( STATE_ON, ) from homeassistant.core import callback -from homeassistant.helpers import collection, entity_registry +from homeassistant.helpers import collection import homeassistant.helpers.config_validation as cv from homeassistant.helpers.entity import ToggleEntity from homeassistant.helpers.entity_component import EntityComponent @@ -113,18 +113,8 @@ async def async_setup(hass: HomeAssistantType, config: ConfigType) -> bool: storage_collection, DOMAIN, DOMAIN, CREATE_FIELDS, UPDATE_FIELDS ).async_setup(hass) - async def _collection_changed( - change_type: str, item_id: str, config: typing.Optional[typing.Dict] - ) -> None: - """Handle a collection change: clean up entity registry on removals.""" - if change_type != collection.CHANGE_REMOVED: - return - - ent_reg = await entity_registry.async_get_registry(hass) - ent_reg.async_remove(ent_reg.async_get_entity_id(DOMAIN, DOMAIN, item_id)) - - yaml_collection.async_add_listener(_collection_changed) - storage_collection.async_add_listener(_collection_changed) + collection.attach_entity_registry_cleaner(hass, DOMAIN, DOMAIN, yaml_collection) + collection.attach_entity_registry_cleaner(hass, DOMAIN, DOMAIN, storage_collection) async def reload_service_handler(service_call: ServiceCallType) -> None: """Remove all input booleans and load new ones from config.""" diff --git a/homeassistant/components/input_datetime/__init__.py b/homeassistant/components/input_datetime/__init__.py index fdb80591bbe..03b468313f8 100644 --- a/homeassistant/components/input_datetime/__init__.py +++ b/homeassistant/components/input_datetime/__init__.py @@ -15,7 +15,7 @@ from homeassistant.const import ( SERVICE_RELOAD, ) from homeassistant.core import callback -from homeassistant.helpers import collection, entity_registry +from homeassistant.helpers import collection import homeassistant.helpers.config_validation as cv from homeassistant.helpers.entity_component import EntityComponent from homeassistant.helpers.restore_state import RestoreEntity @@ -117,18 +117,8 @@ async def async_setup(hass: HomeAssistantType, config: ConfigType) -> bool: storage_collection, DOMAIN, DOMAIN, CREATE_FIELDS, UPDATE_FIELDS ).async_setup(hass) - async def _collection_changed( - change_type: str, item_id: str, config: typing.Optional[typing.Dict] - ) -> None: - """Handle a collection change: clean up entity registry on removals.""" - if change_type != collection.CHANGE_REMOVED: - return - - ent_reg = await entity_registry.async_get_registry(hass) - ent_reg.async_remove(ent_reg.async_get_entity_id(DOMAIN, DOMAIN, item_id)) - - yaml_collection.async_add_listener(_collection_changed) - storage_collection.async_add_listener(_collection_changed) + collection.attach_entity_registry_cleaner(hass, DOMAIN, DOMAIN, yaml_collection) + collection.attach_entity_registry_cleaner(hass, DOMAIN, DOMAIN, storage_collection) async def reload_service_handler(service_call: ServiceCallType) -> None: """Reload yaml entities.""" diff --git a/homeassistant/components/input_number/__init__.py b/homeassistant/components/input_number/__init__.py index 4205389d9b2..f78fc485e40 100644 --- a/homeassistant/components/input_number/__init__.py +++ b/homeassistant/components/input_number/__init__.py @@ -15,7 +15,7 @@ from homeassistant.const import ( SERVICE_RELOAD, ) from homeassistant.core import callback -from homeassistant.helpers import collection, entity_registry +from homeassistant.helpers import collection import homeassistant.helpers.config_validation as cv from homeassistant.helpers.entity_component import EntityComponent from homeassistant.helpers.restore_state import RestoreEntity @@ -142,18 +142,8 @@ async def async_setup(hass: HomeAssistantType, config: ConfigType) -> bool: storage_collection, DOMAIN, DOMAIN, CREATE_FIELDS, UPDATE_FIELDS ).async_setup(hass) - async def _collection_changed( - change_type: str, item_id: str, config: typing.Optional[typing.Dict] - ) -> None: - """Handle a collection change: clean up entity registry on removals.""" - if change_type != collection.CHANGE_REMOVED: - return - - ent_reg = await entity_registry.async_get_registry(hass) - ent_reg.async_remove(ent_reg.async_get_entity_id(DOMAIN, DOMAIN, item_id)) - - yaml_collection.async_add_listener(_collection_changed) - storage_collection.async_add_listener(_collection_changed) + collection.attach_entity_registry_cleaner(hass, DOMAIN, DOMAIN, yaml_collection) + collection.attach_entity_registry_cleaner(hass, DOMAIN, DOMAIN, storage_collection) async def reload_service_handler(service_call: ServiceCallType) -> None: """Reload yaml entities.""" diff --git a/homeassistant/components/input_select/__init__.py b/homeassistant/components/input_select/__init__.py index 8d86904e5f9..26a07e600f3 100644 --- a/homeassistant/components/input_select/__init__.py +++ b/homeassistant/components/input_select/__init__.py @@ -12,7 +12,7 @@ from homeassistant.const import ( SERVICE_RELOAD, ) from homeassistant.core import callback -from homeassistant.helpers import collection, entity_registry +from homeassistant.helpers import collection import homeassistant.helpers.config_validation as cv from homeassistant.helpers.entity_component import EntityComponent from homeassistant.helpers.restore_state import RestoreEntity @@ -116,18 +116,8 @@ async def async_setup(hass: HomeAssistantType, config: ConfigType) -> bool: storage_collection, DOMAIN, DOMAIN, CREATE_FIELDS, UPDATE_FIELDS ).async_setup(hass) - async def _collection_changed( - change_type: str, item_id: str, config: typing.Optional[typing.Dict] - ) -> None: - """Handle a collection change: clean up entity registry on removals.""" - if change_type != collection.CHANGE_REMOVED: - return - - ent_reg = await entity_registry.async_get_registry(hass) - ent_reg.async_remove(ent_reg.async_get_entity_id(DOMAIN, DOMAIN, item_id)) - - yaml_collection.async_add_listener(_collection_changed) - storage_collection.async_add_listener(_collection_changed) + collection.attach_entity_registry_cleaner(hass, DOMAIN, DOMAIN, yaml_collection) + collection.attach_entity_registry_cleaner(hass, DOMAIN, DOMAIN, storage_collection) async def reload_service_handler(service_call: ServiceCallType) -> None: """Reload yaml entities.""" diff --git a/homeassistant/components/input_text/__init__.py b/homeassistant/components/input_text/__init__.py index c439d177224..bdb3e8a4bc9 100644 --- a/homeassistant/components/input_text/__init__.py +++ b/homeassistant/components/input_text/__init__.py @@ -15,7 +15,7 @@ from homeassistant.const import ( SERVICE_RELOAD, ) from homeassistant.core import callback -from homeassistant.helpers import collection, entity_registry +from homeassistant.helpers import collection import homeassistant.helpers.config_validation as cv from homeassistant.helpers.entity_component import EntityComponent from homeassistant.helpers.restore_state import RestoreEntity @@ -144,18 +144,8 @@ async def async_setup(hass: HomeAssistantType, config: ConfigType) -> bool: storage_collection, DOMAIN, DOMAIN, CREATE_FIELDS, UPDATE_FIELDS ).async_setup(hass) - async def _collection_changed( - change_type: str, item_id: str, config: typing.Optional[typing.Dict] - ) -> None: - """Handle a collection change: clean up entity registry on removals.""" - if change_type != collection.CHANGE_REMOVED: - return - - ent_reg = await entity_registry.async_get_registry(hass) - ent_reg.async_remove(ent_reg.async_get_entity_id(DOMAIN, DOMAIN, item_id)) - - yaml_collection.async_add_listener(_collection_changed) - storage_collection.async_add_listener(_collection_changed) + collection.attach_entity_registry_cleaner(hass, DOMAIN, DOMAIN, yaml_collection) + collection.attach_entity_registry_cleaner(hass, DOMAIN, DOMAIN, storage_collection) async def reload_service_handler(service_call: ServiceCallType) -> None: """Reload yaml entities.""" diff --git a/homeassistant/components/person/__init__.py b/homeassistant/components/person/__init__.py index c34fb89a718..dabcc046f7a 100644 --- a/homeassistant/components/person/__init__.py +++ b/homeassistant/components/person/__init__.py @@ -159,7 +159,6 @@ class PersonStorageCollection(collection.StorageCollection): ): """Initialize a person storage collection.""" super().__init__(store, logger, id_manager) - self.async_add_listener(self._collection_changed) self.yaml_collection = yaml_collection async def async_load(self) -> None: @@ -228,16 +227,6 @@ class PersonStorageCollection(collection.StorageCollection): if any(person for person in persons if person.get(CONF_USER_ID) == user_id): raise ValueError("User already taken") - async def _collection_changed( - self, change_type: str, item_id: str, config: Optional[dict] - ) -> None: - """Handle a collection change.""" - if change_type != collection.CHANGE_REMOVED: - return - - ent_reg = await entity_registry.async_get_registry(self.hass) - ent_reg.async_remove(ent_reg.async_get_entity_id(DOMAIN, DOMAIN, item_id)) - async def filter_yaml_data(hass: HomeAssistantType, persons: List[dict]) -> List[dict]: """Validate YAML data that we can't validate via schema.""" @@ -294,6 +283,8 @@ async def async_setup(hass: HomeAssistantType, config: ConfigType): collection.attach_entity_component_collection( entity_component, storage_collection, lambda conf: Person(conf, True) ) + collection.attach_entity_registry_cleaner(hass, DOMAIN, DOMAIN, yaml_collection) + collection.attach_entity_registry_cleaner(hass, DOMAIN, DOMAIN, storage_collection) await yaml_collection.async_load( await filter_yaml_data(hass, config.get(DOMAIN, [])) diff --git a/homeassistant/components/timer/__init__.py b/homeassistant/components/timer/__init__.py index 1216bc8a239..abf3a6ab0f7 100644 --- a/homeassistant/components/timer/__init__.py +++ b/homeassistant/components/timer/__init__.py @@ -13,7 +13,7 @@ from homeassistant.const import ( SERVICE_RELOAD, ) from homeassistant.core import callback -from homeassistant.helpers import collection, entity_registry +from homeassistant.helpers import collection import homeassistant.helpers.config_validation as cv from homeassistant.helpers.entity_component import EntityComponent from homeassistant.helpers.event import async_track_point_in_utc_time @@ -119,18 +119,8 @@ async def async_setup(hass: HomeAssistantType, config: ConfigType) -> bool: storage_collection, DOMAIN, DOMAIN, CREATE_FIELDS, UPDATE_FIELDS ).async_setup(hass) - async def _collection_changed( - change_type: str, item_id: str, config: typing.Optional[typing.Dict] - ) -> None: - """Handle a collection change: clean up entity registry on removals.""" - if change_type != collection.CHANGE_REMOVED: - return - - ent_reg = await entity_registry.async_get_registry(hass) - ent_reg.async_remove(ent_reg.async_get_entity_id(DOMAIN, DOMAIN, item_id)) - - yaml_collection.async_add_listener(_collection_changed) - storage_collection.async_add_listener(_collection_changed) + collection.attach_entity_registry_cleaner(hass, DOMAIN, DOMAIN, yaml_collection) + collection.attach_entity_registry_cleaner(hass, DOMAIN, DOMAIN, storage_collection) async def reload_service_handler(service_call: ServiceCallType) -> None: """Reload yaml entities.""" diff --git a/homeassistant/helpers/collection.py b/homeassistant/helpers/collection.py index dd0edbd09b9..80790a6d831 100644 --- a/homeassistant/helpers/collection.py +++ b/homeassistant/helpers/collection.py @@ -10,9 +10,11 @@ from homeassistant.components import websocket_api from homeassistant.const import CONF_ID from homeassistant.core import HomeAssistant, callback from homeassistant.exceptions import HomeAssistantError +from homeassistant.helpers import entity_registry from homeassistant.helpers.entity import Entity from homeassistant.helpers.entity_component import EntityComponent from homeassistant.helpers.storage import Store +from homeassistant.helpers.typing import HomeAssistantType from homeassistant.util import slugify STORAGE_VERSION = 1 @@ -259,6 +261,30 @@ def attach_entity_component_collection( collection.async_add_listener(_collection_changed) +@callback +def attach_entity_registry_cleaner( + hass: HomeAssistantType, + domain: str, + platform: str, + collection: ObservableCollection, +) -> None: + """Attach a listener to clean up entity registry on collection changes.""" + + async def _collection_changed( + change_type: str, item_id: str, config: Optional[Dict] + ) -> None: + """Handle a collection change: clean up entity registry on removals.""" + if change_type != CHANGE_REMOVED: + return + + ent_reg = await entity_registry.async_get_registry(hass) + ent_to_remove = ent_reg.async_get_entity_id(domain, platform, item_id) + if ent_to_remove is not None: + ent_reg.async_remove(ent_to_remove) + + collection.async_add_listener(_collection_changed) + + class StorageCollectionWebsocket: """Class to expose storage collection management over websocket.""" diff --git a/tests/components/input_datetime/test_init.py b/tests/components/input_datetime/test_init.py index b2f92a333ea..67a23a61d8b 100644 --- a/tests/components/input_datetime/test_init.py +++ b/tests/components/input_datetime/test_init.py @@ -360,22 +360,33 @@ async def test_input_datetime_context(hass, hass_admin_user): async def test_reload(hass, hass_admin_user, hass_read_only_user): """Test reload service.""" count_start = len(hass.states.async_entity_ids()) + ent_reg = await entity_registry.async_get_registry(hass) assert await async_setup_component( hass, DOMAIN, - {DOMAIN: {"dt1": {"has_time": False, "has_date": True, "initial": "2019-1-1"}}}, + { + DOMAIN: { + "dt1": {"has_time": False, "has_date": True, "initial": "2019-1-1"}, + "dt3": {CONF_HAS_TIME: True, CONF_HAS_DATE: True}, + } + }, ) - assert count_start + 1 == len(hass.states.async_entity_ids()) + assert count_start + 2 == len(hass.states.async_entity_ids()) state_1 = hass.states.get("input_datetime.dt1") state_2 = hass.states.get("input_datetime.dt2") + state_3 = hass.states.get("input_datetime.dt3") dt_obj = datetime.datetime(2019, 1, 1, 0, 0) assert state_1 is not None assert state_2 is None + assert state_3 is not None assert str(dt_obj.date()) == state_1.state + assert ent_reg.async_get_entity_id(DOMAIN, DOMAIN, "dt1") == f"{DOMAIN}.dt1" + assert ent_reg.async_get_entity_id(DOMAIN, DOMAIN, "dt2") is None + assert ent_reg.async_get_entity_id(DOMAIN, DOMAIN, "dt3") == f"{DOMAIN}.dt3" with patch( "homeassistant.config.load_yaml_config_file", @@ -406,12 +417,18 @@ async def test_reload(hass, hass_admin_user, hass_read_only_user): state_1 = hass.states.get("input_datetime.dt1") state_2 = hass.states.get("input_datetime.dt2") + state_3 = hass.states.get("input_datetime.dt3") assert state_1 is not None assert state_2 is not None + assert state_3 is None assert str(DEFAULT_TIME) == state_1.state assert str(datetime.datetime(1970, 1, 1, 0, 0)) == state_2.state + assert ent_reg.async_get_entity_id(DOMAIN, DOMAIN, "dt1") == f"{DOMAIN}.dt1" + assert ent_reg.async_get_entity_id(DOMAIN, DOMAIN, "dt2") == f"{DOMAIN}.dt2" + assert ent_reg.async_get_entity_id(DOMAIN, DOMAIN, "dt3") is None + async def test_load_from_storage(hass, storage_setup): """Test set up from storage."""