From 2cdaf632a49d6b7ee2dbeaad9cfe88abdabadfbb Mon Sep 17 00:00:00 2001
From: Franck Nijhof <git@frenck.dev>
Date: Mon, 23 Aug 2021 20:05:29 +0200
Subject: [PATCH] Restore last_triggered state in scripts (#55071)

---
 homeassistant/components/script/__init__.py | 10 ++++-
 tests/components/script/test_init.py        | 50 ++++++++++++++++++++-
 2 files changed, 57 insertions(+), 3 deletions(-)

diff --git a/homeassistant/components/script/__init__.py b/homeassistant/components/script/__init__.py
index 483b4065be2..f3f34a0ad53 100644
--- a/homeassistant/components/script/__init__.py
+++ b/homeassistant/components/script/__init__.py
@@ -32,6 +32,7 @@ import homeassistant.helpers.config_validation as cv
 from homeassistant.helpers.config_validation import make_entity_service_schema
 from homeassistant.helpers.entity import ToggleEntity
 from homeassistant.helpers.entity_component import EntityComponent
+from homeassistant.helpers.restore_state import RestoreEntity
 from homeassistant.helpers.script import (
     ATTR_CUR,
     ATTR_MAX,
@@ -42,6 +43,7 @@ from homeassistant.helpers.script import (
 from homeassistant.helpers.service import async_set_service_schema
 from homeassistant.helpers.trace import trace_get, trace_path
 from homeassistant.loader import bind_hass
+from homeassistant.util.dt import parse_datetime
 
 from .config import ScriptConfig, async_validate_config_item
 from .const import (
@@ -296,7 +298,7 @@ async def _async_process_config(hass, config, component) -> bool:
     return blueprints_used
 
 
-class ScriptEntity(ToggleEntity):
+class ScriptEntity(ToggleEntity, RestoreEntity):
     """Representation of a script entity."""
 
     icon = None
@@ -415,6 +417,12 @@ class ScriptEntity(ToggleEntity):
         """
         await self.script.async_stop()
 
+    async def async_added_to_hass(self) -> None:
+        """Restore last triggered on startup."""
+        if state := await self.async_get_last_state():
+            if last_triggered := state.attributes.get("last_triggered"):
+                self.script.last_triggered = parse_datetime(last_triggered)
+
     async def async_will_remove_from_hass(self):
         """Stop script and remove service when it will be removed from Home Assistant."""
         await self.script.async_stop()
diff --git a/tests/components/script/test_init.py b/tests/components/script/test_init.py
index 6070daeb8af..9190b033f44 100644
--- a/tests/components/script/test_init.py
+++ b/tests/components/script/test_init.py
@@ -15,16 +15,25 @@ from homeassistant.const import (
     SERVICE_TOGGLE,
     SERVICE_TURN_OFF,
     SERVICE_TURN_ON,
+    STATE_OFF,
+)
+from homeassistant.core import (
+    Context,
+    CoreState,
+    HomeAssistant,
+    State,
+    callback,
+    split_entity_id,
 )
-from homeassistant.core import Context, callback, split_entity_id
 from homeassistant.exceptions import ServiceNotFound
 from homeassistant.helpers import template
 from homeassistant.helpers.event import async_track_state_change
 from homeassistant.helpers.service import async_get_all_descriptions
 from homeassistant.loader import bind_hass
 from homeassistant.setup import async_setup_component, setup_component
+import homeassistant.util.dt as dt_util
 
-from tests.common import async_mock_service, get_test_home_assistant
+from tests.common import async_mock_service, get_test_home_assistant, mock_restore_cache
 from tests.components.logbook.test_init import MockLazyEventPartialState
 
 ENTITY_ID = "script.test"
@@ -84,6 +93,7 @@ class TestScriptComponent(unittest.TestCase):
 
     def test_passing_variables(self):
         """Test different ways of passing in variables."""
+        mock_restore_cache(self.hass, ())
         calls = []
         context = Context()
 
@@ -796,3 +806,39 @@ async def test_script_this_var_always(hass, caplog):
     # Verify this available to all templates
     assert mock_calls[0].data.get("this_template") == "script.script1"
     assert "Error rendering variables" not in caplog.text
+
+
+async def test_script_restore_last_triggered(hass: HomeAssistant) -> None:
+    """Test if last triggered is restored on start."""
+    time = dt_util.utcnow()
+    mock_restore_cache(
+        hass,
+        (
+            State("script.no_last_triggered", STATE_OFF),
+            State("script.last_triggered", STATE_OFF, {"last_triggered": time}),
+        ),
+    )
+    hass.state = CoreState.starting
+
+    assert await async_setup_component(
+        hass,
+        "script",
+        {
+            "script": {
+                "no_last_triggered": {
+                    "sequence": [{"delay": {"seconds": 5}}],
+                },
+                "last_triggered": {
+                    "sequence": [{"delay": {"seconds": 5}}],
+                },
+            },
+        },
+    )
+
+    state = hass.states.get("script.no_last_triggered")
+    assert state
+    assert state.attributes["last_triggered"] is None
+
+    state = hass.states.get("script.last_triggered")
+    assert state
+    assert state.attributes["last_triggered"] == time