Fix MQTT expire_after effects after reloading (#65359)

* Cleanup sensor expire triggers after reload

* fix test binary_sensor

* Also trigger cleanup parent classes

* Restore an expiring state after a reload

* correct discovery_update

* restore expiring state with remaining time

* Update homeassistant/components/mqtt/binary_sensor.py

description

Co-authored-by: Erik Montnemery <erik@montnemery.com>

* Log remaining time

* Move check

* check and tests reload

* remove self.async_write_ha_state()

Co-authored-by: Erik Montnemery <erik@montnemery.com>

homeassistant/components/mqtt/binary_sensor.py

@@ -20,6 +20,8 @@ from homeassistant.const import (
     CONF_PAYLOAD_OFF,
     CONF_PAYLOAD_ON,
     CONF_VALUE_TEMPLATE,
+    STATE_UNAVAILABLE,
+    STATE_UNKNOWN,
 )
 from homeassistant.core import HomeAssistant, callback
 import homeassistant.helpers.config_validation as cv
@@ -27,6 +29,7 @@ from homeassistant.helpers.entity_platform import AddEntitiesCallback
 import homeassistant.helpers.event as evt
 from homeassistant.helpers.event import async_track_point_in_utc_time
 from homeassistant.helpers.reload import async_setup_reload_service
+from homeassistant.helpers.restore_state import RestoreEntity
 from homeassistant.helpers.typing import ConfigType, DiscoveryInfoType
 from homeassistant.util import dt as dt_util
 
@@ -96,7 +99,7 @@ async def _async_setup_entity(
     async_add_entities([MqttBinarySensor(hass, config, config_entry, discovery_data)])
 
 
-class MqttBinarySensor(MqttEntity, BinarySensorEntity):
+class MqttBinarySensor(MqttEntity, BinarySensorEntity, RestoreEntity):
     """Representation a binary sensor that is updated by MQTT."""
 
     _entity_id_format = binary_sensor.ENTITY_ID_FORMAT
@@ -114,6 +117,42 @@ class MqttBinarySensor(MqttEntity, BinarySensorEntity):
 
         MqttEntity.__init__(self, hass, config, config_entry, discovery_data)
 
+    async def async_added_to_hass(self) -> None:
+        """Restore state for entities with expire_after set."""
+        await super().async_added_to_hass()
+        if (
+            (expire_after := self._config.get(CONF_EXPIRE_AFTER)) is not None
+            and expire_after > 0
+            and (last_state := await self.async_get_last_state()) is not None
+            and last_state.state not in [STATE_UNKNOWN, STATE_UNAVAILABLE]
+        ):
+            expiration_at = last_state.last_changed + timedelta(seconds=expire_after)
+            if expiration_at < (time_now := dt_util.utcnow()):
+                # Skip reactivating the binary_sensor
+                _LOGGER.debug("Skip state recovery after reload for %s", self.entity_id)
+                return
+            self._expired = False
+            self._state = last_state.state
+
+            self._expiration_trigger = async_track_point_in_utc_time(
+                self.hass, self._value_is_expired, expiration_at
+            )
+            _LOGGER.debug(
+                "State recovered after reload for %s, remaining time before expiring %s",
+                self.entity_id,
+                expiration_at - time_now,
+            )
+
+    async def async_will_remove_from_hass(self) -> None:
+        """Remove exprire triggers."""
+        # Clean up expire triggers
+        if self._expiration_trigger:
+            _LOGGER.debug("Clean up expire after trigger for %s", self.entity_id)
+            self._expiration_trigger()
+            self._expiration_trigger = None
+            self._expired = False
+        await MqttEntity.async_will_remove_from_hass(self)
+
     @staticmethod
     def config_schema():
         """Return the config schema."""

homeassistant/components/mqtt/mixins.py

@@ -524,6 +524,11 @@ class MqttDiscoveryUpdate(Entity):
     async def async_removed_from_registry(self) -> None:
         """Clear retained discovery topic in broker."""
         if not self._removed_from_hass:
+            # Stop subscribing to discovery updates to not trigger when we clear the
+            # discovery topic
+            self._cleanup_discovery_on_remove()
+
+            # Clear the discovery topic so the entity is not rediscovered after a restart
             discovery_topic = self._discovery_data[ATTR_DISCOVERY_TOPIC]
             publish(self.hass, discovery_topic, "", retain=True)
 

homeassistant/components/mqtt/sensor.py

@@ -23,12 +23,15 @@ from homeassistant.const import (
     CONF_NAME,
     CONF_UNIT_OF_MEASUREMENT,
     CONF_VALUE_TEMPLATE,
+    STATE_UNAVAILABLE,
+    STATE_UNKNOWN,
 )
 from homeassistant.core import HomeAssistant, callback
 import homeassistant.helpers.config_validation as cv
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.event import async_track_point_in_utc_time
 from homeassistant.helpers.reload import async_setup_reload_service
+from homeassistant.helpers.restore_state import RestoreEntity
 from homeassistant.helpers.typing import ConfigType, DiscoveryInfoType
 from homeassistant.util import dt as dt_util
 
@@ -140,7 +143,7 @@ async def _async_setup_entity(
     async_add_entities([MqttSensor(hass, config, config_entry, discovery_data)])
 
 
-class MqttSensor(MqttEntity, SensorEntity):
+class MqttSensor(MqttEntity, SensorEntity, RestoreEntity):
     """Representation of a sensor that can be updated using MQTT."""
 
     _entity_id_format = ENTITY_ID_FORMAT
@@ -160,6 +163,42 @@ class MqttSensor(MqttEntity, SensorEntity):
 
         MqttEntity.__init__(self, hass, config, config_entry, discovery_data)
 
+    async def async_added_to_hass(self) -> None:
+        """Restore state for entities with expire_after set."""
+        await super().async_added_to_hass()
+        if (
+            (expire_after := self._config.get(CONF_EXPIRE_AFTER)) is not None
+            and expire_after > 0
+            and (last_state := await self.async_get_last_state()) is not None
+            and last_state.state not in [STATE_UNKNOWN, STATE_UNAVAILABLE]
+        ):
+            expiration_at = last_state.last_changed + timedelta(seconds=expire_after)
+            if expiration_at < (time_now := dt_util.utcnow()):
+                # Skip reactivating the sensor
+                _LOGGER.debug("Skip state recovery after reload for %s", self.entity_id)
+                return
+            self._expired = False
+            self._state = last_state.state
+
+            self._expiration_trigger = async_track_point_in_utc_time(
+                self.hass, self._value_is_expired, expiration_at
+            )
+            _LOGGER.debug(
+                "State recovered after reload for %s, remaining time before expiring %s",
+                self.entity_id,
+                expiration_at - time_now,
+            )
+
+    async def async_will_remove_from_hass(self) -> None:
+        """Remove exprire triggers."""
+        # Clean up expire triggers
+        if self._expiration_trigger:
+            _LOGGER.debug("Clean up expire after trigger for %s", self.entity_id)
+            self._expiration_trigger()
+            self._expiration_trigger = None
+            self._expired = False
+        await MqttEntity.async_will_remove_from_hass(self)
+
     @staticmethod
     def config_schema():
         """Return the config schema."""

tests/components/mqtt/test_binary_sensor.py

@@ -36,6 +36,7 @@ from .test_common import (
     help_test_entity_device_info_with_identifier,
     help_test_entity_id_update_discovery_update,
     help_test_entity_id_update_subscriptions,
+    help_test_reload_with_config,
     help_test_reloadable,
     help_test_setting_attribute_via_mqtt_json_message,
     help_test_setting_attribute_with_template,
@@ -44,7 +45,11 @@ from .test_common import (
     help_test_update_with_json_attrs_not_dict,
 )
 
-from tests.common import async_fire_mqtt_message, async_fire_time_changed
+from tests.common import (
+    assert_setup_component,
+    async_fire_mqtt_message,
+    async_fire_time_changed,
+)
 
 DEFAULT_CONFIG = {
     binary_sensor.DOMAIN: {
@@ -872,3 +877,87 @@ async def test_reloadable(hass, mqtt_mock, caplog, tmp_path):
     domain = binary_sensor.DOMAIN
     config = DEFAULT_CONFIG[domain]
     await help_test_reloadable(hass, mqtt_mock, caplog, tmp_path, domain, config)
+
+
+async def test_cleanup_triggers_and_restoring_state(
+    hass, mqtt_mock, caplog, tmp_path, freezer
+):
+    """Test cleanup old triggers at reloading and restoring the state."""
+    domain = binary_sensor.DOMAIN
+    config1 = copy.deepcopy(DEFAULT_CONFIG[domain])
+    config1["name"] = "test1"
+    config1["expire_after"] = 30
+    config1["state_topic"] = "test-topic1"
+    config2 = copy.deepcopy(DEFAULT_CONFIG[domain])
+    config2["name"] = "test2"
+    config2["expire_after"] = 5
+    config2["state_topic"] = "test-topic2"
+
+    freezer.move_to("2022-02-02 12:01:00+01:00")
+
+    assert await async_setup_component(
+        hass,
+        binary_sensor.DOMAIN,
+        {binary_sensor.DOMAIN: [config1, config2]},
+    )
+    await hass.async_block_till_done()
+    async_fire_mqtt_message(hass, "test-topic1", "ON")
+    state = hass.states.get("binary_sensor.test1")
+    assert state.state == "on"
+
+    async_fire_mqtt_message(hass, "test-topic2", "ON")
+    state = hass.states.get("binary_sensor.test2")
+    assert state.state == "on"
+
+    freezer.move_to("2022-02-02 12:01:10+01:00")
+
+    await help_test_reload_with_config(
+        hass, caplog, tmp_path, domain, [config1, config2]
+    )
+    assert "Clean up expire after trigger for binary_sensor.test1" in caplog.text
+    assert "Clean up expire after trigger for binary_sensor.test2" not in caplog.text
+    assert (
+        "State recovered after reload for binary_sensor.test1, remaining time before expiring"
+        in caplog.text
+    )
+    assert "State recovered after reload for binary_sensor.test2" not in caplog.text
+
+    state = hass.states.get("binary_sensor.test1")
+    assert state.state == "on"
+
+    state = hass.states.get("binary_sensor.test2")
+    assert state.state == STATE_UNAVAILABLE
+
+    async_fire_mqtt_message(hass, "test-topic1", "OFF")
+    state = hass.states.get("binary_sensor.test1")
+    assert state.state == "off"
+
+    async_fire_mqtt_message(hass, "test-topic2", "OFF")
+    state = hass.states.get("binary_sensor.test2")
+    assert state.state == "off"
+
+
+async def test_skip_restoring_state_with_over_due_expire_trigger(
+    hass, mqtt_mock, caplog, freezer
+):
+    """Test restoring a state with over due expire timer."""
+
+    freezer.move_to("2022-02-02 12:02:00+01:00")
+    domain = binary_sensor.DOMAIN
+    config3 = copy.deepcopy(DEFAULT_CONFIG[domain])
+    config3["name"] = "test3"
+    config3["expire_after"] = 10
+    config3["state_topic"] = "test-topic3"
+    fake_state = ha.State(
+        "binary_sensor.test3",
+        "on",
+        {},
+        last_changed=datetime.fromisoformat("2022-02-02 12:01:35+01:00"),
+    )
+    with patch(
+        "homeassistant.helpers.restore_state.RestoreEntity.async_get_last_state",
+        return_value=fake_state,
+    ), assert_setup_component(1, domain):
+        assert await async_setup_component(hass, domain, {domain: config3})
+        await hass.async_block_till_done()
+    assert "Skip state recovery after reload for binary_sensor.test3" in caplog.text

tests/components/mqtt/test_common.py

@@ -1525,6 +1525,25 @@ async def help_test_publishing_with_custom_encoding(
     mqtt_mock.async_publish.reset_mock()
 
 
+async def help_test_reload_with_config(hass, caplog, tmp_path, domain, config):
+    """Test reloading with supplied config."""
+    new_yaml_config_file = tmp_path / "configuration.yaml"
+    new_yaml_config = yaml.dump({domain: config})
+    new_yaml_config_file.write_text(new_yaml_config)
+    assert new_yaml_config_file.read_text() == new_yaml_config
+
+    with patch.object(hass_config, "YAML_CONFIG_FILE", new_yaml_config_file):
+        await hass.services.async_call(
+            "mqtt",
+            SERVICE_RELOAD,
+            {},
+            blocking=True,
+        )
+        await hass.async_block_till_done()
+
+    assert "<Event event_mqtt_reloaded[L]>" in caplog.text
+
+
 async def help_test_reloadable(hass, mqtt_mock, caplog, tmp_path, domain, config):
     """Test reloading an MQTT platform."""
     # Create and test an old config of 2 entities based on the config supplied
@@ -1549,21 +1568,10 @@ async def help_test_reloadable(hass, mqtt_mock, caplog, tmp_path, domain, config
     new_config_2["name"] = "test_new_2"
     new_config_3 = copy.deepcopy(config)
     new_config_3["name"] = "test_new_3"
-    new_yaml_config_file = tmp_path / "configuration.yaml"
-    new_yaml_config = yaml.dump({domain: [new_config_1, new_config_2, new_config_3]})
-    new_yaml_config_file.write_text(new_yaml_config)
-    assert new_yaml_config_file.read_text() == new_yaml_config
 
-    with patch.object(hass_config, "YAML_CONFIG_FILE", new_yaml_config_file):
+    await help_test_reload_with_config(
-        await hass.services.async_call(
+        hass, caplog, tmp_path, domain, [new_config_1, new_config_2, new_config_3]
-            "mqtt",
+    )
-            SERVICE_RELOAD,
-            {},
-            blocking=True,
-        )
-        await hass.async_block_till_done()
-
-    assert "<Event event_mqtt_reloaded[L]>" in caplog.text
 
     assert len(hass.states.async_all(domain)) == 3
 

tests/components/mqtt/test_sensor.py

@@ -43,6 +43,7 @@ from .test_common import (
     help_test_entity_disabled_by_default,
     help_test_entity_id_update_discovery_update,
     help_test_entity_id_update_subscriptions,
+    help_test_reload_with_config,
     help_test_reloadable,
     help_test_setting_attribute_via_mqtt_json_message,
     help_test_setting_attribute_with_template,
@@ -52,7 +53,11 @@ from .test_common import (
     help_test_update_with_json_attrs_not_dict,
 )
 
-from tests.common import async_fire_mqtt_message, async_fire_time_changed
+from tests.common import (
+    assert_setup_component,
+    async_fire_mqtt_message,
+    async_fire_time_changed,
+)
 
 DEFAULT_CONFIG = {
     sensor.DOMAIN: {"platform": "mqtt", "name": "test", "state_topic": "test-topic"}
@@ -935,6 +940,92 @@ async def test_reloadable(hass, mqtt_mock, caplog, tmp_path):
     await help_test_reloadable(hass, mqtt_mock, caplog, tmp_path, domain, config)
 
 
+async def test_cleanup_triggers_and_restoring_state(
+    hass, mqtt_mock, caplog, tmp_path, freezer
+):
+    """Test cleanup old triggers at reloading and restoring the state."""
+    domain = sensor.DOMAIN
+    config1 = copy.deepcopy(DEFAULT_CONFIG[domain])
+    config1["name"] = "test1"
+    config1["expire_after"] = 30
+    config1["state_topic"] = "test-topic1"
+    config2 = copy.deepcopy(DEFAULT_CONFIG[domain])
+    config2["name"] = "test2"
+    config2["expire_after"] = 5
+    config2["state_topic"] = "test-topic2"
+
+    freezer.move_to("2022-02-02 12:01:00+01:00")
+
+    assert await async_setup_component(
+        hass,
+        domain,
+        {domain: [config1, config2]},
+    )
+    await hass.async_block_till_done()
+    async_fire_mqtt_message(hass, "test-topic1", "100")
+    state = hass.states.get("sensor.test1")
+    assert state.state == "100"
+
+    async_fire_mqtt_message(hass, "test-topic2", "200")
+    state = hass.states.get("sensor.test2")
+    assert state.state == "200"
+
+    freezer.move_to("2022-02-02 12:01:10+01:00")
+
+    await help_test_reload_with_config(
+        hass, caplog, tmp_path, domain, [config1, config2]
+    )
+    await hass.async_block_till_done()
+
+    assert "Clean up expire after trigger for sensor.test1" in caplog.text
+    assert "Clean up expire after trigger for sensor.test2" not in caplog.text
+    assert (
+        "State recovered after reload for sensor.test1, remaining time before expiring"
+        in caplog.text
+    )
+    assert "State recovered after reload for sensor.test2" not in caplog.text
+
+    state = hass.states.get("sensor.test1")
+    assert state.state == "100"
+
+    state = hass.states.get("sensor.test2")
+    assert state.state == STATE_UNAVAILABLE
+
+    async_fire_mqtt_message(hass, "test-topic1", "101")
+    state = hass.states.get("sensor.test1")
+    assert state.state == "101"
+
+    async_fire_mqtt_message(hass, "test-topic2", "201")
+    state = hass.states.get("sensor.test2")
+    assert state.state == "201"
+
+
+async def test_skip_restoring_state_with_over_due_expire_trigger(
+    hass, mqtt_mock, caplog, freezer
+):
+    """Test restoring a state with over due expire timer."""
+
+    freezer.move_to("2022-02-02 12:02:00+01:00")
+    domain = sensor.DOMAIN
+    config3 = copy.deepcopy(DEFAULT_CONFIG[domain])
+    config3["name"] = "test3"
+    config3["expire_after"] = 10
+    config3["state_topic"] = "test-topic3"
+    fake_state = ha.State(
+        "sensor.test3",
+        "300",
+        {},
+        last_changed=datetime.fromisoformat("2022-02-02 12:01:35+01:00"),
+    )
+    with patch(
+        "homeassistant.helpers.restore_state.RestoreEntity.async_get_last_state",
+        return_value=fake_state,
+    ), assert_setup_component(1, domain):
+        assert await async_setup_component(hass, domain, {domain: config3})
+        await hass.async_block_till_done()
+    assert "Skip state recovery after reload for sensor.test3" in caplog.text
+
+
 @pytest.mark.parametrize(
     "topic,value,attribute,attribute_value",
     [