From c851dfa2c7bd77601976102d566d695006499fab Mon Sep 17 00:00:00 2001
From: Diogo Gomes <diogogomes@gmail.com>
Date: Wed, 2 May 2018 22:29:07 +0100
Subject: [PATCH] Restores switch state, case the switch is optimistic (#14151)

* Add restore_state to optimistic switch

* no need to schedule update

* test added

* lint

* new async syntax

* lint
---
 homeassistant/components/switch/mqtt.py |  9 +++++++-
 tests/components/switch/test_mqtt.py    | 30 +++++++++++++++----------
 2 files changed, 26 insertions(+), 13 deletions(-)

diff --git a/homeassistant/components/switch/mqtt.py b/homeassistant/components/switch/mqtt.py
index 15dc6f1d0f4..69f12536c5f 100644
--- a/homeassistant/components/switch/mqtt.py
+++ b/homeassistant/components/switch/mqtt.py
@@ -16,9 +16,10 @@ from homeassistant.components.mqtt import (
 from homeassistant.components.switch import SwitchDevice
 from homeassistant.const import (
     CONF_NAME, CONF_OPTIMISTIC, CONF_VALUE_TEMPLATE, CONF_PAYLOAD_OFF,
-    CONF_PAYLOAD_ON, CONF_ICON)
+    CONF_PAYLOAD_ON, CONF_ICON, STATE_ON)
 import homeassistant.components.mqtt as mqtt
 import homeassistant.helpers.config_validation as cv
+from homeassistant.helpers.restore_state import async_get_last_state
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -112,6 +113,12 @@ class MqttSwitch(MqttAvailability, SwitchDevice):
                 self.hass, self._state_topic, state_message_received,
                 self._qos)
 
+        if self._optimistic:
+            last_state = await async_get_last_state(self.hass,
+                                                    self.entity_id)
+            if last_state:
+                self._state = last_state.state == STATE_ON
+
     @property
     def should_poll(self):
         """Return the polling state."""
diff --git a/tests/components/switch/test_mqtt.py b/tests/components/switch/test_mqtt.py
index f79d0706321..b5e2a0b0395 100644
--- a/tests/components/switch/test_mqtt.py
+++ b/tests/components/switch/test_mqtt.py
@@ -1,12 +1,14 @@
 """The tests for the MQTT switch platform."""
 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.core as ha
 import homeassistant.components.switch as switch
 from tests.common import (
-    mock_mqtt_component, fire_mqtt_message, get_test_home_assistant)
+    mock_mqtt_component, fire_mqtt_message, get_test_home_assistant, mock_coro)
 
 
 class TestSwitchMQTT(unittest.TestCase):
@@ -52,19 +54,23 @@ class TestSwitchMQTT(unittest.TestCase):
 
     def test_sending_mqtt_commands_and_optimistic(self):
         """Test the sending MQTT commands in optimistic mode."""
-        assert setup_component(self.hass, switch.DOMAIN, {
-            switch.DOMAIN: {
-                'platform': 'mqtt',
-                'name': 'test',
-                'command_topic': 'command-topic',
-                'payload_on': 'beer on',
-                'payload_off': 'beer off',
-                'qos': '2'
-            }
-        })
+        fake_state = ha.State('switch.test', 'on')
+
+        with patch('homeassistant.components.switch.mqtt.async_get_last_state',
+                   return_value=mock_coro(fake_state)):
+            assert setup_component(self.hass, switch.DOMAIN, {
+                switch.DOMAIN: {
+                    'platform': 'mqtt',
+                    'name': 'test',
+                    'command_topic': 'command-topic',
+                    'payload_on': 'beer on',
+                    'payload_off': 'beer off',
+                    'qos': '2'
+                }
+            })
 
         state = self.hass.states.get('switch.test')
-        self.assertEqual(STATE_OFF, state.state)
+        self.assertEqual(STATE_ON, state.state)
         self.assertTrue(state.attributes.get(ATTR_ASSUMED_STATE))
 
         switch.turn_on(self.hass, 'switch.test')