From ad5101c5c035ff5a8674e74fb98c7cc6d6927432 Mon Sep 17 00:00:00 2001
From: definitio <37266727+definitio@users.noreply.github.com>
Date: Sun, 7 Jun 2020 11:21:16 +0400
Subject: [PATCH] Set state for MQTT entities to 'unavailable' when no
 connection to broker (#36479)

* Report 'unavailable' state when not connected

to MQTT broker

* Fix tests

* Rewrite to remove the polling

* Add tests

* Add some fixes
---
 homeassistant/components/mqtt/__init__.py      | 15 ++++++++++++++-
 homeassistant/components/mqtt/const.py         |  3 +++
 .../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 ++++++++
 17 files changed, 147 insertions(+), 1 deletion(-)

diff --git a/homeassistant/components/mqtt/__init__.py b/homeassistant/components/mqtt/__init__.py
index 2be31895979..26c844f106a 100644
--- a/homeassistant/components/mqtt/__init__.py
+++ b/homeassistant/components/mqtt/__init__.py
@@ -32,7 +32,7 @@ from homeassistant.const import (
 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 @@ from .const import (
     CONF_STATE_TOPIC,
     DEFAULT_DISCOVERY,
     DEFAULT_QOS,
+    MQTT_CONNECTED,
+    MQTT_DISCONNECTED,
     PROTOCOL_311,
 )
 from .debug_info import log_messages
@@ -923,6 +925,7 @@ class MQTT:
             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 @@ class MQTT:
     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 @@ class MqttAvailability(Entity):
         """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 @@ class MqttAvailability(Entity):
             },
         )
 
+    @callback
+    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(
@@ -1141,6 +1152,8 @@ class MqttAvailability(Entity):
     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/const.py b/homeassistant/components/mqtt/const.py
index 5d1fe2e2505..acb24f4bdda 100644
--- a/homeassistant/components/mqtt/const.py
+++ b/homeassistant/components/mqtt/const.py
@@ -9,3 +9,6 @@ ATTR_DISCOVERY_TOPIC = "discovery_topic"
 CONF_STATE_TOPIC = "state_topic"
 PROTOCOL_311 = "3.1.1"
 DEFAULT_QOS = 0
+
+MQTT_CONNECTED = "mqtt_connected"
+MQTT_DISCONNECTED = "mqtt_disconnected"
diff --git a/tests/components/mqtt/test_alarm_control_panel.py b/tests/components/mqtt/test_alarm_control_panel.py
index 03e8133bde9..29ecb277928 100644
--- a/tests/components/mqtt/test_alarm_control_panel.py
+++ b/tests/components/mqtt/test_alarm_control_panel.py
@@ -19,6 +19,7 @@ from homeassistant.const import (
 )
 
 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 8c68fabf214..dccf8396c5a 100644
--- a/tests/components/mqtt/test_binary_sensor.py
+++ b/tests/components/mqtt/test_binary_sensor.py
@@ -18,6 +18,7 @@ from homeassistant.setup import async_setup_component
 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 11b846d4c38..98e85e2df7c 100644
--- a/tests/components/mqtt/test_camera.py
+++ b/tests/components/mqtt/test_camera.py
@@ -8,6 +8,7 @@ from homeassistant.components.mqtt.discovery import async_start
 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 30018c7c175..1736174d232 100644
--- a/tests/components/mqtt/test_climate.py
+++ b/tests/components/mqtt/test_climate.py
@@ -25,6 +25,7 @@ from homeassistant.components.climate.const import (
 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 d0ddc1d4830..d91c236680d 100644
--- a/tests/components/mqtt/test_common.py
+++ b/tests/components/mqtt/test_common.py
@@ -6,8 +6,10 @@ from unittest import mock
 
 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 @@ DEFAULT_CONFIG_DEVICE_INFO_MAC = {
 }
 
 
+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
+
+    mqtt_mock.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 eb758ebf93a..9ef74565206 100644
--- a/tests/components/mqtt/test_cover.py
+++ b/tests/components/mqtt/test_cover.py
@@ -30,6 +30,7 @@ from homeassistant.const import (
 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 7f6eb79e85e..bec771b6004 100644
--- a/tests/components/mqtt/test_fan.py
+++ b/tests/components/mqtt/test_fan.py
@@ -11,6 +11,7 @@ from homeassistant.const import (
 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 032a55edee4..9d02295db2e 100644
--- a/tests/components/mqtt/test_legacy_vacuum.py
+++ b/tests/components/mqtt/test_legacy_vacuum.py
@@ -23,6 +23,7 @@ from homeassistant.const import CONF_NAME, CONF_PLATFORM, STATE_OFF, STATE_ON
 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 f832e235915..0e8d0c19a19 100644
--- a/tests/components/mqtt/test_light.py
+++ b/tests/components/mqtt/test_light.py
@@ -162,6 +162,7 @@ import homeassistant.core as ha
 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 bb9e2afb0e5..19c531ecb30 100644
--- a/tests/components/mqtt/test_light_json.py
+++ b/tests/components/mqtt/test_light_json.py
@@ -102,6 +102,7 @@ import homeassistant.core as ha
 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 cb5aff40b4b..d1e27d4c516 100644
--- a/tests/components/mqtt/test_light_template.py
+++ b/tests/components/mqtt/test_light_template.py
@@ -39,6 +39,7 @@ import homeassistant.core as ha
 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 0e9a9af850f..abea215f837 100644
--- a/tests/components/mqtt/test_lock.py
+++ b/tests/components/mqtt/test_lock.py
@@ -12,6 +12,7 @@ from homeassistant.const import ATTR_ASSUMED_STATE, ATTR_ENTITY_ID
 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 d711b9e3bb8..be4b5012817 100644
--- a/tests/components/mqtt/test_sensor.py
+++ b/tests/components/mqtt/test_sensor.py
@@ -13,6 +13,7 @@ from homeassistant.setup import async_setup_component
 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 @@ async def test_force_update_enabled(hass, mqtt_mock):
     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 f77a1a11ca1..1a59e84c2be 100644
--- a/tests/components/mqtt/test_state_vacuum.py
+++ b/tests/components/mqtt/test_state_vacuum.py
@@ -33,6 +33,7 @@ from homeassistant.const import (
 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 da66e2a7f60..b90dfa68413 100644
--- a/tests/components/mqtt/test_switch.py
+++ b/tests/components/mqtt/test_switch.py
@@ -7,6 +7,7 @@ import homeassistant.core as ha
 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(