Restore states when removing/adding entities (#18890)

homeassistant/helpers/restore_state.py

@@ -165,13 +165,19 @@ class RestoreStateData():
                 self.async_dump_states()))
 
     @callback
-    def async_register_entity(self, entity_id: str) -> None:
+    def async_restore_entity_added(self, entity_id: str) -> None:
         """Store this entity's state when hass is shutdown."""
         self.entity_ids.add(entity_id)
 
     @callback
-    def async_unregister_entity(self, entity_id: str) -> None:
+    def async_restore_entity_removed(self, entity_id: str) -> None:
         """Unregister this entity from saving state."""
+        # When an entity is being removed from hass, store its last state. This
+        # allows us to support state restoration if the entity is removed, then
+        # re-added while hass is still running.
+        self.last_states[entity_id] = StoredState(
+            self.hass.states.get(entity_id), dt_util.utcnow())
+
         self.entity_ids.remove(entity_id)
 
 
@@ -184,7 +190,7 @@ class RestoreEntity(Entity):
             super().async_added_to_hass(),
             RestoreStateData.async_get_instance(self.hass),
         )
-        data.async_register_entity(self.entity_id)
+        data.async_restore_entity_added(self.entity_id)
 
     async def async_will_remove_from_hass(self) -> None:
         """Run when entity will be removed from hass."""
@@ -192,7 +198,7 @@ class RestoreEntity(Entity):
             super().async_will_remove_from_hass(),
             RestoreStateData.async_get_instance(self.hass),
         )
-        data.async_unregister_entity(self.entity_id)
+        data.async_restore_entity_removed(self.entity_id)
 
     async def async_get_last_state(self) -> Optional[State]:
         """Get the entity state from the previous run."""

tests/helpers/test_restore_state.py

@@ -198,3 +198,23 @@ async def test_load_error(hass):
         state = await entity.async_get_last_state()
 
     assert state is None
+
+
+async def test_state_saved_on_remove(hass):
+    """Test that we save entity state on removal."""
+    entity = RestoreEntity()
+    entity.hass = hass
+    entity.entity_id = 'input_boolean.b0'
+    await entity.async_added_to_hass()
+
+    hass.states.async_set('input_boolean.b0', 'on')
+
+    data = await RestoreStateData.async_get_instance(hass)
+
+    # No last states should currently be saved
+    assert not data.last_states
+
+    await entity.async_will_remove_from_hass()
+
+    # We should store the input boolean state when it is removed
+    assert data.last_states['input_boolean.b0'].state.state == 'on'