Optimistic MQTT light (#14401)

* Restores light state, case the light is optimistic

* lint

* hound

* hound

* Added mqtt_json

* hound

* added mqtt_template

* lint

* cleanup

* use ATTR

homeassistant/components/light/mqtt.py

@@ -4,7 +4,6 @@ Support for MQTT lights.
 For more details about this platform, please refer to the documentation at
 https://home-assistant.io/components/light.mqtt/
 """
-import asyncio
 import logging
 
 import voluptuous as vol
@@ -17,12 +16,13 @@ from homeassistant.components.light import (
     SUPPORT_EFFECT, SUPPORT_COLOR, SUPPORT_WHITE_VALUE)
 from homeassistant.const import (
     CONF_BRIGHTNESS, CONF_COLOR_TEMP, CONF_EFFECT, CONF_NAME,
-    CONF_OPTIMISTIC, CONF_PAYLOAD_OFF, CONF_PAYLOAD_ON,
+    CONF_OPTIMISTIC, CONF_PAYLOAD_OFF, CONF_PAYLOAD_ON, STATE_ON,
     CONF_RGB, CONF_STATE, CONF_VALUE_TEMPLATE, CONF_WHITE_VALUE, CONF_XY)
 from homeassistant.components.mqtt import (
     CONF_AVAILABILITY_TOPIC, CONF_COMMAND_TOPIC, CONF_PAYLOAD_AVAILABLE,
     CONF_PAYLOAD_NOT_AVAILABLE, CONF_QOS, CONF_RETAIN, CONF_STATE_TOPIC,
     MqttAvailability)
+from homeassistant.helpers.restore_state import async_get_last_state
 import homeassistant.helpers.config_validation as cv
 import homeassistant.util.color as color_util
 
@@ -100,8 +100,8 @@ PLATFORM_SCHEMA = mqtt.MQTT_RW_PLATFORM_SCHEMA.extend({
 }).extend(mqtt.MQTT_AVAILABILITY_SCHEMA.schema)
 
 
-@asyncio.coroutine
-def async_setup_platform(hass, config, async_add_devices, discovery_info=None):
+async def async_setup_platform(hass, config, async_add_devices,
+                               discovery_info=None):
     """Set up a MQTT Light."""
     if discovery_info is not None:
         config = PLATFORM_SCHEMA(discovery_info)
@@ -213,10 +213,9 @@ class MqttLight(MqttAvailability, Light):
         self._supported_features |= (
             topic[CONF_XY_COMMAND_TOPIC] is not None and SUPPORT_COLOR)
 
-    @asyncio.coroutine
-    def async_added_to_hass(self):
+    async def async_added_to_hass(self):
         """Subscribe to MQTT events."""
-        yield from super().async_added_to_hass()
+        await super().async_added_to_hass()
 
         templates = {}
         for key, tpl in list(self._templates.items()):
@@ -226,6 +225,8 @@ class MqttLight(MqttAvailability, Light):
                 tpl.hass = self.hass
                 templates[key] = tpl.async_render_with_possible_json_value
 
+        last_state = await async_get_last_state(self.hass, self.entity_id)
+
         @callback
         def state_received(topic, payload, qos):
             """Handle new MQTT messages."""
@@ -237,9 +238,11 @@ class MqttLight(MqttAvailability, Light):
             self.async_schedule_update_ha_state()
 
         if self._topic[CONF_STATE_TOPIC] is not None:
-            yield from mqtt.async_subscribe(
+            await mqtt.async_subscribe(
                 self.hass, self._topic[CONF_STATE_TOPIC], state_received,
                 self._qos)
+        elif self._optimistic and last_state:
+            self._state = last_state.state == STATE_ON
 
         @callback
         def brightness_received(topic, payload, qos):
@@ -250,10 +253,13 @@ class MqttLight(MqttAvailability, Light):
             self.async_schedule_update_ha_state()
 
         if self._topic[CONF_BRIGHTNESS_STATE_TOPIC] is not None:
-            yield from mqtt.async_subscribe(
+            await mqtt.async_subscribe(
                 self.hass, self._topic[CONF_BRIGHTNESS_STATE_TOPIC],
                 brightness_received, self._qos)
             self._brightness = 255
+        elif self._optimistic_brightness and last_state\
+                and last_state.attributes.get(ATTR_BRIGHTNESS):
+            self._brightness = last_state.attributes.get(ATTR_BRIGHTNESS)
         elif self._topic[CONF_BRIGHTNESS_COMMAND_TOPIC] is not None:
             self._brightness = 255
         else:
@@ -268,11 +274,14 @@ class MqttLight(MqttAvailability, Light):
             self.async_schedule_update_ha_state()
 
         if self._topic[CONF_RGB_STATE_TOPIC] is not None:
-            yield from mqtt.async_subscribe(
+            await mqtt.async_subscribe(
                 self.hass, self._topic[CONF_RGB_STATE_TOPIC], rgb_received,
                 self._qos)
             self._hs = (0, 0)
-        if self._topic[CONF_RGB_COMMAND_TOPIC] is not None:
+        if self._optimistic_rgb and last_state\
+                and last_state.attributes.get(ATTR_HS_COLOR):
+            self._hs = last_state.attributes.get(ATTR_HS_COLOR)
+        elif self._topic[CONF_RGB_COMMAND_TOPIC] is not None:
             self._hs = (0, 0)
 
         @callback
@@ -282,11 +291,14 @@ class MqttLight(MqttAvailability, Light):
             self.async_schedule_update_ha_state()
 
         if self._topic[CONF_COLOR_TEMP_STATE_TOPIC] is not None:
-            yield from mqtt.async_subscribe(
+            await mqtt.async_subscribe(
                 self.hass, self._topic[CONF_COLOR_TEMP_STATE_TOPIC],
                 color_temp_received, self._qos)
             self._color_temp = 150
-        if self._topic[CONF_COLOR_TEMP_COMMAND_TOPIC] is not None:
+        if self._optimistic_color_temp and last_state\
+                and last_state.attributes.get(ATTR_COLOR_TEMP):
+            self._color_temp = last_state.attributes.get(ATTR_COLOR_TEMP)
+        elif self._topic[CONF_COLOR_TEMP_COMMAND_TOPIC] is not None:
             self._color_temp = 150
         else:
             self._color_temp = None
@@ -298,11 +310,14 @@ class MqttLight(MqttAvailability, Light):
             self.async_schedule_update_ha_state()
 
         if self._topic[CONF_EFFECT_STATE_TOPIC] is not None:
-            yield from mqtt.async_subscribe(
+            await mqtt.async_subscribe(
                 self.hass, self._topic[CONF_EFFECT_STATE_TOPIC],
                 effect_received, self._qos)
             self._effect = 'none'
-        if self._topic[CONF_EFFECT_COMMAND_TOPIC] is not None:
+        if self._optimistic_effect and last_state\
+                and last_state.attributes.get(ATTR_EFFECT):
+            self._effect = last_state.attributes.get(ATTR_EFFECT)
+        elif self._topic[CONF_EFFECT_COMMAND_TOPIC] is not None:
             self._effect = 'none'
         else:
             self._effect = None
@@ -316,10 +331,13 @@ class MqttLight(MqttAvailability, Light):
             self.async_schedule_update_ha_state()
 
         if self._topic[CONF_WHITE_VALUE_STATE_TOPIC] is not None:
-            yield from mqtt.async_subscribe(
+            await mqtt.async_subscribe(
                 self.hass, self._topic[CONF_WHITE_VALUE_STATE_TOPIC],
                 white_value_received, self._qos)
             self._white_value = 255
+        elif self._optimistic_white_value and last_state\
+                and last_state.attributes.get(ATTR_WHITE_VALUE):
+            self._white_value = last_state.attributes.get(ATTR_WHITE_VALUE)
         elif self._topic[CONF_WHITE_VALUE_COMMAND_TOPIC] is not None:
             self._white_value = 255
         else:
@@ -334,11 +352,14 @@ class MqttLight(MqttAvailability, Light):
             self.async_schedule_update_ha_state()
 
         if self._topic[CONF_XY_STATE_TOPIC] is not None:
-            yield from mqtt.async_subscribe(
+            await mqtt.async_subscribe(
                 self.hass, self._topic[CONF_XY_STATE_TOPIC], xy_received,
                 self._qos)
             self._hs = (0, 0)
-        if self._topic[CONF_XY_COMMAND_TOPIC] is not None:
+        if self._optimistic_xy and last_state\
+                and last_state.attributes.get(ATTR_HS_COLOR):
+            self._hs = last_state.attributes.get(ATTR_HS_COLOR)
+        elif self._topic[CONF_XY_COMMAND_TOPIC] is not None:
             self._hs = (0, 0)
 
     @property
@@ -396,8 +417,7 @@ class MqttLight(MqttAvailability, Light):
         """Flag supported features."""
         return self._supported_features
 
-    @asyncio.coroutine
-    def async_turn_on(self, **kwargs):
+    async def async_turn_on(self, **kwargs):
         """Turn the device on.
 
         This method is a coroutine.
@@ -517,8 +537,7 @@ class MqttLight(MqttAvailability, Light):
         if should_update:
             self.async_schedule_update_ha_state()
 
-    @asyncio.coroutine
-    def async_turn_off(self, **kwargs):
+    async def async_turn_off(self, **kwargs):
         """Turn the device off.
 
         This method is a coroutine.

homeassistant/components/light/mqtt_json.py

@@ -18,7 +18,7 @@ from homeassistant.components.light import (
     SUPPORT_TRANSITION, SUPPORT_WHITE_VALUE)
 from homeassistant.components.light.mqtt import CONF_BRIGHTNESS_SCALE
 from homeassistant.const import (
-    CONF_BRIGHTNESS, CONF_COLOR_TEMP, CONF_EFFECT,
+    CONF_BRIGHTNESS, CONF_COLOR_TEMP, CONF_EFFECT, STATE_ON,
     CONF_NAME, CONF_OPTIMISTIC, CONF_RGB, CONF_WHITE_VALUE, CONF_XY)
 from homeassistant.components.mqtt import (
     CONF_AVAILABILITY_TOPIC, CONF_STATE_TOPIC, CONF_COMMAND_TOPIC,
@@ -26,6 +26,7 @@ from homeassistant.components.mqtt import (
     MqttAvailability)
 import homeassistant.helpers.config_validation as cv
 from homeassistant.helpers.typing import HomeAssistantType, ConfigType
+from homeassistant.helpers.restore_state import async_get_last_state
 import homeassistant.util.color as color_util
 
 _LOGGER = logging.getLogger(__name__)
@@ -177,6 +178,8 @@ class MqttJson(MqttAvailability, Light):
         """Subscribe to MQTT events."""
         await super().async_added_to_hass()
 
+        last_state = await async_get_last_state(self.hass, self.entity_id)
+
         @callback
         def state_received(topic, payload, qos):
             """Handle new MQTT messages."""
@@ -260,6 +263,19 @@ class MqttJson(MqttAvailability, Light):
                 self.hass, self._topic[CONF_STATE_TOPIC], state_received,
                 self._qos)
 
+        if self._optimistic and last_state:
+            self._state = last_state.state == STATE_ON
+            if last_state.attributes.get(ATTR_BRIGHTNESS):
+                self._brightness = last_state.attributes.get(ATTR_BRIGHTNESS)
+            if last_state.attributes.get(ATTR_HS_COLOR):
+                self._hs = last_state.attributes.get(ATTR_HS_COLOR)
+            if last_state.attributes.get(ATTR_COLOR_TEMP):
+                self._color_temp = last_state.attributes.get(ATTR_COLOR_TEMP)
+            if last_state.attributes.get(ATTR_EFFECT):
+                self._effect = last_state.attributes.get(ATTR_EFFECT)
+            if last_state.attributes.get(ATTR_WHITE_VALUE):
+                self._white_value = last_state.attributes.get(ATTR_WHITE_VALUE)
+
     @property
     def brightness(self):
         """Return the brightness of this light between 0..255."""

homeassistant/components/light/mqtt_template.py

@@ -4,7 +4,6 @@ Support for MQTT Template lights.
 For more details about this platform, please refer to the documentation at
 https://home-assistant.io/components/light.mqtt_template/
 """
-import asyncio
 import logging
 import voluptuous as vol
 
@@ -22,6 +21,7 @@ from homeassistant.components.mqtt import (
     MqttAvailability)
 import homeassistant.helpers.config_validation as cv
 import homeassistant.util.color as color_util
+from homeassistant.helpers.restore_state import async_get_last_state
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -66,8 +66,8 @@ PLATFORM_SCHEMA = PLATFORM_SCHEMA.extend({
 }).extend(mqtt.MQTT_AVAILABILITY_SCHEMA.schema)
 
 
-@asyncio.coroutine
-def async_setup_platform(hass, config, async_add_devices, discovery_info=None):
+async def async_setup_platform(hass, config, async_add_devices,
+                               discovery_info=None):
     """Set up a MQTT Template light."""
     if discovery_info is not None:
         config = PLATFORM_SCHEMA(discovery_info)
@@ -152,10 +152,11 @@ class MqttTemplate(MqttAvailability, Light):
             if tpl is not None:
                 tpl.hass = hass
 
-    @asyncio.coroutine
-    def async_added_to_hass(self):
+    async def async_added_to_hass(self):
         """Subscribe to MQTT events."""
-        yield from super().async_added_to_hass()
+        await super().async_added_to_hass()
+
+        last_state = await async_get_last_state(self.hass, self.entity_id)
 
         @callback
         def state_received(topic, payload, qos):
@@ -223,10 +224,23 @@ class MqttTemplate(MqttAvailability, Light):
             self.async_schedule_update_ha_state()
 
         if self._topics[CONF_STATE_TOPIC] is not None:
-            yield from mqtt.async_subscribe(
+            await mqtt.async_subscribe(
                 self.hass, self._topics[CONF_STATE_TOPIC], state_received,
                 self._qos)
 
+        if self._optimistic and last_state:
+            self._state = last_state.state == STATE_ON
+            if last_state.attributes.get(ATTR_BRIGHTNESS):
+                self._brightness = last_state.attributes.get(ATTR_BRIGHTNESS)
+            if last_state.attributes.get(ATTR_HS_COLOR):
+                self._hs = last_state.attributes.get(ATTR_HS_COLOR)
+            if last_state.attributes.get(ATTR_COLOR_TEMP):
+                self._color_temp = last_state.attributes.get(ATTR_COLOR_TEMP)
+            if last_state.attributes.get(ATTR_EFFECT):
+                self._effect = last_state.attributes.get(ATTR_EFFECT)
+            if last_state.attributes.get(ATTR_WHITE_VALUE):
+                self._white_value = last_state.attributes.get(ATTR_WHITE_VALUE)
+
     @property
     def brightness(self):
         """Return the brightness of this light between 0..255."""
@@ -280,8 +294,7 @@ class MqttTemplate(MqttAvailability, Light):
         """Return the current effect."""
         return self._effect
 
-    @asyncio.coroutine
-    def async_turn_on(self, **kwargs):
+    async def async_turn_on(self, **kwargs):
         """Turn the entity on.
 
         This method is a coroutine.
@@ -339,8 +352,7 @@ class MqttTemplate(MqttAvailability, Light):
         if self._optimistic:
             self.async_schedule_update_ha_state()
 
-    @asyncio.coroutine
-    def async_turn_off(self, **kwargs):
+    async def async_turn_off(self, **kwargs):
         """Turn the entity off.
 
         This method is a coroutine.

tests/components/light/test_mqtt.py

@@ -140,14 +140,16 @@ light:
 """
 import unittest
 from unittest import mock
+from unittest.mock import patch
 
 from homeassistant.setup import setup_component
 from homeassistant.const import (
     STATE_ON, STATE_OFF, STATE_UNAVAILABLE, ATTR_ASSUMED_STATE)
 import homeassistant.components.light as light
+import homeassistant.core as ha
 from tests.common import (
     assert_setup_component, get_test_home_assistant, mock_mqtt_component,
-    fire_mqtt_message)
+    fire_mqtt_message, mock_coro)
 
 
 class TestLightMQTT(unittest.TestCase):
@@ -481,12 +483,23 @@ class TestLightMQTT(unittest.TestCase):
             'payload_on': 'on',
             'payload_off': 'off'
         }}
-
+        fake_state = ha.State('light.test', 'on', {'brightness': 95,
+                                                   'hs_color': [100, 100],
+                                                   'effect': 'random',
+                                                   'color_temp': 100,
+                                                   'white_value': 50})
+        with patch('homeassistant.components.light.mqtt.async_get_last_state',
+                   return_value=mock_coro(fake_state)):
             with assert_setup_component(1, light.DOMAIN):
                 assert setup_component(self.hass, light.DOMAIN, config)
 
         state = self.hass.states.get('light.test')
-        self.assertEqual(STATE_OFF, state.state)
+        self.assertEqual(STATE_ON, state.state)
+        self.assertEqual(95, state.attributes.get('brightness'))
+        self.assertEqual((100, 100), state.attributes.get('hs_color'))
+        self.assertEqual('random', state.attributes.get('effect'))
+        self.assertEqual(100, state.attributes.get('color_temp'))
+        self.assertEqual(50, state.attributes.get('white_value'))
         self.assertTrue(state.attributes.get(ATTR_ASSUMED_STATE))
 
         light.turn_on(self.hass, 'light.test')

tests/components/light/test_mqtt_json.py

@@ -90,15 +90,17 @@ light:
 
 import json
 import unittest
+from unittest.mock import patch
 
 from homeassistant.setup import setup_component
 from homeassistant.const import (
     STATE_ON, STATE_OFF, STATE_UNAVAILABLE, ATTR_ASSUMED_STATE,
     ATTR_SUPPORTED_FEATURES)
 import homeassistant.components.light as light
+import homeassistant.core as ha
 from tests.common import (
     get_test_home_assistant, mock_mqtt_component, fire_mqtt_message,
-    assert_setup_component)
+    assert_setup_component, mock_coro)
 
 
 class TestLightMQTTJSON(unittest.TestCase):
@@ -284,6 +286,15 @@ class TestLightMQTTJSON(unittest.TestCase):
     def test_sending_mqtt_commands_and_optimistic(self): \
             # pylint: disable=invalid-name
         """Test the sending of command in optimistic mode."""
+        fake_state = ha.State('light.test', 'on', {'brightness': 95,
+                                                   'hs_color': [100, 100],
+                                                   'effect': 'random',
+                                                   'color_temp': 100,
+                                                   'white_value': 50})
+
+        with patch('homeassistant.components.light.mqtt_json'
+                   '.async_get_last_state',
+                   return_value=mock_coro(fake_state)):
             assert setup_component(self.hass, light.DOMAIN, {
                 light.DOMAIN: {
                     'platform': 'mqtt_json',
@@ -299,7 +310,12 @@ class TestLightMQTTJSON(unittest.TestCase):
             })
 
         state = self.hass.states.get('light.test')
-        self.assertEqual(STATE_OFF, state.state)
+        self.assertEqual(STATE_ON, state.state)
+        self.assertEqual(95, state.attributes.get('brightness'))
+        self.assertEqual((100, 100), state.attributes.get('hs_color'))
+        self.assertEqual('random', state.attributes.get('effect'))
+        self.assertEqual(100, state.attributes.get('color_temp'))
+        self.assertEqual(50, state.attributes.get('white_value'))
         self.assertEqual(191, state.attributes.get(ATTR_SUPPORTED_FEATURES))
         self.assertTrue(state.attributes.get(ATTR_ASSUMED_STATE))
 

tests/components/light/test_mqtt_template.py

@@ -27,14 +27,16 @@ If your light doesn't support white value feature, omit `white_value_template`.
 If your light doesn't support RGB feature, omit `(red|green|blue)_template`.
 """
 import unittest
+from unittest.mock import patch
 
 from homeassistant.setup import setup_component
 from homeassistant.const import (
     STATE_ON, STATE_OFF, STATE_UNAVAILABLE, ATTR_ASSUMED_STATE)
 import homeassistant.components.light as light
+import homeassistant.core as ha
 from tests.common import (
     get_test_home_assistant, mock_mqtt_component, fire_mqtt_message,
-    assert_setup_component)
+    assert_setup_component, mock_coro)
 
 
 class TestLightMQTTTemplate(unittest.TestCase):
@@ -207,6 +209,15 @@ class TestLightMQTTTemplate(unittest.TestCase):
     def test_optimistic(self): \
             # pylint: disable=invalid-name
         """Test optimistic mode."""
+        fake_state = ha.State('light.test', 'on', {'brightness': 95,
+                                                   'hs_color': [100, 100],
+                                                   'effect': 'random',
+                                                   'color_temp': 100,
+                                                   'white_value': 50})
+
+        with patch('homeassistant.components.light.mqtt_template'
+                   '.async_get_last_state',
+                   return_value=mock_coro(fake_state)):
             with assert_setup_component(1, light.DOMAIN):
                 assert setup_component(self.hass, light.DOMAIN, {
                     light.DOMAIN: {
@@ -226,7 +237,12 @@ class TestLightMQTTTemplate(unittest.TestCase):
                 })
 
         state = self.hass.states.get('light.test')
-        self.assertEqual(STATE_OFF, state.state)
+        self.assertEqual(STATE_ON, state.state)
+        self.assertEqual(95, state.attributes.get('brightness'))
+        self.assertEqual((100, 100), state.attributes.get('hs_color'))
+        self.assertEqual('random', state.attributes.get('effect'))
+        self.assertEqual(100, state.attributes.get('color_temp'))
+        self.assertEqual(50, state.attributes.get('white_value'))
         self.assertTrue(state.attributes.get(ATTR_ASSUMED_STATE))
 
         # turn on the light