From fb67123d77292a7a1ffd50139938dd2816dfec05 Mon Sep 17 00:00:00 2001 From: Jan Bouwhuis Date: Fri, 9 Sep 2022 15:24:26 +0200 Subject: [PATCH] Clear MQTT discovery topic when a disabled entity is removed (#77757) * Cleanup discovery on entity removal * Add test * Cleanup and test * Test with clearing payload not unique id * Address comments * Tests cover and typing * Just pass hass * reuse code * Follow up comments revert changes to cover tests * Add test unique_id has priority over disabled * Update homeassistant/components/mqtt/__init__.py Co-authored-by: Erik Montnemery Co-authored-by: Erik Montnemery --- homeassistant/components/mqtt/__init__.py | 16 +- homeassistant/components/mqtt/const.py | 1 + homeassistant/components/mqtt/mixins.py | 51 ++++++- tests/components/mqtt/test_discovery.py | 170 ++++++++++++++++++++++ 4 files changed, 232 insertions(+), 6 deletions(-) diff --git a/homeassistant/components/mqtt/__init__.py b/homeassistant/components/mqtt/__init__.py index 3ec9a7e9d4e..842e5b6405f 100644 --- a/homeassistant/components/mqtt/__init__.py +++ b/homeassistant/components/mqtt/__init__.py @@ -20,7 +20,13 @@ from homeassistant.const import ( CONF_USERNAME, SERVICE_RELOAD, ) -from homeassistant.core import HassJob, HomeAssistant, ServiceCall, callback +from homeassistant.core import ( + CALLBACK_TYPE, + HassJob, + HomeAssistant, + ServiceCall, + callback, +) from homeassistant.exceptions import TemplateError, Unauthorized from homeassistant.helpers import ( config_validation as cv, @@ -68,6 +74,7 @@ from .const import ( # noqa: F401 CONFIG_ENTRY_IS_SETUP, DATA_MQTT, DATA_MQTT_CONFIG, + DATA_MQTT_DISCOVERY_REGISTRY_HOOKS, DATA_MQTT_RELOAD_DISPATCHERS, DATA_MQTT_RELOAD_ENTRY, DATA_MQTT_RELOAD_NEEDED, @@ -315,6 +322,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: # Bail out return False + hass.data[DATA_MQTT_DISCOVERY_REGISTRY_HOOKS] = {} hass.data[DATA_MQTT] = MQTT(hass, entry, conf) # Restore saved subscriptions if DATA_MQTT_SUBSCRIPTIONS_TO_RESTORE in hass.data: @@ -638,6 +646,12 @@ async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: hass.data[DATA_MQTT_RELOAD_ENTRY] = True # Reload the legacy yaml platform to make entities unavailable await async_reload_integration_platforms(hass, DOMAIN, RELOADABLE_PLATFORMS) + # Cleanup entity registry hooks + registry_hooks: dict[tuple, CALLBACK_TYPE] = hass.data[ + DATA_MQTT_DISCOVERY_REGISTRY_HOOKS + ] + while registry_hooks: + registry_hooks.popitem()[1]() # Wait for all ACKs and stop the loop await mqtt_client.async_disconnect() # Store remaining subscriptions to be able to restore or reload them diff --git a/homeassistant/components/mqtt/const.py b/homeassistant/components/mqtt/const.py index 0c711e097d5..c8af58862e0 100644 --- a/homeassistant/components/mqtt/const.py +++ b/homeassistant/components/mqtt/const.py @@ -33,6 +33,7 @@ CONF_TLS_VERSION = "tls_version" CONFIG_ENTRY_IS_SETUP = "mqtt_config_entry_is_setup" DATA_MQTT = "mqtt" DATA_MQTT_SUBSCRIPTIONS_TO_RESTORE = "mqtt_client_subscriptions" +DATA_MQTT_DISCOVERY_REGISTRY_HOOKS = "mqtt_discovery_registry_hooks" DATA_MQTT_CONFIG = "mqtt_config" MQTT_DATA_DEVICE_TRACKER_LEGACY = "mqtt_device_tracker_legacy" DATA_MQTT_RELOAD_DISPATCHERS = "mqtt_reload_dispatchers" diff --git a/homeassistant/components/mqtt/mixins.py b/homeassistant/components/mqtt/mixins.py index 75532f75c13..fddbe838303 100644 --- a/homeassistant/components/mqtt/mixins.py +++ b/homeassistant/components/mqtt/mixins.py @@ -28,7 +28,13 @@ from homeassistant.const import ( CONF_UNIQUE_ID, CONF_VALUE_TEMPLATE, ) -from homeassistant.core import Event, HomeAssistant, async_get_hass, callback +from homeassistant.core import ( + CALLBACK_TYPE, + Event, + HomeAssistant, + async_get_hass, + callback, +) from homeassistant.helpers import ( config_validation as cv, device_registry as dr, @@ -48,6 +54,7 @@ from homeassistant.helpers.entity import ( async_generate_entity_id, ) from homeassistant.helpers.entity_platform import AddEntitiesCallback +from homeassistant.helpers.event import async_track_entity_registry_updated_event from homeassistant.helpers.issue_registry import IssueSeverity, async_create_issue from homeassistant.helpers.json import json_loads from homeassistant.helpers.typing import ConfigType, DiscoveryInfoType @@ -64,6 +71,7 @@ from .const import ( CONF_TOPIC, DATA_MQTT, DATA_MQTT_CONFIG, + DATA_MQTT_DISCOVERY_REGISTRY_HOOKS, DATA_MQTT_RELOAD_DISPATCHERS, DATA_MQTT_RELOAD_ENTRY, DATA_MQTT_UPDATED_CONFIG, @@ -654,6 +662,17 @@ async def async_remove_discovery_payload(hass: HomeAssistant, discovery_data: di await async_publish(hass, discovery_topic, "", retain=True) +async def async_clear_discovery_topic_if_entity_removed( + hass: HomeAssistant, + discovery_data: dict[str, Any], + event: Event, +) -> None: + """Clear the discovery topic if the entity is removed.""" + if event.data["action"] == "remove": + # publish empty payload to config topic to avoid re-adding + await async_remove_discovery_payload(hass, discovery_data) + + class MqttDiscoveryDeviceUpdate: """Add support for auto discovery for platforms without an entity.""" @@ -787,7 +806,8 @@ class MqttDiscoveryUpdate(Entity): def __init__( self, - discovery_data: dict, + hass: HomeAssistant, + discovery_data: dict | None, discovery_update: Callable | None = None, ) -> None: """Initialize the discovery update mixin.""" @@ -795,6 +815,14 @@ class MqttDiscoveryUpdate(Entity): self._discovery_update = discovery_update self._remove_discovery_updated: Callable | None = None self._removed_from_hass = False + if discovery_data is None: + return + self._registry_hooks: dict[tuple, CALLBACK_TYPE] = hass.data[ + DATA_MQTT_DISCOVERY_REGISTRY_HOOKS + ] + discovery_hash: tuple[str, str] = discovery_data[ATTR_DISCOVERY_HASH] + if discovery_hash in self._registry_hooks: + self._registry_hooks.pop(discovery_hash)() async def async_added_to_hass(self) -> None: """Subscribe to discovery updates.""" @@ -857,7 +885,7 @@ class MqttDiscoveryUpdate(Entity): async def async_removed_from_registry(self) -> None: """Clear retained discovery topic in broker.""" - if not self._removed_from_hass: + if not self._removed_from_hass and self._discovery_data is not None: # Stop subscribing to discovery updates to not trigger when we clear the # discovery topic self._cleanup_discovery_on_remove() @@ -868,7 +896,20 @@ class MqttDiscoveryUpdate(Entity): @callback def add_to_platform_abort(self) -> None: """Abort adding an entity to a platform.""" - if self._discovery_data: + if self._discovery_data is not None: + discovery_hash: tuple = self._discovery_data[ATTR_DISCOVERY_HASH] + if self.registry_entry is not None: + self._registry_hooks[ + discovery_hash + ] = async_track_entity_registry_updated_event( + self.hass, + self.entity_id, + partial( + async_clear_discovery_topic_if_entity_removed, + self.hass, + self._discovery_data, + ), + ) stop_discovery_updates(self.hass, self._discovery_data) send_discovery_done(self.hass, self._discovery_data) super().add_to_platform_abort() @@ -976,7 +1017,7 @@ class MqttEntity( # Initialize mixin classes MqttAttributes.__init__(self, config) MqttAvailability.__init__(self, config) - MqttDiscoveryUpdate.__init__(self, discovery_data, self.discovery_update) + MqttDiscoveryUpdate.__init__(self, hass, discovery_data, self.discovery_update) MqttEntityDeviceInfo.__init__(self, config.get(CONF_DEVICE), config_entry) def _init_entity_id(self): diff --git a/tests/components/mqtt/test_discovery.py b/tests/components/mqtt/test_discovery.py index 29ca1f11743..c625d0a21f9 100644 --- a/tests/components/mqtt/test_discovery.py +++ b/tests/components/mqtt/test_discovery.py @@ -1,4 +1,5 @@ """The tests for the MQTT discovery.""" +import copy import json from pathlib import Path import re @@ -23,6 +24,8 @@ from homeassistant.const import ( import homeassistant.core as ha from homeassistant.setup import async_setup_component +from .test_common import help_test_unload_config_entry + from tests.common import ( MockConfigEntry, async_capture_events, @@ -1356,3 +1359,170 @@ async def test_mqtt_discovery_unsubscribe_once( await hass.async_block_till_done() await hass.async_block_till_done() mqtt_client_mock.unsubscribe.assert_called_once_with("comp/discovery/#") + + +@patch("homeassistant.components.mqtt.PLATFORMS", [Platform.SENSOR]) +async def test_clear_config_topic_disabled_entity( + hass, mqtt_mock_entry_no_yaml_config, device_reg, caplog +): + """Test the discovery topic is removed when a disabled entity is removed.""" + mqtt_mock = await mqtt_mock_entry_no_yaml_config() + # discover an entity that is not enabled by default + config = { + "name": "sbfspot_12345", + "state_topic": "homeassistant_test/sensor/sbfspot_0/sbfspot_12345/", + "unique_id": "sbfspot_12345", + "enabled_by_default": False, + "device": { + "identifiers": ["sbfspot_12345"], + "name": "sbfspot_12345", + "sw_version": "1.0", + "connections": [["mac", "12:34:56:AB:CD:EF"]], + }, + } + async_fire_mqtt_message( + hass, + "homeassistant/sensor/sbfspot_0/sbfspot_12345/config", + json.dumps(config), + ) + await hass.async_block_till_done() + # discover an entity that is not unique (part 1), will be added + config_not_unique1 = copy.deepcopy(config) + config_not_unique1["name"] = "sbfspot_12345_1" + config_not_unique1["unique_id"] = "not_unique" + config_not_unique1.pop("enabled_by_default") + async_fire_mqtt_message( + hass, + "homeassistant/sensor/sbfspot_0/sbfspot_12345_1/config", + json.dumps(config_not_unique1), + ) + # discover an entity that is not unique (part 2), will not be added + config_not_unique2 = copy.deepcopy(config_not_unique1) + config_not_unique2["name"] = "sbfspot_12345_2" + async_fire_mqtt_message( + hass, + "homeassistant/sensor/sbfspot_0/sbfspot_12345_2/config", + json.dumps(config_not_unique2), + ) + await hass.async_block_till_done() + assert "Platform mqtt does not generate unique IDs" in caplog.text + + assert hass.states.get("sensor.sbfspot_12345") is None # disabled + assert hass.states.get("sensor.sbfspot_12345_1") is not None # enabled + assert hass.states.get("sensor.sbfspot_12345_2") is None # not unique + + # Verify device is created + device_entry = device_reg.async_get_device(set(), {("mac", "12:34:56:AB:CD:EF")}) + assert device_entry is not None + + # Remove the device from the registry + device_reg.async_remove_device(device_entry.id) + await hass.async_block_till_done() + await hass.async_block_till_done() + + # Assert all valid discovery topics are cleared + assert mqtt_mock.async_publish.call_count == 2 + assert ( + call("homeassistant/sensor/sbfspot_0/sbfspot_12345/config", "", 0, True) + in mqtt_mock.async_publish.mock_calls + ) + assert ( + call("homeassistant/sensor/sbfspot_0/sbfspot_12345_1/config", "", 0, True) + in mqtt_mock.async_publish.mock_calls + ) + + +@patch("homeassistant.components.mqtt.PLATFORMS", [Platform.SENSOR]) +async def test_clean_up_registry_monitoring( + hass, mqtt_mock_entry_no_yaml_config, device_reg, tmp_path +): + """Test registry monitoring hook is removed after a reload.""" + await mqtt_mock_entry_no_yaml_config() + hooks: dict = hass.data[mqtt.const.DATA_MQTT_DISCOVERY_REGISTRY_HOOKS] + # discover an entity that is not enabled by default + config1 = { + "name": "sbfspot_12345", + "state_topic": "homeassistant_test/sensor/sbfspot_0/sbfspot_12345/", + "unique_id": "sbfspot_12345", + "enabled_by_default": False, + "device": { + "identifiers": ["sbfspot_12345"], + "name": "sbfspot_12345", + "sw_version": "1.0", + "connections": [["mac", "12:34:56:AB:CD:EF"]], + }, + } + # Publish it config + # Since it is not enabled_by_default the sensor will not be loaded + # it should register a hook for monitoring the entiry registry + async_fire_mqtt_message( + hass, + "homeassistant/sensor/sbfspot_0/sbfspot_12345/config", + json.dumps(config1), + ) + await hass.async_block_till_done() + assert len(hooks) == 1 + + # Publish it again no new monitor should be started + async_fire_mqtt_message( + hass, + "homeassistant/sensor/sbfspot_0/sbfspot_12345/config", + json.dumps(config1), + ) + await hass.async_block_till_done() + assert len(hooks) == 1 + + # Verify device is created + device_entry = device_reg.async_get_device(set(), {("mac", "12:34:56:AB:CD:EF")}) + assert device_entry is not None + + # Enload the entry + # The monitoring should be cleared + await help_test_unload_config_entry(hass, tmp_path, {}) + assert len(hooks) == 0 + + +@patch("homeassistant.components.mqtt.PLATFORMS", [Platform.SENSOR]) +async def test_unique_id_collission_has_priority( + hass, mqtt_mock_entry_no_yaml_config, entity_reg +): + """Test tehe unique_id collision detection has priority over registry disabled items.""" + await mqtt_mock_entry_no_yaml_config() + config = { + "name": "sbfspot_12345", + "state_topic": "homeassistant_test/sensor/sbfspot_0/sbfspot_12345/", + "unique_id": "sbfspot_12345", + "enabled_by_default": False, + "device": { + "identifiers": ["sbfspot_12345"], + "name": "sbfspot_12345", + "sw_version": "1.0", + "connections": [["mac", "12:34:56:AB:CD:EF"]], + }, + } + # discover an entity that is not unique and disabled by default (part 1), will be added + config_not_unique1 = copy.deepcopy(config) + config_not_unique1["name"] = "sbfspot_12345_1" + config_not_unique1["unique_id"] = "not_unique" + async_fire_mqtt_message( + hass, + "homeassistant/sensor/sbfspot_0/sbfspot_12345_1/config", + json.dumps(config_not_unique1), + ) + # discover an entity that is not unique (part 2), will not be added, and the registry entry is cleared + config_not_unique2 = copy.deepcopy(config_not_unique1) + config_not_unique2["name"] = "sbfspot_12345_2" + async_fire_mqtt_message( + hass, + "homeassistant/sensor/sbfspot_0/sbfspot_12345_2/config", + json.dumps(config_not_unique2), + ) + await hass.async_block_till_done() + + assert hass.states.get("sensor.sbfspot_12345_1") is None # not enabled + assert hass.states.get("sensor.sbfspot_12345_2") is None # not unique + + # Verify the first entity is created + assert entity_reg.async_get("sensor.sbfspot_12345_1") is not None + # Verify the second entity is not created because it is not unique + assert entity_reg.async_get("sensor.sbfspot_12345_2") is None