From a10cbadb57157f49752b79d4433322e7fd8e77dc Mon Sep 17 00:00:00 2001
From: Adam Mills <adam@armills.info>
Date: Sun, 2 Dec 2018 04:51:15 -0500
Subject: [PATCH] Restore states when removing/adding entities (#18890)

---
 homeassistant/helpers/restore_state.py | 14 ++++++++++----
 tests/helpers/test_restore_state.py    | 20 ++++++++++++++++++++
 2 files changed, 30 insertions(+), 4 deletions(-)

diff --git a/homeassistant/helpers/restore_state.py b/homeassistant/helpers/restore_state.py
index cabaf64d859..33b612b555a 100644
--- a/homeassistant/helpers/restore_state.py
+++ b/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."""
diff --git a/tests/helpers/test_restore_state.py b/tests/helpers/test_restore_state.py
index e6693d2cf61..b13bc87421b 100644
--- a/tests/helpers/test_restore_state.py
+++ b/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'