From c2db3c50fd30100274abd52b9d0c1b2a87d2099b Mon Sep 17 00:00:00 2001 From: "J. Nick Koston" Date: Mon, 19 Feb 2024 16:44:35 -0600 Subject: [PATCH 1/7] Fix race in removing entities from the registry There was a race here because `_async_registry_updated` would previously run in a task and `self.registry_entry` would not get unset soon enough so it would write unavailable state instead of removing the entity. I've seen this problem happen in production but I never could figure out why until now since its much more likely/reproducible as there are less tasks in play. Any time the `_async_registry_updated` ran after the `async_remove` task, the entity wouldn't get removed and the state would get written as unavailable instead because `self.registry_entry` hasn't been unset yet. To avoid this and not break the api (`self.async_removed_from_registry` might still need to access `self.registry_entry`) a new flag is added called `self._removed_from_registry` --- homeassistant/helpers/entity.py | 29 ++++++++++++++++++++++- tests/helpers/test_entity.py | 41 +++++++++++++++++++++++++++++++++ 2 files changed, 69 insertions(+), 1 deletion(-) diff --git a/homeassistant/helpers/entity.py b/homeassistant/helpers/entity.py index 32aa97ab8fef42..f7497e77a943a8 100644 --- a/homeassistant/helpers/entity.py +++ b/homeassistant/helpers/entity.py @@ -496,6 +496,9 @@ class Entity( # Entry in the entity registry registry_entry: er.RegistryEntry | None = None + # If the entity is removed from the entity registry + _removed_from_registry: bool = False + # The device entry for this entity device_entry: dr.DeviceEntry | None = None @@ -1361,6 +1364,17 @@ async def __async_remove_impl(self, force_remove: bool) -> None: not force_remove and self.registry_entry and not self.registry_entry.disabled + # Check if entity is still in the entity registry + # by checking self._removed_from_registry + # + # Because self.registry_entry is unset in a task, + # its possible that the entity has been removed but + # the task has not yet been executed. + # + # self._removed_from_registry is set to True in a + # callback which does not have the same issue. + # + and not self._removed_from_registry ): # Set the entity's state will to unavailable + ATTR_RESTORED: True self.registry_entry.write_unavailable_state(self.hass) @@ -1430,10 +1444,23 @@ async def async_internal_will_remove_from_hass(self) -> None: if self.platform: self.hass.data[DATA_ENTITY_SOURCE].pop(self.entity_id) - async def _async_registry_updated( + @callback + def _async_registry_updated( self, event: EventType[er.EventEntityRegistryUpdatedData] ) -> None: """Handle entity registry update.""" + action = event.data["action"] + is_remove = action == "remove" + self._removed_from_registry = is_remove + if action == "update" or is_remove: + self.hass.async_create_task( + self._async_process_registry_update_or_remove(event) + ) + + async def _async_process_registry_update_or_remove( + self, event: EventType[er.EventEntityRegistryUpdatedData] + ) -> None: + """Handle entity registry update or remove.""" data = event.data if data["action"] == "remove": await self.async_removed_from_registry() diff --git a/tests/helpers/test_entity.py b/tests/helpers/test_entity.py index 19600506ae2e88..c21caa83fa103e 100644 --- a/tests/helpers/test_entity.py +++ b/tests/helpers/test_entity.py @@ -2468,3 +2468,44 @@ class MockEntityFeatures(IntFlag): "is using deprecated supported features values which will be removed" not in caplog.text ) + + +async def test_remove_entity_registry( + hass: HomeAssistant, entity_registry: er.EntityRegistry +) -> None: + """Test removing an entity from the registry.""" + result = [] + + entry = entity_registry.async_get_or_create( + "test", "test_platform", "5678", suggested_object_id="test" + ) + assert entry.entity_id == "test.test" + + class MockEntity(entity.Entity): + _attr_unique_id = "5678" + + def __init__(self) -> None: + self.added_calls = [] + self.remove_calls = [] + + async def async_added_to_hass(self): + self.added_calls.append(None) + self.async_on_remove(lambda: result.append(1)) + + async def async_will_remove_from_hass(self): + self.remove_calls.append(None) + + platform = MockEntityPlatform(hass, domain="test") + ent = MockEntity() + await platform.async_add_entities([ent]) + assert hass.states.get("test.test").state == STATE_UNKNOWN + assert len(ent.added_calls) == 1 + + entry = entity_registry.async_remove(entry.entity_id) + await hass.async_block_till_done() + + assert len(result) == 1 + assert len(ent.added_calls) == 1 + assert len(ent.remove_calls) == 1 + + assert hass.states.get("test.test") is None From 40bb0ed52d6cb792e4fc4526ef5d25516d79b901 Mon Sep 17 00:00:00 2001 From: "J. Nick Koston" Date: Mon, 19 Feb 2024 17:47:54 -0600 Subject: [PATCH 2/7] Fix pytest assertions from tessie test common If these tests fail they will not show anything other than AssertionError without the rewrite they are currently failing in https://github.com/home-assistant/core/pull/110978 and I have not been able to find the cause --- tests/components/tessie/__init__.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/components/tessie/__init__.py b/tests/components/tessie/__init__.py index df17fe027d9e80..92781c8c87bdf9 100644 --- a/tests/components/tessie/__init__.py +++ b/tests/components/tessie/__init__.py @@ -1 +1,5 @@ """Tests for the Tessie integration.""" + +import pytest + +pytest.register_assert_rewrite("tests.components.tessie.common") From 3df9c145f91a9dd7f913d7b2d2aa8db06aa2bf29 Mon Sep 17 00:00:00 2001 From: "J. Nick Koston" Date: Mon, 19 Feb 2024 20:34:15 -0600 Subject: [PATCH 3/7] fix snapshots --- tests/components/tessie/snapshots/test_sensor.ambr | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/components/tessie/snapshots/test_sensor.ambr b/tests/components/tessie/snapshots/test_sensor.ambr index df2351cafeb0b9..a928316bed47a3 100644 --- a/tests/components/tessie/snapshots/test_sensor.ambr +++ b/tests/components/tessie/snapshots/test_sensor.ambr @@ -124,6 +124,8 @@ 'hidden_by': None, 'icon': None, 'id': , + 'labels': set({ + }), 'name': None, 'options': dict({ 'sensor': dict({ @@ -178,6 +180,8 @@ 'hidden_by': None, 'icon': None, 'id': , + 'labels': set({ + }), 'name': None, 'options': dict({ 'sensor': dict({ From 0e884d4594061693055ce35bd29173a370a40fb0 Mon Sep 17 00:00:00 2001 From: "J. Nick Koston" Date: Tue, 20 Feb 2024 09:49:56 -0600 Subject: [PATCH 4/7] fix race --- homeassistant/helpers/event.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/homeassistant/helpers/event.py b/homeassistant/helpers/event.py index f9c2e47dc962cf..e7a0624b26ce2c 100644 --- a/homeassistant/helpers/event.py +++ b/homeassistant/helpers/event.py @@ -328,6 +328,7 @@ def _async_track_state_change_event( _async_dispatch_entity_id_event, _async_state_change_filter, action, + False, ) @@ -378,6 +379,7 @@ def _async_track_event( bool, ], action: Callable[[EventType[_TypedDictT]], None], + run_immediately: bool, ) -> CALLBACK_TYPE: """Track an event by a specific key.""" if not keys: @@ -399,6 +401,7 @@ def _async_track_event( event_type, ft.partial(dispatcher_callable, hass, callbacks), event_filter=ft.partial(filter_callable, hass, callbacks), + run_immediately=run_immediately, ) job = HassJob(action, f"track {event_type} event {keys}") @@ -473,6 +476,7 @@ def async_track_entity_registry_updated_event( _async_dispatch_old_entity_id_or_entity_id_event, _async_entity_registry_updated_filter, action, + True, ) @@ -529,6 +533,7 @@ def async_track_device_registry_updated_event( _async_dispatch_device_id_event, _async_device_registry_updated_filter, action, + True, ) @@ -590,6 +595,7 @@ def _async_track_state_added_domain( _async_dispatch_domain_event, _async_domain_added_filter, action, + False, ) @@ -622,6 +628,7 @@ def async_track_state_removed_domain( _async_dispatch_domain_event, _async_domain_removed_filter, action, + False, ) From 9b0b3daf01f2bb940a56b0d1c0539e26dd76af86 Mon Sep 17 00:00:00 2001 From: "J. Nick Koston" Date: Tue, 20 Feb 2024 09:52:11 -0600 Subject: [PATCH 5/7] fix race --- homeassistant/helpers/event.py | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) diff --git a/homeassistant/helpers/event.py b/homeassistant/helpers/event.py index e7a0624b26ce2c..55757c7d1f7f73 100644 --- a/homeassistant/helpers/event.py +++ b/homeassistant/helpers/event.py @@ -390,10 +390,8 @@ def _async_track_event( hass_data = hass.data - callbacks: dict[ - str, list[HassJob[[EventType[_TypedDictT]], Any]] - ] | None = hass_data.get(callbacks_key) - if not callbacks: + callbacks: dict[str, list[HassJob[[EventType[_TypedDictT]], Any]]] | None + if not (callbacks := hass_data.get(callbacks_key)): callbacks = hass_data[callbacks_key] = {} if listeners_key not in hass_data: @@ -407,8 +405,7 @@ def _async_track_event( job = HassJob(action, f"track {event_type} event {keys}") for key in keys: - callback_list = callbacks.get(key) - if callback_list: + if callback_list := callbacks.get(key): callback_list.append(job) else: callbacks[key] = [job] From a9280417947112693169a493626f87f7d688e33f Mon Sep 17 00:00:00 2001 From: "J. Nick Koston" Date: Tue, 20 Feb 2024 09:54:10 -0600 Subject: [PATCH 6/7] fully fixed --- tests/helpers/test_entity.py | 47 ++++++++++++++++++++++++++++++++++++ 1 file changed, 47 insertions(+) diff --git a/tests/helpers/test_entity.py b/tests/helpers/test_entity.py index c21caa83fa103e..4de38cc814dd9f 100644 --- a/tests/helpers/test_entity.py +++ b/tests/helpers/test_entity.py @@ -2509,3 +2509,50 @@ async def async_will_remove_from_hass(self): assert len(ent.remove_calls) == 1 assert hass.states.get("test.test") is None + + +async def test_reset_right_after_remove_entity_registry( + hass: HomeAssistant, entity_registry: er.EntityRegistry +) -> None: + """Test resetting the platform right after removing an entity from the registry. + + A reset commonly happens during a reload. + """ + result = [] + + entry = entity_registry.async_get_or_create( + "test", "test_platform", "5678", suggested_object_id="test" + ) + assert entry.entity_id == "test.test" + + class MockEntity(entity.Entity): + _attr_unique_id = "5678" + + def __init__(self) -> None: + self.added_calls = [] + self.remove_calls = [] + + async def async_added_to_hass(self): + self.added_calls.append(None) + self.async_on_remove(lambda: result.append(1)) + + async def async_will_remove_from_hass(self): + self.remove_calls.append(None) + + platform = MockEntityPlatform(hass, domain="test") + ent = MockEntity() + await platform.async_add_entities([ent]) + assert hass.states.get("test.test").state == STATE_UNKNOWN + assert len(ent.added_calls) == 1 + + entry = entity_registry.async_remove(entry.entity_id) + + # Reset the platform immediately after removing the entity from the registry + await platform.async_reset() + await hass.async_block_till_done() + + assert len(result) == 1 + assert len(ent.added_calls) == 1 + assert len(ent.remove_calls) == 1 + + assert hass.states.get("test.test") is None From aea1648db52c1206f32111536755fac3a81a11e8 Mon Sep 17 00:00:00 2001 From: "J. Nick Koston" Date: Tue, 20 Feb 2024 20:07:13 -0600 Subject: [PATCH 7/7] tweak docstring --- homeassistant/helpers/event.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/homeassistant/helpers/event.py b/homeassistant/helpers/event.py index 55757c7d1f7f73..35493dbcacb0df 100644 --- a/homeassistant/helpers/event.py +++ b/homeassistant/helpers/event.py @@ -381,7 +381,14 @@ def _async_track_event( action: Callable[[EventType[_TypedDictT]], None], run_immediately: bool, ) -> CALLBACK_TYPE: - """Track an event by a specific key.""" + """Track an event by a specific key. + + This function is intended for internal use only. + + The dispatcher_callable, filter_callable, event_type, and run_immediately + must always be the same for the listener_key as the first call to this + function will set the listener_key in hass.data. + """ if not keys: return _remove_empty_listener