From 9fde7bcad387f8399bc5d944929767081d0cf65d Mon Sep 17 00:00:00 2001
From: definitio <37266727+definitio@users.noreply.github.com>
Date: Fri, 5 Jun 2020 17:50:00 +0400
Subject: [PATCH 1/5] Report 'unavailable' state when not connected

to MQTT broker
---
 homeassistant/components/mqtt/__init__.py              | 2 ++
 homeassistant/components/mqtt/alarm_control_panel.py   | 5 -----
 homeassistant/components/mqtt/binary_sensor.py         | 5 -----
 homeassistant/components/mqtt/climate.py               | 5 -----
 homeassistant/components/mqtt/cover.py                 | 5 -----
 homeassistant/components/mqtt/fan.py                   | 5 -----
 homeassistant/components/mqtt/light/schema_basic.py    | 5 -----
 homeassistant/components/mqtt/light/schema_json.py     | 5 -----
 homeassistant/components/mqtt/light/schema_template.py | 8 --------
 homeassistant/components/mqtt/lock.py                  | 5 -----
 homeassistant/components/mqtt/sensor.py                | 5 -----
 homeassistant/components/mqtt/switch.py                | 5 -----
 homeassistant/components/mqtt/vacuum/schema_legacy.py  | 5 -----
 13 files changed, 2 insertions(+), 63 deletions(-)

diff --git a/homeassistant/components/mqtt/__init__.py b/homeassistant/components/mqtt/__init__.py
index 2be3189597939..593cc16756844 100644
--- a/homeassistant/components/mqtt/__init__.py
+++ b/homeassistant/components/mqtt/__init__.py
@@ -1141,6 +1141,8 @@ async def async_will_remove_from_hass(self):
     def available(self) -> bool:
         """Return if the device is available."""
         availability_topic = self._avail_config.get(CONF_AVAILABILITY_TOPIC)
+        if not self.hass.data[DATA_MQTT].connected:
+            return False
         return availability_topic is None or self._available
 
 
diff --git a/homeassistant/components/mqtt/alarm_control_panel.py b/homeassistant/components/mqtt/alarm_control_panel.py
index dae94b5a781e0..b4c82a8da2ada 100644
--- a/homeassistant/components/mqtt/alarm_control_panel.py
+++ b/homeassistant/components/mqtt/alarm_control_panel.py
@@ -222,11 +222,6 @@ async def async_will_remove_from_hass(self):
         await MqttAvailability.async_will_remove_from_hass(self)
         await MqttDiscoveryUpdate.async_will_remove_from_hass(self)
 
-    @property
-    def should_poll(self):
-        """No polling needed."""
-        return False
-
     @property
     def name(self):
         """Return the name of the device."""
diff --git a/homeassistant/components/mqtt/binary_sensor.py b/homeassistant/components/mqtt/binary_sensor.py
index 4e335dda95954..899bc3c3f755b 100644
--- a/homeassistant/components/mqtt/binary_sensor.py
+++ b/homeassistant/components/mqtt/binary_sensor.py
@@ -245,11 +245,6 @@ def value_is_expired(self, *_):
 
         self.async_write_ha_state()
 
-    @property
-    def should_poll(self):
-        """Return the polling state."""
-        return False
-
     @property
     def name(self):
         """Return the name of the binary sensor."""
diff --git a/homeassistant/components/mqtt/climate.py b/homeassistant/components/mqtt/climate.py
index 8cc6a3ffbdbe7..c8e4409e90aaa 100644
--- a/homeassistant/components/mqtt/climate.py
+++ b/homeassistant/components/mqtt/climate.py
@@ -566,11 +566,6 @@ async def async_will_remove_from_hass(self):
         await MqttAvailability.async_will_remove_from_hass(self)
         await MqttDiscoveryUpdate.async_will_remove_from_hass(self)
 
-    @property
-    def should_poll(self):
-        """Return the polling state."""
-        return False
-
     @property
     def name(self):
         """Return the name of the climate device."""
diff --git a/homeassistant/components/mqtt/cover.py b/homeassistant/components/mqtt/cover.py
index c9fc388e1a356..b58503ee5d025 100644
--- a/homeassistant/components/mqtt/cover.py
+++ b/homeassistant/components/mqtt/cover.py
@@ -374,11 +374,6 @@ async def async_will_remove_from_hass(self):
         await MqttAvailability.async_will_remove_from_hass(self)
         await MqttDiscoveryUpdate.async_will_remove_from_hass(self)
 
-    @property
-    def should_poll(self):
-        """No polling needed."""
-        return False
-
     @property
     def assumed_state(self):
         """Return true if we do optimistic updates."""
diff --git a/homeassistant/components/mqtt/fan.py b/homeassistant/components/mqtt/fan.py
index 9ce3809bfe848..d4cb85f77894a 100644
--- a/homeassistant/components/mqtt/fan.py
+++ b/homeassistant/components/mqtt/fan.py
@@ -322,11 +322,6 @@ async def async_will_remove_from_hass(self):
         await MqttAvailability.async_will_remove_from_hass(self)
         await MqttDiscoveryUpdate.async_will_remove_from_hass(self)
 
-    @property
-    def should_poll(self):
-        """No polling needed for a MQTT fan."""
-        return False
-
     @property
     def assumed_state(self):
         """Return true if we do optimistic updates."""
diff --git a/homeassistant/components/mqtt/light/schema_basic.py b/homeassistant/components/mqtt/light/schema_basic.py
index 375318764d1ff..9dab567405a23 100644
--- a/homeassistant/components/mqtt/light/schema_basic.py
+++ b/homeassistant/components/mqtt/light/schema_basic.py
@@ -586,11 +586,6 @@ def white_value(self):
             white_value = min(round(white_value), 255)
         return white_value
 
