Base Script on entity

homeassistant/components/script.py

@@ -8,154 +8,198 @@ by the user or automatically based upon automation events, etc.
 import logging
 from datetime import timedelta
 import homeassistant.util.dt as date_util
+from itertools import islice
 import threading
 
-from homeassistant.helpers.event import track_point_in_time
+from homeassistant.helpers.entity_component import EntityComponent
+from homeassistant.helpers.entity import ToggleEntity
+from homeassistant.helpers.event import track_point_in_utc_time
 from homeassistant.util import split_entity_id
 from homeassistant.const import (
-    STATE_ON, STATE_OFF, SERVICE_TURN_ON, SERVICE_TURN_OFF, EVENT_TIME_CHANGED)
+    ATTR_ENTITY_ID, EVENT_TIME_CHANGED, STATE_ON, SERVICE_TURN_ON,
+    SERVICE_TURN_OFF)
 
 DOMAIN = "script"
+ENTITY_ID_FORMAT = DOMAIN + '.{}'
 DEPENDENCIES = ["group"]
 
+STATE_NOT_RUNNING = 'Not Running'
+
 CONF_ALIAS = "alias"
-CONF_SERVICE = "execute_service"
+CONF_SERVICE = "service"
+CONF_SERVICE_OLD = "execute_service"
 CONF_SERVICE_DATA = "service_data"
 CONF_SEQUENCE = "sequence"
 CONF_EVENT = "event"
 CONF_EVENT_DATA = "event_data"
 CONF_DELAY = "delay"
-ATTR_ENTITY_ID = "entity_id"
+
+ATTR_LAST_ACTION = 'last_action'
 
 _LOGGER = logging.getLogger(__name__)
 
 
+def is_on(hass, entity_id):
+    """ Returns if the switch is on based on the statemachine. """
+    return hass.states.is_state(entity_id, STATE_ON)
+
+
+def turn_on(hass, entity_id):
+    """ Turn script on. """
+    _, object_id = split_entity_id(entity_id)
+
+    hass.services.call(DOMAIN, object_id)
+
+
+def turn_off(hass, entity_id):
+    """ Turn script on. """
+    hass.services.call(DOMAIN, SERVICE_TURN_OFF, {ATTR_ENTITY_ID: entity_id})
+
+
 def setup(hass, config):
     """ Load the scripts from the configuration. """
 
-    scripts = []
+    component = EntityComponent(_LOGGER, DOMAIN, hass)
+
+    def service_handler(service):
+        """ Execute a service call to script.<script name>. """
+        entity_id = ENTITY_ID_FORMAT.format(service.service)
+        script = component.entities.get(entity_id)
+        if script:
+            script.turn_on()
+
     for name, cfg in config[DOMAIN].items():
-        if CONF_SEQUENCE not in cfg:
+        if not cfg.get(CONF_SEQUENCE):
             _LOGGER.warn("Missing key 'sequence' for script %s", name)
             continue
         alias = cfg.get(CONF_ALIAS, name)
-        entity_id = "{}.{}".format(DOMAIN, name)
+        script = Script(hass, alias, cfg[CONF_SEQUENCE])
-        script = Script(hass, entity_id, alias, cfg[CONF_SEQUENCE])
+        component.add_entities((script,))
-        hass.services.register(DOMAIN, name, script)
+        _, object_id = split_entity_id(script.entity_id)
-        scripts.append(script)
+        hass.services.register(DOMAIN, object_id, service_handler)
 
-    def _get_entities(service):
+    def turn_on_service(service):
-        """  Make sure that we always get a list of entities """
+        """ Calls a service to turn script on. """
-        if isinstance(service.data[ATTR_ENTITY_ID], list):
+        # We could turn on script directly here, but we only want to offer
-            return service.data[ATTR_ENTITY_ID]
+        # one way to do it. Otherwise no easy way to call invocations.
-        else:
+        for script in component.extract_from_service(service):
-            return [service.data[ATTR_ENTITY_ID]]
+            turn_on(hass, script.entity_id)
 
