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