-    @property
-    def should_poll(self):
-        """No polling needed for a MQTT light."""
-        return False
-
     @property
     def name(self):
         """Return the name of the device if any."""
diff --git a/homeassistant/components/mqtt/light/schema_json.py b/homeassistant/components/mqtt/light/schema_json.py
index 89ae0601fdefc..1ab557d875d0d 100644
--- a/homeassistant/components/mqtt/light/schema_json.py
+++ b/homeassistant/components/mqtt/light/schema_json.py
@@ -393,11 +393,6 @@ def white_value(self):
         """Return the white property."""
         return self._white_value
 
-    @property
-    def should_poll(self):
-        """No polling needed for a MQTT light."""
-        return False
-
     @property
     def name(self):
         """Return the name of the device if any."""
diff --git a/homeassistant/components/mqtt/light/schema_template.py b/homeassistant/components/mqtt/light/schema_template.py
index a9f18e7039b20..28d0c682fce83 100644
--- a/homeassistant/components/mqtt/light/schema_template.py
+++ b/homeassistant/components/mqtt/light/schema_template.py
@@ -361,14 +361,6 @@ def white_value(self):
         """Return the white property."""
         return self._white_value
 
-    @property
-    def should_poll(self):
-        """Return True if entity has to be polled for state.
-
-        False if entity pushes its state to HA.
-        """
-        return False
-
     @property
     def name(self):
         """Return the name of the entity."""
diff --git a/homeassistant/components/mqtt/lock.py b/homeassistant/components/mqtt/lock.py
index 34905fe8aa472..be14d34a8a84d 100644
--- a/homeassistant/components/mqtt/lock.py
+++ b/homeassistant/components/mqtt/lock.py
@@ -195,11 +195,6 @@ async def async_will_remove_from_hass(self):
         await MqttAvailability.async_will_remove_from_hass(self)
         await MqttDiscoveryUpdate.async_will_remove_from_hass(self)
 
-    @property
-    def should_poll(self):
-        """No polling needed."""
-        return False
-
     @property
     def name(self):
         """Return the name of the lock."""
diff --git a/homeassistant/components/mqtt/sensor.py b/homeassistant/components/mqtt/sensor.py
index 2704c5ae3a1f7..50aca98e5a079 100644
--- a/homeassistant/components/mqtt/sensor.py
+++ b/homeassistant/components/mqtt/sensor.py
@@ -192,11 +192,6 @@ def value_is_expired(self, *_):
         self._state = None
         self.async_write_ha_state()
 
-    @property
-    def should_poll(self):
-        """No polling needed."""
-        return False
-
     @property
     def name(self):
         """Return the name of the sensor."""
diff --git a/homeassistant/components/mqtt/switch.py b/homeassistant/components/mqtt/switch.py
index 44d028829d06d..8bebcd5da24fc 100644
--- a/homeassistant/components/mqtt/switch.py
+++ b/homeassistant/components/mqtt/switch.py
@@ -206,11 +206,6 @@ async def async_will_remove_from_hass(self):
         await MqttAvailability.async_will_remove_from_hass(self)
         await MqttDiscoveryUpdate.async_will_remove_from_hass(self)
 
-    @property
-    def should_poll(self):
-        """Return the polling state."""
-        return False
-
     @property
     def name(self):
         """Return the name of the switch."""
diff --git a/homeassistant/components/mqtt/vacuum/schema_legacy.py b/homeassistant/components/mqtt/vacuum/schema_legacy.py
index b69565f711443..ea5c8b3642870 100644
--- a/homeassistant/components/mqtt/vacuum/schema_legacy.py
+++ b/homeassistant/components/mqtt/vacuum/schema_legacy.py
@@ -377,11 +377,6 @@ def name(self):
         """Return the name of the vacuum."""
         return self._name
 
-    @property
-    def should_poll(self):
-        """No polling needed for an MQTT vacuum."""
-        return False
-
     @property
     def is_on(self):
         """Return true if vacuum is on."""

From 508c93ef89534c5851b0f70c023d2e4cbbe9a9d5 Mon Sep 17 00:00:00 2001
From: definitio <37266727+definitio@users.noreply.github.com>
Date: Fri, 5 Jun 2020 23:58:06 +0400
Subject: [PATCH 2/5] Fix tests

---
 tests/components/mqtt/test_binary_sensor.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tests/components/mqtt/test_binary_sensor.py b/tests/components/mqtt/test_binary_sensor.py
index 8c68fabf2148a..7b84f4b152f88 100644
--- a/tests/components/mqtt/test_binary_sensor.py
+++ b/tests/components/mqtt/test_binary_sensor.py
@@ -402,7 +402,7 @@ async def test_off_delay(hass, mqtt_mock):
                 "state_topic": "test-topic",
                 "payload_on": "ON",
                 "payload_off": "OFF",
-                "off_delay": 30,
+                "off_delay": 3,
                 "force_update": True,
             }
         },
@@ -430,7 +430,7 @@ def callback(event):
     assert state.state == STATE_ON
     assert len(events) == 2
 
-    async_fire_time_changed(hass, dt_util.utcnow() + timedelta(seconds=30))
+    async_fire_time_changed(hass, dt_util.utcnow() + timedelta(seconds=3))
     await hass.async_block_till_done()
     state = hass.states.get("binary_sensor.test")
     assert state.state == STATE_OFF

From 46e64dc3858ce18b7034cf8bdacf5306925cf444 Mon Sep 17 00:00:00 2001
From: definitio <37266727+definitio@users.noreply.github.com>
Date: Sat, 6 Jun 2020 14:13:02 +0400
Subject: [PATCH 3/5] Rewrite to remove the polling

