Add Context support for async_entity removal (#33209)

* Add Context for async_remove

* Check context in state automation on entity removal

homeassistant/components/configurator/__init__.py

@@ -14,7 +14,7 @@ from homeassistant.const import (
     ATTR_FRIENDLY_NAME,
     EVENT_TIME_CHANGED,
 )
-from homeassistant.core import callback as async_callback
+from homeassistant.core import Event, callback as async_callback
 from homeassistant.helpers.entity import async_generate_entity_id
 from homeassistant.loader import bind_hass
 from homeassistant.util.async_ import run_callback_threadsafe
@@ -214,9 +214,9 @@ class Configurator:
         # it shortly after so that it is deleted when the client updates.
         self.hass.states.async_set(entity_id, STATE_CONFIGURED)
 
-        def deferred_remove(event):
+        def deferred_remove(event: Event):
             """Remove the request state."""
-            self.hass.states.async_remove(entity_id)
+            self.hass.states.async_remove(entity_id, context=event.context)
 
         self.hass.bus.async_listen_once(EVENT_TIME_CHANGED, deferred_remove)
 

homeassistant/components/microsoft_face/__init__.py

@@ -113,7 +113,7 @@ async def async_setup(hass, config):
             face.store.pop(g_id)
 
             entity = entities.pop(g_id)
-            hass.states.async_remove(entity.entity_id)
+            hass.states.async_remove(entity.entity_id, service.context)
         except HomeAssistantError as err:
             _LOGGER.error("Can't delete group '%s' with error: %s", g_id, err)
 

homeassistant/components/persistent_notification/__init__.py

@@ -158,7 +158,7 @@ async def async_setup(hass: HomeAssistant, config: dict) -> bool:
         if entity_id not in persistent_notifications:
             return
 
-        hass.states.async_remove(entity_id)
+        hass.states.async_remove(entity_id, call.context)
 
         del persistent_notifications[entity_id]
         hass.bus.async_fire(EVENT_PERSISTENT_NOTIFICATIONS_UPDATED)

homeassistant/core.py

@@ -901,7 +901,7 @@ class StateMachine:
         ).result()
 
     @callback
-    def async_remove(self, entity_id: str) -> bool:
+    def async_remove(self, entity_id: str, context: Optional[Context] = None) -> bool:
         """Remove the state of an entity.
 
         Returns boolean to indicate if an entity was removed.
@@ -917,6 +917,8 @@ class StateMachine:
         self._bus.async_fire(
             EVENT_STATE_CHANGED,
             {"entity_id": entity_id, "old_state": old_state, "new_state": None},
+            EventOrigin.local,
+            context=context,
         )
         return True
 

homeassistant/helpers/entity.py

@@ -504,7 +504,7 @@ class Entity(ABC):
             while self._on_remove:
                 self._on_remove.pop()()
 
-        self.hass.states.async_remove(self.entity_id)
+        self.hass.states.async_remove(self.entity_id, context=self._context)
 
     async def async_added_to_hass(self) -> None:
         """Run when entity about to be added to hass.

homeassistant/helpers/entity_registry.py

@@ -518,7 +518,7 @@ def async_setup_entity_restore(
         if state is None or not state.attributes.get(ATTR_RESTORED):
             return
 
-        hass.states.async_remove(event.data["entity_id"])
+        hass.states.async_remove(event.data["entity_id"], context=event.context)
 
     hass.bus.async_listen(EVENT_ENTITY_REGISTRY_UPDATED, cleanup_restored_states)
 

tests/components/automation/test_state.py

@@ -521,6 +521,7 @@ async def test_if_fires_on_entity_change_with_for(hass, calls):
 
 async def test_if_fires_on_entity_removal(hass, calls):
     """Test for firing on entity removal, when new_state is None."""
+    context = Context()
     hass.states.async_set("test.entity", "hello")
     await hass.async_block_till_done()
 
@@ -536,9 +537,10 @@ async def test_if_fires_on_entity_removal(hass, calls):
     )
     await hass.async_block_till_done()
 
-    assert hass.states.async_remove("test.entity")
+    assert hass.states.async_remove("test.entity", context=context)
     await hass.async_block_till_done()
     assert 1 == len(calls)
+    assert calls[0].context.parent_id == context.id
 
 
 async def test_if_fires_on_for_condition(hass, calls):