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.

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."""

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."""

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."""

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."""

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."""

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, []))

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."""

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."""
 

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."""