-    def turn_on(service):
+    def turn_off_service(service):
-        """ Calls a script. """
-        for entity_id in _get_entities(service):
-            domain, service = split_entity_id(entity_id)
-            hass.services.call(domain, service, {})
-
-    def turn_off(service):
         """ Cancels a script. """
-        for entity_id in _get_entities(service):
+        for script in component.extract_from_service(service):
-            for script in scripts:
+            script.turn_off()
-                if script.entity_id == entity_id:
-                    script.cancel()
 
-    hass.services.register(DOMAIN, SERVICE_TURN_ON, turn_on)
+    hass.services.register(DOMAIN, SERVICE_TURN_ON, turn_on_service)
-    hass.services.register(DOMAIN, SERVICE_TURN_OFF, turn_off)
+    hass.services.register(DOMAIN, SERVICE_TURN_OFF, turn_off_service)
 
     return True
 
 
-class Script(object):
+class Script(ToggleEntity):
-    # pylint: disable=attribute-defined-outside-init
+    """ Represents a script. """
-    # pylint: disable=too-many-instance-attributes
+    def __init__(self, hass, name, sequence):
-    # pylint: disable=too-few-public-methods
-    """
-    A script contains a sequence of service calls or configured delays
-    that are executed in order.
-    Each script also has a state (on/off) indicating whether the script is
-    running or not.
-    """
-    def __init__(self, hass, entity_id, alias, sequence):
         self.hass = hass
-        self.alias = alias
+        self._name = name
         self.sequence = sequence
-        self.entity_id = entity_id
         self._lock = threading.Lock()
-        self._reset()
+        self._cur = -1
+        self._last_action = None
+        self._listener = None
 
-    def cancel(self):
+    @property
-        """ Cancels a running script and resets the state back to off. """
+    def should_poll(self):
-        _LOGGER.info("Cancelled script %s", self.alias)
+        return False
+
+    @property
+    def name(self):
+        """ Returns the name of the entity. """
+        return self._name
+
+    @property
+    def state_attributes(self):
+        """ Returns the state attributes. """
+        attrs = {}
+
+        if self._last_action:
+            attrs[ATTR_LAST_ACTION] = self._last_action
+
+        return attrs
+
+    @property
+    def is_on(self):
+        """ True if entity is on. """
+        return self._cur != -1
+
+    def turn_on(self, **kwargs):
+        """ Turn the entity on. """
+        _LOGGER.info("Executing script %s", self._name)
         with self._lock:
-            if self.listener:
+            if self._cur == -1:
-                self.hass.bus.remove_listener(EVENT_TIME_CHANGED,
+                self._cur = 0
-                                              self.listener)
-                self.listener = None
-            self._reset()
 
-    def _reset(self):
+            # Unregister callback if we were in a delay but turn on is called
-        """ Resets a script back to default state so that it is ready to
+            # again. In that case we just continue execution.
-            run from the start again. """
+            self._remove_listener()
-        self.actions = None
-        self.listener = None
-        self.last_action = "Not Running"
-        self.hass.states.set(self.entity_id, STATE_OFF, {
-            "friendly_name": self.alias,
-            "last_action": self.last_action
-        })
 
