From fb820975b534a9d8721d4e1c01d0614ad6ee6559 Mon Sep 17 00:00:00 2001
From: zewelor <zewelor@gmail.com>
Date: Tue, 19 Feb 2019 19:06:40 +0100
Subject: [PATCH] Add yeelight flow action support (#21195)

---
 homeassistant/components/light/services.yaml |  3 +++
 homeassistant/components/light/yeelight.py   | 10 +++++++++-
 2 files changed, 12 insertions(+), 1 deletion(-)

diff --git a/homeassistant/components/light/services.yaml b/homeassistant/components/light/services.yaml
index 9836bf97f90..10cbeb42aa4 100644
--- a/homeassistant/components/light/services.yaml
+++ b/homeassistant/components/light/services.yaml
@@ -202,6 +202,9 @@ yeelight_start_flow:
     count:
       description: The number of times to run this flow (0 to run forever).
       example: 0
+    action:
+      description: The action to take after the flow stops. Can be 'recover', 'stay', 'off'. (default 'recover')
+      example: 'stay'
     transitions:
       description: Array of transitions, for desired effect. Examples https://yeelight.readthedocs.io/en/stable/flow.html
       example: '[{ "TemperatureTransition": [1900, 1000, 80] }, { "TemperatureTransition": [1900, 1000, 10] }]'
diff --git a/homeassistant/components/light/yeelight.py b/homeassistant/components/light/yeelight.py
index b678fcd2799..d22d94f7bbc 100644
--- a/homeassistant/components/light/yeelight.py
+++ b/homeassistant/components/light/yeelight.py
@@ -46,8 +46,13 @@ DATA_KEY = 'light.yeelight'
 
 ATTR_MODE = 'mode'
 ATTR_COUNT = 'count'
+ATTR_ACTION = 'action'
 ATTR_TRANSITIONS = 'transitions'
 
+ACTION_RECOVER = 'recover'
+ACTION_STAY = 'stay'
+ACTION_OFF = 'off'
+
 YEELIGHT_RGB_TRANSITION = 'RGBTransition'
 YEELIGHT_HSV_TRANSACTION = 'HSVTransition'
 YEELIGHT_TEMPERATURE_TRANSACTION = 'TemperatureTransition'
@@ -59,6 +64,8 @@ YEELIGHT_SERVICE_SCHEMA = vol.Schema({
 
 YEELIGHT_FLOW_TRANSITION_SCHEMA = {
     vol.Optional(ATTR_COUNT, default=0): cv.positive_int,
+    vol.Optional(ATTR_ACTION, default=ACTION_RECOVER):
+        vol.Any(ACTION_RECOVER, ACTION_OFF, ACTION_STAY),
     vol.Required(ATTR_TRANSITIONS): [{
         vol.Exclusive(YEELIGHT_RGB_TRANSITION, CONF_TRANSITION):
             vol.All(cv.ensure_list, [cv.positive_int]),
@@ -604,13 +611,14 @@ class YeelightLight(Light):
 
         return transition_objects
 
-    def start_flow(self, transitions, count=0):
+    def start_flow(self, transitions, count=0, action=ACTION_RECOVER):
         """Start flow."""
         import yeelight
 
         try:
             flow = yeelight.Flow(
                 count=count,
+                action=yeelight.Flow.actions[action],
                 transitions=self.transitions_config_parser(transitions))
 
             self._bulb.start_flow(flow)