From 23fb8c4cddeba13841504c84cace1ac0f4ad3b99 Mon Sep 17 00:00:00 2001
From: Paulus Schoutsen <paulus@paulusschoutsen.nl>
Date: Thu, 17 Nov 2016 21:50:01 -0800
Subject: [PATCH] Convert script component to async (#4427)

---
 homeassistant/components/script.py | 58 +++++++++++++++++++-----------
 tests/components/test_script.py    |  5 +++
 2 files changed, 42 insertions(+), 21 deletions(-)

diff --git a/homeassistant/components/script.py b/homeassistant/components/script.py
index 8c521c33856..df46fb5a03d 100644
--- a/homeassistant/components/script.py
+++ b/homeassistant/components/script.py
@@ -7,6 +7,7 @@ by the user or automatically based upon automation events, etc.
 For more details about this component, please refer to the documentation at
 https://home-assistant.io/components/script/
 """
+import asyncio
 import logging
 
 import voluptuous as vol
@@ -72,11 +73,13 @@ def toggle(hass, entity_id):
     hass.services.call(DOMAIN, SERVICE_TOGGLE, {ATTR_ENTITY_ID: entity_id})
 
 
-def setup(hass, config):
+@asyncio.coroutine
+def async_setup(hass, config):
     """Load the scripts from the configuration."""
     component = EntityComponent(_LOGGER, DOMAIN, hass,
                                 group_name=GROUP_NAME_ALL_SCRIPTS)
 
+    @asyncio.coroutine
     def service_handler(service):
         """Execute a service call to script.<script name>."""
         entity_id = ENTITY_ID_FORMAT.format(service.service)
@@ -84,38 +87,48 @@ def setup(hass, config):
         if script.is_on:
             _LOGGER.warning("Script %s already running.", entity_id)
             return
-        script.turn_on(variables=service.data)
+        yield from script.async_turn_on(variables=service.data)
+
+    scripts = []
 
     for object_id, cfg in config[DOMAIN].items():
         alias = cfg.get(CONF_ALIAS, object_id)
         script = ScriptEntity(hass, object_id, alias, cfg[CONF_SEQUENCE])
-        component.add_entities((script,))
-        hass.services.register(DOMAIN, object_id, service_handler,
-                               schema=SCRIPT_SERVICE_SCHEMA)
+        scripts.append(script)
+        hass.services.async_register(DOMAIN, object_id, service_handler,
+                                     schema=SCRIPT_SERVICE_SCHEMA)
 
+    yield from component.async_add_entities(scripts)
+
+    @asyncio.coroutine
     def turn_on_service(service):
         """Call a service to turn script on."""
         # We could turn on script directly here, but we only want to offer
         # one way to do it. Otherwise no easy way to detect invocations.
-        for script in component.extract_from_service(service):
-            turn_on(hass, script.entity_id, service.data.get(ATTR_VARIABLES))
+        var = service.data.get(ATTR_VARIABLES)
+        for script in component.async_extract_from_service(service):
+            yield from hass.services.async_call(DOMAIN, script.object_id, var)
 
+    @asyncio.coroutine
     def turn_off_service(service):
         """Cancel a script."""
-        for script in component.extract_from_service(service):
-            script.turn_off()
+        # Stopping a script is ok to be done in parallel
+        yield from asyncio.wait(
+            [script.async_turn_off() for script
+             in component.async_extract_from_service(service)], loop=hass.loop)
 
+    @asyncio.coroutine
     def toggle_service(service):
         """Toggle a script."""
-        for script in component.extract_from_service(service):
-            script.toggle()
+        for script in component.async_extract_from_service(service):
+            yield from script.async_toggle()
 
-    hass.services.register(DOMAIN, SERVICE_TURN_ON, turn_on_service,
-                           schema=SCRIPT_TURN_ONOFF_SCHEMA)
-    hass.services.register(DOMAIN, SERVICE_TURN_OFF, turn_off_service,
-                           schema=SCRIPT_TURN_ONOFF_SCHEMA)
-    hass.services.register(DOMAIN, SERVICE_TOGGLE, toggle_service,
-                           schema=SCRIPT_TURN_ONOFF_SCHEMA)
+    hass.services.async_register(DOMAIN, SERVICE_TURN_ON, turn_on_service,
+                                 schema=SCRIPT_TURN_ONOFF_SCHEMA)
+    hass.services.async_register(DOMAIN, SERVICE_TURN_OFF, turn_off_service,
+                                 schema=SCRIPT_TURN_ONOFF_SCHEMA)
+    hass.services.async_register(DOMAIN, SERVICE_TOGGLE, toggle_service,
+                                 schema=SCRIPT_TURN_ONOFF_SCHEMA)
     return True
 
 
@@ -124,6 +137,7 @@ class ScriptEntity(ToggleEntity):
 
     def __init__(self, hass, object_id, name, sequence):
         """Initialize the script."""
+        self.object_id = object_id
         self.entity_id = ENTITY_ID_FORMAT.format(object_id)
         self.script = Script(hass, sequence, name, self.async_update_ha_state)
 
@@ -152,10 +166,12 @@ class ScriptEntity(ToggleEntity):
         """Return true if script is on."""
         return self.script.is_running
 
-    def turn_on(self, **kwargs):
+    @asyncio.coroutine
+    def async_turn_on(self, **kwargs):
         """Turn the script on."""
-        self.script.run(kwargs.get(ATTR_VARIABLES))
+        yield from self.script.async_run(kwargs.get(ATTR_VARIABLES))
 
-    def turn_off(self, **kwargs):
+    @asyncio.coroutine
+    def async_turn_off(self, **kwargs):
         """Turn script off."""
-        self.script.stop()
+        self.script.async_stop()
diff --git a/tests/components/test_script.py b/tests/components/test_script.py
index 76df0a24055..438ae75e9c6 100644
--- a/tests/components/test_script.py
+++ b/tests/components/test_script.py
@@ -86,6 +86,11 @@ class TestScriptComponent(unittest.TestCase):
         self.hass.block_till_done()
         self.assertEqual(0, len(events))
 
+        script.turn_off(self.hass, ENTITY_ID)
+        self.hass.block_till_done()
+        self.assertFalse(script.is_on(self.hass, ENTITY_ID))
+        self.assertEqual(0, len(events))
+
         state = self.hass.states.get('group.all_scripts')
         assert state is not None
         assert state.attributes.get('entity_id') == (ENTITY_ID,)