---
 homeassistant/components/mqtt/__init__.py           | 13 ++++++++++++-
 .../components/mqtt/alarm_control_panel.py          |  5 +++++
 homeassistant/components/mqtt/binary_sensor.py      |  5 +++++
 homeassistant/components/mqtt/climate.py            |  5 +++++
 homeassistant/components/mqtt/const.py              |  3 +++
 homeassistant/components/mqtt/cover.py              |  5 +++++
 homeassistant/components/mqtt/fan.py                |  5 +++++
 homeassistant/components/mqtt/light/schema_basic.py |  5 +++++
 homeassistant/components/mqtt/light/schema_json.py  |  5 +++++
 .../components/mqtt/light/schema_template.py        |  8 ++++++++
 homeassistant/components/mqtt/lock.py               |  5 +++++
 homeassistant/components/mqtt/sensor.py             |  5 +++++
 homeassistant/components/mqtt/switch.py             |  5 +++++
 .../components/mqtt/vacuum/schema_legacy.py         |  5 +++++
 tests/components/mqtt/test_binary_sensor.py         |  4 ++--
 15 files changed, 80 insertions(+), 3 deletions(-)

diff --git a/homeassistant/components/mqtt/__init__.py b/homeassistant/components/mqtt/__init__.py
index 593cc16756844..2dbd7a0380161 100644
--- a/homeassistant/components/mqtt/__init__.py
+++ b/homeassistant/components/mqtt/__init__.py
@@ -32,7 +32,7 @@
 from homeassistant.core import Event, ServiceCall, callback
 from homeassistant.exceptions import HomeAssistantError, Unauthorized
 from homeassistant.helpers import config_validation as cv, event, template
-from homeassistant.helpers.dispatcher import async_dispatcher_connect
+from homeassistant.helpers.dispatcher import async_dispatcher_connect, dispatcher_send
 from homeassistant.helpers.entity import Entity
 from homeassistant.helpers.typing import ConfigType, HomeAssistantType, ServiceDataType
 from homeassistant.loader import bind_hass
@@ -51,6 +51,8 @@
     CONF_STATE_TOPIC,
     DEFAULT_DISCOVERY,
     DEFAULT_QOS,
+    MQTT_CONNECTED,
+    MQTT_DISCONNECTED,
     PROTOCOL_311,
 )
 from .debug_info import log_messages
@@ -923,6 +925,7 @@ def _mqtt_on_connect(self, _mqttc, _userdata, _flags, result_code: int) -> None:
             return
 
         self.connected = True
+        dispatcher_send(self.hass, MQTT_CONNECTED)
         _LOGGER.info("Connected to MQTT server (%s)", result_code)
 
         # Group subscriptions to only re-subscribe once for each topic.
@@ -990,6 +993,7 @@ def _mqtt_handle_message(self, msg) -> None:
     def _mqtt_on_disconnect(self, _mqttc, _userdata, result_code: int) -> None:
         """Disconnected callback."""
         self.connected = False
+        dispatcher_send(self.hass, MQTT_DISCONNECTED)
         _LOGGER.warning("Disconnected from MQTT server (%s)", result_code)
 
 
@@ -1099,6 +1103,8 @@ async def async_added_to_hass(self) -> None:
         """Subscribe MQTT events."""
         await super().async_added_to_hass()
         await self._availability_subscribe_topics()
+        async_dispatcher_connect(self.hass, MQTT_CONNECTED, self.async_mqtt_connect)
+        async_dispatcher_connect(self.hass, MQTT_DISCONNECTED, self.async_mqtt_connect)
 
     async def availability_discovery_update(self, config: dict):
         """Handle updated discovery message."""
@@ -1131,6 +1137,11 @@ def availability_message_received(msg: Message) -> None:
             },
         )
 