-    def _execute_until_done(self):
+            for cur, action in islice(enumerate(self.sequence), self._cur,
-        """ Executes a sequence of actions until finished or until a delay
+                                      None):
-            is encountered.  If a delay action is encountered, the script
-            registers itself to be called again in the future, when
-            _execute_until_done will resume.
 
-            Returns True if finished, False otherwise. """
+                if CONF_SERVICE in action or CONF_SERVICE_OLD in action:
-        for action in self.actions:
-            if CONF_SERVICE in action:
                     self._call_service(action)
+
                 elif CONF_EVENT in action:
                     self._fire_event(action)
+
                 elif CONF_DELAY in action:
+                    # Call ourselves in the future to continue work
+                    def script_delay(now):
+                        """ Called after delay is done. """
+                        self._listener = None
+                        self.turn_on()
+
                     delay = timedelta(**action[CONF_DELAY])
-                point_in_time = date_util.now() + delay
+                    self._listener = track_point_in_utc_time(
-                self.listener = track_point_in_time(
+                        self.hass, script_delay, date_util.utcnow() + delay)
-                    self.hass, self, point_in_time)
+                    self._cur = cur + 1
-                return False
+                    self.update_ha_state()
-        return True
+                    return
 
-    def __call__(self, *args, **kwargs):
+            self._cur = -1
-        """ Executes the script. """
+            self._last_action = None
-        _LOGGER.info("Executing script %s", self.alias)
+            self.update_ha_state()
+
+    def turn_off(self, **kwargs):
+        """ Turn script off. """
+        _LOGGER.info("Cancelled script %s", self._name)
         with self._lock:
-            if self.actions is None:
+            if self._cur == -1:
-                self.actions = (action for action in self.sequence)
+                return
 
-            if not self._execute_until_done():
+            self._cur = -1
-                state = self.hass.states.get(self.entity_id)
+            self.update_ha_state()
-                state.attributes['last_action'] = self.last_action
+            self._remove_listener()
-                self.hass.states.set(self.entity_id, STATE_ON,
-                                     state.attributes)
-            else:
-                self._reset()
 
     def _call_service(self, action):
         """ Calls the service specified in the action. """
-        self.last_action = action.get(CONF_ALIAS, action[CONF_SERVICE])
+        conf_service = action.get(CONF_SERVICE, action.get(CONF_SERVICE_OLD))
-        _LOGGER.info("Executing script %s step %s", self.alias,
+        self._last_action = action.get(CONF_ALIAS, conf_service)
-                     self.last_action)
+        _LOGGER.info("Executing script %s step %s", self._name,
-        domain, service = split_entity_id(action[CONF_SERVICE])
+                     self._last_action)
+        domain, service = split_entity_id(conf_service)
         data = action.get(CONF_SERVICE_DATA, {})
         self.hass.services.call(domain, service, data)
 
     def _fire_event(self, action):
         """ Fires an event. """
-        self.last_action = action.get(CONF_ALIAS, action[CONF_EVENT])
+        self._last_action = action.get(CONF_ALIAS, action[CONF_EVENT])
-        _LOGGER.info("Executing script %s step %s", self.alias,
+        _LOGGER.info("Executing script %s step %s", self._name,
-                     self.last_action)
+                     self._last_action)
         self.hass.bus.fire(action[CONF_EVENT], action.get(CONF_EVENT_DATA))
+
+    def _remove_listener(self):
+        """ Remove point in time listener, if any. """
+        if self._listener:
+            self.hass.bus.remove_listener(EVENT_TIME_CHANGED,
+                                          self._listener)
+            self._listener = None

tests/components/test_script.py

@@ -0,0 +1,217 @@
+"""
+tests.components.test_script
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+Tests script component.
+"""
+# pylint: disable=too-many-public-methods,protected-access
+from datetime import timedelta
+import unittest
+
+from homeassistant.components import script
+import homeassistant.util.dt as dt_util
+
+from tests.common import fire_time_changed, get_test_home_assistant
+
+
+ENTITY_ID = 'script.test'
+
+
+class TestScript(unittest.TestCase):
+    """ Test the switch module. """
+
+    def setUp(self):  # pylint: disable=invalid-name
+        self.hass = get_test_home_assistant()
+
+    def tearDown(self):  # pylint: disable=invalid-name
+        """ Stop down stuff we started. """
+        self.hass.stop()
+
+    def test_setup_with_empty_sequence(self):
+        self.assertTrue(script.setup(self.hass, {
+            'script': {
+                'test': {
+                    'sequence': []
+                }
+            }
+        }))
+
+        self.assertIsNone(self.hass.states.get(ENTITY_ID))
+
+    def test_setup_with_missing_sequence(self):
+        self.assertTrue(script.setup(self.hass, {
+            'script': {
+                'test': {
+                    'sequence': []
+                }
+            }
+        }))
+
+        self.assertIsNone(self.hass.states.get(ENTITY_ID))
+
+    def test_firing_event(self):
+        event = 'test_event'
+        calls = []
+
+        def record_event(event):
+            calls.append(event)
+
+        self.hass.bus.listen(event, record_event)
+
+        self.assertTrue(script.setup(self.hass, {
+            'script': {
+                'test': {
+                    'sequence': [{
+                        'event': event,
+                        'event_data': {
+                            'hello': 'world'
+                        }
+                    }]
+                }
+            }
+        }))
+
+        script.turn_on(self.hass, ENTITY_ID)
+        self.hass.pool.block_till_done()
+
+        self.assertEqual(1, len(calls))
+        self.assertEqual('world', calls[0].data.get('hello'))
+
+    def test_calling_service_old(self):
+        calls = []
+
+        def record_call(service):
+            calls.append(service)
+
+        self.hass.services.register('test', 'script', record_call)
+
+        self.assertTrue(script.setup(self.hass, {
+            'script': {
+                'test': {
+                    'sequence': [{
+                        'execute_service': 'test.script',
+                        'service_data': {
+                            'hello': 'world'
+                        }
+                    }]
+                }
+            }
+        }))
+
+        script.turn_on(self.hass, ENTITY_ID)
+        self.hass.pool.block_till_done()
+
+        self.assertEqual(1, len(calls))
+        self.assertEqual('world', calls[0].data.get('hello'))
+
+    def test_calling_service(self):
+        calls = []
+
+        def record_call(service):
+            calls.append(service)
+
+        self.hass.services.register('test', 'script', record_call)
+
+        self.assertTrue(script.setup(self.hass, {
+            'script': {
+                'test': {
+                    'sequence': [{
+                        'service': 'test.script',
+                        'service_data': {
+                            'hello': 'world'
+                        }
+                    }]
+                }
+            }
+        }))
+
+        script.turn_on(self.hass, ENTITY_ID)
+        self.hass.pool.block_till_done()
+
+        self.assertEqual(1, len(calls))
+        self.assertEqual('world', calls[0].data.get('hello'))
+
+    def test_delay(self):
+        event = 'test_event'
+        calls = []
+
+        def record_event(event):
+            calls.append(event)
+
+        self.hass.bus.listen(event, record_event)
+
+        self.assertTrue(script.setup(self.hass, {
+            'script': {
+                'test': {
+                    'sequence': [{
+                        'event': event
+                    }, {
+                        'delay': {
+                            'seconds': 5
+                        }
+                    }, {
+                        'event': event,
+                    }]
+                }
+            }
+        }))
+
+        script.turn_on(self.hass, ENTITY_ID)
+        self.hass.pool.block_till_done()
+
+        self.assertTrue(script.is_on(self.hass, ENTITY_ID))
+
+        self.assertEqual(
+            event,
+            self.hass.states.get(ENTITY_ID).attributes.get('last_action'))
+        self.assertEqual(1, len(calls))
+
+        future = dt_util.utcnow() + timedelta(seconds=5)
+        fire_time_changed(self.hass, future)
+        self.hass.pool.block_till_done()
+
+        self.assertFalse(script.is_on(self.hass, ENTITY_ID))
+
+        self.assertEqual(2, len(calls))
+
+    def test_cancel_while_delay(self):
+        event = 'test_event'
+        calls = []
+
+        def record_event(event):
+            calls.append(event)
+
+        self.hass.bus.listen(event, record_event)
+
+        self.assertTrue(script.setup(self.hass, {
+            'script': {
+                'test': {
+                    'sequence': [{
+                        'delay': {
+                            'seconds': 5
+                        }
+                    }, {
+                        'event': event,
+                    }]
+                }
+            }
+        }))
+
+        script.turn_on(self.hass, ENTITY_ID)
+        self.hass.pool.block_till_done()
+
+        self.assertTrue(script.is_on(self.hass, ENTITY_ID))
+
+        self.assertEqual(0, len(calls))
+
+        script.turn_off(self.hass, ENTITY_ID)
+        self.hass.pool.block_till_done()
+        self.assertFalse(script.is_on(self.hass, ENTITY_ID))
+
+        future = dt_util.utcnow() + timedelta(seconds=5)
+        fire_time_changed(self.hass, future)
+        self.hass.pool.block_till_done()
+
+        self.assertFalse(script.is_on(self.hass, ENTITY_ID))
+
+        self.assertEqual(0, len(calls))