From 7f08dd851eac812b8cf18a8d4d61d7858617af72 Mon Sep 17 00:00:00 2001
From: Erik Montnemery <erik@montnemery.com>
Date: Wed, 28 Sep 2022 17:05:29 +0200
Subject: [PATCH] Make script entities register their own service (#79202)

---
 homeassistant/components/script/__init__.py | 44 +++++++++------------
 tests/components/script/test_init.py        | 14 +++++++
 2 files changed, 33 insertions(+), 25 deletions(-)

diff --git a/homeassistant/components/script/__init__.py b/homeassistant/components/script/__init__.py
index f72e6dfc6ef..0effdc2754b 100644
--- a/homeassistant/components/script/__init__.py
+++ b/homeassistant/components/script/__init__.py
@@ -271,30 +271,6 @@ async def _async_process_config(hass, config, component) -> bool:
 
     await component.async_add_entities(entities)
 
-    async def service_handler(service: ServiceCall) -> None:
-        """Execute a service call to script.<script name>."""
-        entity_registry = er.async_get(hass)
-        entity_id = entity_registry.async_get_entity_id(DOMAIN, DOMAIN, service.service)
-        script_entity = component.get_entity(entity_id)
-        await script_entity.async_turn_on(
-            variables=service.data, context=service.context
-        )
-
-    # Register services for all entities that were created successfully.
-    for entity in entities:
-        hass.services.async_register(
-            DOMAIN, entity.unique_id, service_handler, schema=SCRIPT_SERVICE_SCHEMA
-        )
-
-        # Register the service description
-        service_desc = {
-            CONF_NAME: entity.name,
-            CONF_DESCRIPTION: entity.description,
-            CONF_FIELDS: entity.fields,
-        }
-        unique_id = cast(str, entity.unique_id)
-        async_set_service_schema(hass, DOMAIN, unique_id, service_desc)
-
     return blueprints_used
 
 
@@ -429,8 +405,26 @@ class ScriptEntity(ToggleEntity, RestoreEntity):
         """
         await self.script.async_stop()
 
+    async def _service_handler(self, service: ServiceCall) -> None:
+        """Execute a service call to script.<script name>."""
+        await self.async_turn_on(variables=service.data, context=service.context)
+
     async def async_added_to_hass(self) -> None:
-        """Restore last triggered on startup."""
+        """Restore last triggered on startup and register service."""
+
+        unique_id = cast(str, self.unique_id)
+        self.hass.services.async_register(
+            DOMAIN, unique_id, self._service_handler, schema=SCRIPT_SERVICE_SCHEMA
+        )
+
+        # Register the service description
+        service_desc = {
+            CONF_NAME: cast(er.RegistryEntry, self.registry_entry).name or self.name,
+            CONF_DESCRIPTION: self.description,
+            CONF_FIELDS: self.fields,
+        }
+        async_set_service_schema(self.hass, DOMAIN, unique_id, service_desc)
+
         if state := await self.async_get_last_state():
             if last_triggered := state.attributes.get("last_triggered"):
                 self.script.last_triggered = parse_datetime(last_triggered)
diff --git a/tests/components/script/test_init.py b/tests/components/script/test_init.py
index cc9c9a22302..b48a65275b7 100644
--- a/tests/components/script/test_init.py
+++ b/tests/components/script/test_init.py
@@ -1084,6 +1084,7 @@ async def test_script_service_changed_entity_id(hass: HomeAssistant) -> None:
 
     hass.services.async_register("test", "script", record_call)
 
+    # Make sure the service of a script with overridden entity_id works
     assert await async_setup_component(
         hass,
         "script",
@@ -1105,3 +1106,16 @@ async def test_script_service_changed_entity_id(hass: HomeAssistant) -> None:
 
     assert len(calls) == 1
     assert calls[0].data["entity_id"] == "script.custom_entity_id"
+
+    # Change entity while the script entity is loaded, and make sure the service still works
+    entry = entity_reg.async_update_entity(
+        entry.entity_id, new_entity_id="script.custom_entity_id_2"
+    )
+    assert entry.entity_id == "script.custom_entity_id_2"
+    await hass.async_block_till_done()
+
+    await hass.services.async_call(DOMAIN, "test", {"greeting": "world"})
+    await hass.async_block_till_done()
+
+    assert len(calls) == 2
+    assert calls[1].data["entity_id"] == "script.custom_entity_id_2"