+    @callback
+    async def async_mqtt_connect(self):
+        """Update state on connection/disconnection to MQTT broker."""
+        self.async_write_ha_state()
+
     async def async_will_remove_from_hass(self):
         """Unsubscribe when removed."""
         self._availability_sub_state = await async_unsubscribe_topics(
diff --git a/homeassistant/components/mqtt/alarm_control_panel.py b/homeassistant/components/mqtt/alarm_control_panel.py
index b4c82a8da2ada..dae94b5a781e0 100644
--- a/homeassistant/components/mqtt/alarm_control_panel.py
+++ b/homeassistant/components/mqtt/alarm_control_panel.py
@@ -222,6 +222,11 @@ async def async_will_remove_from_hass(self):
         await MqttAvailability.async_will_remove_from_hass(self)
         await MqttDiscoveryUpdate.async_will_remove_from_hass(self)
 
+    @property
+    def should_poll(self):
+        """No polling needed."""
+        return False
+
     @property
     def name(self):
         """Return the name of the device."""
diff --git a/homeassistant/components/mqtt/binary_sensor.py b/homeassistant/components/mqtt/binary_sensor.py
index 899bc3c3f755b..4e335dda95954 100644
--- a/homeassistant/components/mqtt/binary_sensor.py
+++ b/homeassistant/components/mqtt/binary_sensor.py
@@ -245,6 +245,11 @@ def value_is_expired(self, *_):
 
         self.async_write_ha_state()
 
+    @property
+    def should_poll(self):
+        """Return the polling state."""
+        return False
+
     @property
     def name(self):
         """Return the name of the binary sensor."""
diff --git a/homeassistant/components/mqtt/climate.py b/homeassistant/components/mqtt/climate.py
index c8e4409e90aaa..8cc6a3ffbdbe7 100644
--- a/homeassistant/components/mqtt/climate.py
+++ b/homeassistant/components/mqtt/climate.py
@@ -566,6 +566,11 @@ async def async_will_remove_from_hass(self):
         await MqttAvailability.async_will_remove_from_hass(self)
         await MqttDiscoveryUpdate.async_will_remove_from_hass(self)
 
+    @property
+    def should_poll(self):
+        """Return the polling state."""
+        return False
+
     @property
     def name(self):
         """Return the name of the climate device."""
diff --git a/homeassistant/components/mqtt/const.py b/homeassistant/components/mqtt/const.py
index 5d1fe2e2505ee..acb24f4bddab1 100644
--- a/homeassistant/components/mqtt/const.py
+++ b/homeassistant/components/mqtt/const.py
@@ -9,3 +9,6 @@
 CONF_STATE_TOPIC = "state_topic"
 PROTOCOL_311 = "3.1.1"
 DEFAULT_QOS = 0
+
+MQTT_CONNECTED = "mqtt_connected"
+MQTT_DISCONNECTED = "mqtt_disconnected"
diff --git a/homeassistant/components/mqtt/cover.py b/homeassistant/components/mqtt/cover.py
index b58503ee5d025..c9fc388e1a356 100644
--- a/homeassistant/components/mqtt/cover.py
+++ b/homeassistant/components/mqtt/cover.py
@@ -374,6 +374,11 @@ async def async_will_remove_from_hass(self):
         await MqttAvailability.async_will_remove_from_hass(self)
         await MqttDiscoveryUpdate.async_will_remove_from_hass(self)
 
+    @property
+    def should_poll(self):
+        """No polling needed."""
+        return False
+
     @property
     def assumed_state(self):
         """Return true if we do optimistic updates."""
diff --git a/homeassistant/components/mqtt/fan.py b/homeassistant/components/mqtt/fan.py
index d4cb85f77894a..9ce3809bfe848 100644
--- a/homeassistant/components/mqtt/fan.py
+++ b/homeassistant/components/mqtt/fan.py
@@ -322,6 +322,11 @@ async def async_will_remove_from_hass(self):
         await MqttAvailability.async_will_remove_from_hass(self)
         await MqttDiscoveryUpdate.async_will_remove_from_hass(self)
 
+    @property
+    def should_poll(self):
+        """No polling needed for a MQTT fan."""
+        return False
+
     @property
     def assumed_state(self):
         """Return true if we do optimistic updates."""
diff --git a/homeassistant/components/mqtt/light/schema_basic.py b/homeassistant/components/mqtt/light/schema_basic.py
index 9dab567405a23..375318764d1ff 100644
--- a/homeassistant/components/mqtt/light/schema_basic.py
+++ b/homeassistant/components/mqtt/light/schema_basic.py
@@ -586,6 +586,11 @@ def white_value(self):
             white_value = min(round(white_value), 255)
         return white_value
 
+    @property
+    def should_poll(self):
+        """No polling needed for a MQTT light."""
+        return False
+
     @property
     def name(self):
         """Return the name of the device if any."""
diff --git a/homeassistant/components/mqtt/light/schema_json.py b/homeassistant/components/mqtt/light/schema_json.py
index 1ab557d875d0d..89ae0601fdefc 100644
--- a/homeassistant/components/mqtt/light/schema_json.py
+++ b/homeassistant/components/mqtt/light/schema_json.py
@@ -393,6 +393,11 @@ def white_value(self):
         """Return the white property."""
         return self._white_value
 
+    @property
+    def should_poll(self):
+        """No polling needed for a MQTT light."""
+        return False
+
     @property
     def name(self):
         """Return the name of the device if any."""
diff --git a/homeassistant/components/mqtt/light/schema_template.py b/homeassistant/components/mqtt/light/schema_template.py
index 28d0c682fce83..a9f18e7039b20 100644
--- a/homeassistant/components/mqtt/light/schema_template.py
+++ b/homeassistant/components/mqtt/light/schema_template.py
@@ -361,6 +361,14 @@ def white_value(self):
         """Return the white property."""
         return self._white_value
 
+    @property
+    def should_poll(self):
+        """Return True if entity has to be polled for state.
+
+        False if entity pushes its state to HA.
+        """
+        return False
+
     @property
     def name(self):
         """Return the name of the entity."""
diff --git a/homeassistant/components/mqtt/lock.py b/homeassistant/components/mqtt/lock.py
index be14d34a8a84d..34905fe8aa472 100644
--- a/homeassistant/components/mqtt/lock.py
+++ b/homeassistant/components/mqtt/lock.py
@@ -195,6 +195,11 @@ async def async_will_remove_from_hass(self):
         await MqttAvailability.async_will_remove_from_hass(self)
         await MqttDiscoveryUpdate.async_will_remove_from_hass(self)
 
+    @property
+    def should_poll(self):
+        """No polling needed."""
+        return False
+
     @property
     def name(self):
         """Return the name of the lock."""
diff --git a/homeassistant/components/mqtt/sensor.py b/homeassistant/components/mqtt/sensor.py
index 50aca98e5a079..2704c5ae3a1f7 100644
--- a/homeassistant/components/mqtt/sensor.py
+++ b/homeassistant/components/mqtt/sensor.py
@@ -192,6 +192,11 @@ def value_is_expired(self, *_):
         self._state = None
         self.async_write_ha_state()
 
+    @property
+    def should_poll(self):
+        """No polling needed."""
+        return False
+
     @property
     def name(self):
         """Return the name of the sensor."""
diff --git a/homeassistant/components/mqtt/switch.py b/homeassistant/components/mqtt/switch.py
index 8bebcd5da24fc..44d028829d06d 100644
--- a/homeassistant/components/mqtt/switch.py
+++ b/homeassistant/components/mqtt/switch.py
@@ -206,6 +206,11 @@ async def async_will_remove_from_hass(self):
         await MqttAvailability.async_will_remove_from_hass(self)
         await MqttDiscoveryUpdate.async_will_remove_from_hass(self)
 
+    @property
+    def should_poll(self):
+        """Return the polling state."""
+        return False
+
     @property
     def name(self):
         """Return the name of the switch."""
diff --git a/homeassistant/components/mqtt/vacuum/schema_legacy.py b/homeassistant/components/mqtt/vacuum/schema_legacy.py
index ea5c8b3642870..b69565f711443 100644
--- a/homeassistant/components/mqtt/vacuum/schema_legacy.py
+++ b/homeassistant/components/mqtt/vacuum/schema_legacy.py
@@ -377,6 +377,11 @@ def name(self):
         """Return the name of the vacuum."""
         return self._name
 
+    @property
+    def should_poll(self):
+        """No polling needed for an MQTT vacuum."""
+        return False
+
     @property
     def is_on(self):
         """Return true if vacuum is on."""
diff --git a/tests/components/mqtt/test_binary_sensor.py b/tests/components/mqtt/test_binary_sensor.py
index 7b84f4b152f88..8c68fabf2148a 100644
--- a/tests/components/mqtt/test_binary_sensor.py
+++ b/tests/components/mqtt/test_binary_sensor.py
@@ -402,7 +402,7 @@ async def test_off_delay(hass, mqtt_mock):
                 "state_topic": "test-topic",
                 "payload_on": "ON",
                 "payload_off": "OFF",
-                "off_delay": 3,
+                "off_delay": 30,
                 "force_update": True,
             }
         },
