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 <erik@montnemery.com>

Co-authored-by: Erik Montnemery <erik@montnemery.com>
This commit is contained in:
Jan Bouwhuis 2022-09-09 15:24:26 +02:00 committed by GitHub
parent 0a143ac596
commit fb67123d77
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
4 changed files with 232 additions and 6 deletions

View file

@ -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

View file

@ -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"

View file

@ -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):

View file

@ -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