@@ -430,7 +430,7 @@ def callback(event):
     assert state.state == STATE_ON
     assert len(events) == 2
 
-    async_fire_time_changed(hass, dt_util.utcnow() + timedelta(seconds=3))
+    async_fire_time_changed(hass, dt_util.utcnow() + timedelta(seconds=30))
     await hass.async_block_till_done()
     state = hass.states.get("binary_sensor.test")
     assert state.state == STATE_OFF

From 35e7c18a56d10d4a92f2298ffbda7b8590cef5a9 Mon Sep 17 00:00:00 2001
From: definitio <37266727+definitio@users.noreply.github.com>
Date: Sat, 6 Jun 2020 15:23:27 +0400
Subject: [PATCH 4/5] Add tests

---
 .../mqtt/test_alarm_control_panel.py           |  8 ++++++++
 tests/components/mqtt/test_binary_sensor.py    |  8 ++++++++
 tests/components/mqtt/test_camera.py           |  8 ++++++++
 tests/components/mqtt/test_climate.py          |  8 ++++++++
 tests/components/mqtt/test_common.py           | 18 ++++++++++++++++++
 tests/components/mqtt/test_cover.py            |  8 ++++++++
 tests/components/mqtt/test_fan.py              |  8 ++++++++
 tests/components/mqtt/test_legacy_vacuum.py    |  8 ++++++++
 tests/components/mqtt/test_light.py            |  8 ++++++++
 tests/components/mqtt/test_light_json.py       |  8 ++++++++
 tests/components/mqtt/test_light_template.py   |  8 ++++++++
 tests/components/mqtt/test_lock.py             |  8 ++++++++
 tests/components/mqtt/test_sensor.py           |  8 ++++++++
 tests/components/mqtt/test_state_vacuum.py     |  8 ++++++++
 tests/components/mqtt/test_switch.py           |  8 ++++++++
 15 files changed, 130 insertions(+)

diff --git a/tests/components/mqtt/test_alarm_control_panel.py b/tests/components/mqtt/test_alarm_control_panel.py
index 03e8133bde96e..29ecb277928f4 100644
--- a/tests/components/mqtt/test_alarm_control_panel.py
+++ b/tests/components/mqtt/test_alarm_control_panel.py
@@ -19,6 +19,7 @@
 )
 
 from .test_common import (
+    help_test_availability_when_connection_lost,
     help_test_availability_without_topic,
     help_test_custom_availability_payload,
     help_test_default_availability_payload,
@@ -469,6 +470,13 @@ async def test_attributes_code_text(hass, mqtt_mock):
     )
 
 
+async def test_availability_when_connection_lost(hass, mqtt_mock):
+    """Test availability after MQTT disconnection."""
+    await help_test_availability_when_connection_lost(
+        hass, mqtt_mock, alarm_control_panel.DOMAIN, DEFAULT_CONFIG_CODE
+    )
+
+
 async def test_availability_without_topic(hass, mqtt_mock):
     """Test availability without defined availability topic."""
     await help_test_availability_without_topic(
diff --git a/tests/components/mqtt/test_binary_sensor.py b/tests/components/mqtt/test_binary_sensor.py
index 8c68fabf2148a..dccf8396c5a67 100644
--- a/tests/components/mqtt/test_binary_sensor.py
+++ b/tests/components/mqtt/test_binary_sensor.py
@@ -18,6 +18,7 @@
 import homeassistant.util.dt as dt_util
 
 from .test_common import (
+    help_test_availability_when_connection_lost,
     help_test_availability_without_topic,
     help_test_custom_availability_payload,
     help_test_default_availability_payload,
@@ -298,6 +299,13 @@ async def test_invalid_device_class(hass, mqtt_mock):
     assert state is None
 
 
+async def test_availability_when_connection_lost(hass, mqtt_mock):
+    """Test availability after MQTT disconnection."""
+    await help_test_availability_when_connection_lost(
+        hass, mqtt_mock, binary_sensor.DOMAIN, DEFAULT_CONFIG
+    )
+
+
 async def test_availability_without_topic(hass, mqtt_mock):
     """Test availability without defined availability topic."""
     await help_test_availability_without_topic(
diff --git a/tests/components/mqtt/test_camera.py b/tests/components/mqtt/test_camera.py
index 11b846d4c38f3..98e85e2df7c64 100644
--- a/tests/components/mqtt/test_camera.py
+++ b/tests/components/mqtt/test_camera.py
@@ -8,6 +8,7 @@
 from homeassistant.setup import async_setup_component
 
 from .test_common import (
+    help_test_availability_when_connection_lost,
     help_test_availability_without_topic,
     help_test_custom_availability_payload,
     help_test_default_availability_payload,
@@ -62,6 +63,13 @@ async def test_run_camera_setup(hass, aiohttp_client):
     assert body == "beer"
 
 
+async def test_availability_when_connection_lost(hass, mqtt_mock):
+    """Test availability after MQTT disconnection."""
+    await help_test_availability_when_connection_lost(
+        hass, mqtt_mock, camera.DOMAIN, DEFAULT_CONFIG
+    )
+
+
 async def test_availability_without_topic(hass, mqtt_mock):
     """Test availability without defined availability topic."""
     await help_test_availability_without_topic(
diff --git a/tests/components/mqtt/test_climate.py b/tests/components/mqtt/test_climate.py
index 30018c7c17578..1736174d23241 100644
--- a/tests/components/mqtt/test_climate.py
+++ b/tests/components/mqtt/test_climate.py
@@ -25,6 +25,7 @@
 from homeassistant.const import STATE_OFF
 
 from .test_common import (
+    help_test_availability_when_connection_lost,
     help_test_availability_without_topic,
     help_test_custom_availability_payload,
     help_test_default_availability_payload,
@@ -608,6 +609,13 @@ async def test_set_aux(hass, mqtt_mock):
     assert state.attributes.get("aux_heat") == "off"
 
 
+async def test_availability_when_connection_lost(hass, mqtt_mock):
+    """Test availability after MQTT disconnection."""
+    await help_test_availability_when_connection_lost(
+        hass, mqtt_mock, CLIMATE_DOMAIN, DEFAULT_CONFIG
+    )
+
+
 async def test_availability_without_topic(hass, mqtt_mock):
     """Test availability without defined availability topic."""
     await help_test_availability_without_topic(
diff --git a/tests/components/mqtt/test_common.py b/tests/components/mqtt/test_common.py
index d0ddc1d483027..445b41922d982 100644
--- a/tests/components/mqtt/test_common.py
+++ b/tests/components/mqtt/test_common.py
@@ -6,8 +6,10 @@
 
 from homeassistant.components import mqtt
 from homeassistant.components.mqtt import debug_info
+from homeassistant.components.mqtt.const import MQTT_DISCONNECTED
 from homeassistant.components.mqtt.discovery import async_start
 from homeassistant.const import ATTR_ASSUMED_STATE, STATE_UNAVAILABLE
+from homeassistant.helpers.dispatcher import async_dispatcher_send
 
 from tests.async_mock import ANY
 from tests.common import (
@@ -35,6 +37,22 @@
 }
 
 
+async def help_test_availability_when_connection_lost(hass, mqtt_mock, domain, config):
+    """Test availability after MQTT disconnection."""
+    assert await async_setup_component(hass, domain, config)
+    await hass.async_block_till_done()
+
+    state = hass.states.get(f"{domain}.test")
+    assert state.state != STATE_UNAVAILABLE
+
+    hass.data["mqtt"].connected = False
+    async_dispatcher_send(hass, MQTT_DISCONNECTED)
+    await hass.async_block_till_done()
+
+    state = hass.states.get(f"{domain}.test")
+    assert state.state == STATE_UNAVAILABLE
+
+
 async def help_test_availability_without_topic(hass, mqtt_mock, domain, config):
     """Test availability without defined availability topic."""
     assert "availability_topic" not in config[domain]
diff --git a/tests/components/mqtt/test_cover.py b/tests/components/mqtt/test_cover.py
index eb758ebf93a4a..9ef745652061c 100644
--- a/tests/components/mqtt/test_cover.py
+++ b/tests/components/mqtt/test_cover.py
@@ -30,6 +30,7 @@
 from homeassistant.setup import async_setup_component
 
 from .test_common import (
+    help_test_availability_when_connection_lost,
     help_test_availability_without_topic,
     help_test_custom_availability_payload,
     help_test_default_availability_payload,
@@ -1735,6 +1736,13 @@ async def test_find_in_range_altered_inverted(hass, mqtt_mock):
     assert mqtt_cover.find_in_range_from_percent(60, "cover") == 120
 
 
+async def test_availability_when_connection_lost(hass, mqtt_mock):
+    """Test availability after MQTT disconnection."""
+    await help_test_availability_when_connection_lost(
+        hass, mqtt_mock, cover.DOMAIN, DEFAULT_CONFIG
+    )
+
+
 async def test_availability_without_topic(hass, mqtt_mock):
     """Test availability without defined availability topic."""
     await help_test_availability_without_topic(
diff --git a/tests/components/mqtt/test_fan.py b/tests/components/mqtt/test_fan.py
index 7f6eb79e85e1e..bec771b600471 100644
--- a/tests/components/mqtt/test_fan.py
+++ b/tests/components/mqtt/test_fan.py
@@ -11,6 +11,7 @@
 from homeassistant.setup import async_setup_component
 
 from .test_common import (
+    help_test_availability_when_connection_lost,
     help_test_availability_without_topic,
     help_test_custom_availability_payload,
     help_test_default_availability_payload,
@@ -600,6 +601,13 @@ async def test_supported_features(hass, mqtt_mock):
     )
 
 
+async def test_availability_when_connection_lost(hass, mqtt_mock):
+    """Test availability after MQTT disconnection."""
+    await help_test_availability_when_connection_lost(
+        hass, mqtt_mock, fan.DOMAIN, DEFAULT_CONFIG
+    )
+
+
 async def test_availability_without_topic(hass, mqtt_mock):
     """Test availability without defined availability topic."""
     await help_test_availability_without_topic(
diff --git a/tests/components/mqtt/test_legacy_vacuum.py b/tests/components/mqtt/test_legacy_vacuum.py
index 032a55edee4a0..9d02295db2ed4 100644
--- a/tests/components/mqtt/test_legacy_vacuum.py
+++ b/tests/components/mqtt/test_legacy_vacuum.py
@@ -23,6 +23,7 @@
 from homeassistant.setup import async_setup_component
 
 from .test_common import (
+    help_test_availability_when_connection_lost,
     help_test_availability_without_topic,
     help_test_custom_availability_payload,
     help_test_default_availability_payload,
@@ -543,6 +544,13 @@ async def test_missing_fan_speed_template(hass, mqtt_mock):
     assert state is None
 
 
+async def test_availability_when_connection_lost(hass, mqtt_mock):
+    """Test availability after MQTT disconnection."""
+    await help_test_availability_when_connection_lost(
+        hass, mqtt_mock, vacuum.DOMAIN, DEFAULT_CONFIG_2
+    )
+
+
 async def test_availability_without_topic(hass, mqtt_mock):
     """Test availability without defined availability topic."""
     await help_test_availability_without_topic(
diff --git a/tests/components/mqtt/test_light.py b/tests/components/mqtt/test_light.py
index f832e23591560..0e8d0c19a1994 100644
--- a/tests/components/mqtt/test_light.py
+++ b/tests/components/mqtt/test_light.py
@@ -162,6 +162,7 @@
 from homeassistant.setup import async_setup_component
 
 from .test_common import (
+    help_test_availability_when_connection_lost,
     help_test_availability_without_topic,
     help_test_custom_availability_payload,
     help_test_default_availability_payload,
@@ -1326,6 +1327,13 @@ async def test_effect(hass, mqtt_mock):
     mqtt_mock.async_publish.assert_called_once_with("test_light/set", "OFF", 0, False)
 
 
+async def test_availability_when_connection_lost(hass, mqtt_mock):
+    """Test availability after MQTT disconnection."""
+    await help_test_availability_when_connection_lost(
+        hass, mqtt_mock, light.DOMAIN, DEFAULT_CONFIG
+    )
+
+
 async def test_availability_without_topic(hass, mqtt_mock):
     """Test availability without defined availability topic."""
     await help_test_availability_without_topic(
diff --git a/tests/components/mqtt/test_light_json.py b/tests/components/mqtt/test_light_json.py
index bb9e2afb0e51f..19c531ecb30ff 100644
--- a/tests/components/mqtt/test_light_json.py
+++ b/tests/components/mqtt/test_light_json.py
@@ -102,6 +102,7 @@
 from homeassistant.setup import async_setup_component
 
 from .test_common import (
+    help_test_availability_when_connection_lost,
     help_test_availability_without_topic,
     help_test_custom_availability_payload,
     help_test_default_availability_payload,
@@ -1065,6 +1066,13 @@ async def test_invalid_values(hass, mqtt_mock):
     assert state.attributes.get("color_temp") == 100
 
 
+async def test_availability_when_connection_lost(hass, mqtt_mock):
+    """Test availability after MQTT disconnection."""
+    await help_test_availability_when_connection_lost(
+        hass, mqtt_mock, light.DOMAIN, DEFAULT_CONFIG
+    )
+
+
 async def test_availability_without_topic(hass, mqtt_mock):
     """Test availability without defined availability topic."""
     await help_test_availability_without_topic(
diff --git a/tests/components/mqtt/test_light_template.py b/tests/components/mqtt/test_light_template.py
index cb5aff40b4bf5..d1e27d4c5165b 100644
--- a/tests/components/mqtt/test_light_template.py
+++ b/tests/components/mqtt/test_light_template.py
@@ -39,6 +39,7 @@
 from homeassistant.setup import async_setup_component
 
 from .test_common import (
+    help_test_availability_when_connection_lost,
     help_test_availability_without_topic,
     help_test_custom_availability_payload,
     help_test_default_availability_payload,
@@ -797,6 +798,13 @@ async def test_invalid_values(hass, mqtt_mock):
     assert state.attributes.get("effect") == "rainbow"
 
 
+async def test_availability_when_connection_lost(hass, mqtt_mock):
+    """Test availability after MQTT disconnection."""
+    await help_test_availability_when_connection_lost(
+        hass, mqtt_mock, light.DOMAIN, DEFAULT_CONFIG
+    )
+
+
 async def test_availability_without_topic(hass, mqtt_mock):
     """Test availability without defined availability topic."""
     await help_test_availability_without_topic(
diff --git a/tests/components/mqtt/test_lock.py b/tests/components/mqtt/test_lock.py
index 0e9a9af850fd2..abea215f837a9 100644
--- a/tests/components/mqtt/test_lock.py
+++ b/tests/components/mqtt/test_lock.py
@@ -12,6 +12,7 @@
 from homeassistant.setup import async_setup_component
 
 from .test_common import (
+    help_test_availability_when_connection_lost,
     help_test_availability_without_topic,
     help_test_custom_availability_payload,
     help_test_default_availability_payload,
@@ -272,6 +273,13 @@ async def test_sending_mqtt_commands_and_explicit_optimistic(hass, mqtt_mock):
     assert state.attributes.get(ATTR_ASSUMED_STATE)
 
 
+async def test_availability_when_connection_lost(hass, mqtt_mock):
+    """Test availability after MQTT disconnection."""
+    await help_test_availability_when_connection_lost(
+        hass, mqtt_mock, LOCK_DOMAIN, DEFAULT_CONFIG
+    )
+
+
 async def test_availability_without_topic(hass, mqtt_mock):
     """Test availability without defined availability topic."""
     await help_test_availability_without_topic(
diff --git a/tests/components/mqtt/test_sensor.py b/tests/components/mqtt/test_sensor.py
index d711b9e3bb89c..be4b501281793 100644
--- a/tests/components/mqtt/test_sensor.py
+++ b/tests/components/mqtt/test_sensor.py
@@ -13,6 +13,7 @@
 import homeassistant.util.dt as dt_util
 
 from .test_common import (
+    help_test_availability_when_connection_lost,
     help_test_availability_without_topic,
     help_test_custom_availability_payload,
     help_test_default_availability_payload,
@@ -232,6 +233,13 @@ def callback(event):
     assert len(events) == 2
 
 
+async def test_availability_when_connection_lost(hass, mqtt_mock):
+    """Test availability after MQTT disconnection."""
+    await help_test_availability_when_connection_lost(
+        hass, mqtt_mock, sensor.DOMAIN, DEFAULT_CONFIG
+    )
+
+
 async def test_availability_without_topic(hass, mqtt_mock):
     """Test availability without defined availability topic."""
     await help_test_availability_without_topic(
diff --git a/tests/components/mqtt/test_state_vacuum.py b/tests/components/mqtt/test_state_vacuum.py
index f77a1a11ca191..1a59e84c2be26 100644
--- a/tests/components/mqtt/test_state_vacuum.py
+++ b/tests/components/mqtt/test_state_vacuum.py
@@ -33,6 +33,7 @@
 from homeassistant.setup import async_setup_component
 
 from .test_common import (
+    help_test_availability_when_connection_lost,
     help_test_availability_without_topic,
     help_test_custom_availability_payload,
     help_test_default_availability_payload,
@@ -321,6 +322,13 @@ async def test_status_invalid_json(hass, mqtt_mock):
     assert state.state == STATE_UNKNOWN
 
 
+async def test_availability_when_connection_lost(hass, mqtt_mock):
+    """Test availability after MQTT disconnection."""
+    await help_test_availability_when_connection_lost(
+        hass, mqtt_mock, vacuum.DOMAIN, DEFAULT_CONFIG_2
+    )
+
+
 async def test_availability_without_topic(hass, mqtt_mock):
     """Test availability without defined availability topic."""
     await help_test_availability_without_topic(
diff --git a/tests/components/mqtt/test_switch.py b/tests/components/mqtt/test_switch.py
index da66e2a7f609e..b90dfa684137f 100644
--- a/tests/components/mqtt/test_switch.py
+++ b/tests/components/mqtt/test_switch.py
@@ -7,6 +7,7 @@
 from homeassistant.setup import async_setup_component
 
 from .test_common import (
+    help_test_availability_when_connection_lost,
     help_test_availability_without_topic,
     help_test_custom_availability_payload,
     help_test_default_availability_payload,
@@ -155,6 +156,13 @@ async def test_controlling_state_via_topic_and_json_message(hass, mock_publish):
     assert state.state == STATE_OFF
 
 
+async def test_availability_when_connection_lost(hass, mqtt_mock):
+    """Test availability after MQTT disconnection."""
+    await help_test_availability_when_connection_lost(
+        hass, mqtt_mock, switch.DOMAIN, DEFAULT_CONFIG
+    )
+
+
 async def test_availability_without_topic(hass, mqtt_mock):
     """Test availability without defined availability topic."""
     await help_test_availability_without_topic(

From 62022fa91322469f677a8cab578ab7820b43ef41 Mon Sep 17 00:00:00 2001
From: definitio <37266727+definitio@users.noreply.github.com>
Date: Sat, 6 Jun 2020 16:23:01 +0400
Subject: [PATCH 5/5] Add some fixes

---
 homeassistant/components/mqtt/__init__.py | 2 +-
 tests/components/mqtt/test_common.py      | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/homeassistant/components/mqtt/__init__.py b/homeassistant/components/mqtt/__init__.py
index 2dbd7a0380161..26c844f106af9 100644
--- a/homeassistant/components/mqtt/__init__.py
+++ b/homeassistant/components/mqtt/__init__.py
@@ -1138,7 +1138,7 @@ def availability_message_received(msg: Message) -> None:
         )
 
     @callback
-    async def async_mqtt_connect(self):
+    def async_mqtt_connect(self):
         """Update state on connection/disconnection to MQTT broker."""
         self.async_write_ha_state()
 
diff --git a/tests/components/mqtt/test_common.py b/tests/components/mqtt/test_common.py
index 445b41922d982..d91c236680ddf 100644
--- a/tests/components/mqtt/test_common.py
+++ b/tests/components/mqtt/test_common.py
@@ -45,7 +45,7 @@ async def help_test_availability_when_connection_lost(hass, mqtt_mock, domain, c
     state = hass.states.get(f"{domain}.test")
     assert state.state != STATE_UNAVAILABLE
 
-    hass.data["mqtt"].connected = False
+    mqtt_mock.connected = False
     async_dispatcher_send(hass, MQTT_DISCONNECTED)
     await hass.async_block_till_done()