Refactor MQTT discovery (#67966)

* Proof of concept

* remove notify platform

* remove loose test

* Add rework from #67912 (#1)

* Move notify serviceupdater to Mixins

* Move tag discovery handler to Mixins

* fix tests

* Add typing for async_load_platform_helper

* Add add entry unload support for notify platform

* Simplify discovery updates

* Remove not needed extra logic

* Cleanup inrelevant or duplicate code

* reuse update_device and move to mixins

* Remove notify platform

* revert changes to notify platform

* Rename update class

* unify tag entry setup

* Use shared code for device_trigger `update_device`

* PoC shared dispatcher for device_trigger

* Fix bugs

* Improve typing - remove async_update

* Unload config_entry and tests

* Release dispatcher after setup and deduplicate

* closures to methods, revert `in` to `=`, updates

* Re-add update support for tag platform

* Re-add update support for device-trigger platform

* Cleanup rediscovery code revert related changes

* Undo discovery code shift

* Update homeassistant/components/mqtt/mixins.py

Co-authored-by: Erik Montnemery <[email protected]>

* Update homeassistant/components/mqtt/device_trigger.py

Co-authored-by: Erik Montnemery <[email protected]>

* Update homeassistant/components/mqtt/mixins.py

Co-authored-by: Erik Montnemery <[email protected]>

* revert doc string changes

* move conditions

* typing and check config_entry_id

* Update homeassistant/components/mqtt/mixins.py

Co-authored-by: Erik Montnemery <[email protected]>

* cleanup not used attribute

* Remove entry_unload code and tests

* update  comment

* add second comment

Co-authored-by: Erik Montnemery <[email protected]>

Author: jbouwh

homeassistant/components/mqtt/device_trigger.py

@@ -3,7 +3,7 @@
 
 from collections.abc import Callable
 import logging
-from typing import Any
+from typing import Any, cast
 
 import attr
 import voluptuous as vol
@@ -13,6 +13,7 @@
     AutomationTriggerInfo,
 )
 from homeassistant.components.device_automation import DEVICE_TRIGGER_BASE_SCHEMA
+from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import (
     CONF_DEVICE,
     CONF_DEVICE_ID,
@@ -23,30 +24,19 @@
 )
 from homeassistant.core import CALLBACK_TYPE, HomeAssistant, callback
 from homeassistant.exceptions import HomeAssistantError
-from homeassistant.helpers import config_validation as cv, device_registry as dr
-from homeassistant.helpers.dispatcher import (
-    async_dispatcher_connect,
-    async_dispatcher_send,
-)
+from homeassistant.helpers import config_validation as cv
+from homeassistant.helpers.dispatcher import async_dispatcher_send
 from homeassistant.helpers.typing import ConfigType
 
 from . import debug_info, trigger as mqtt_trigger
 from .. import mqtt
-from .const import (
-    ATTR_DISCOVERY_HASH,
-    ATTR_DISCOVERY_TOPIC,
-    CONF_PAYLOAD,
-    CONF_QOS,
-    CONF_TOPIC,
-    DOMAIN,
-)
-from .discovery import MQTT_DISCOVERY_DONE, MQTT_DISCOVERY_UPDATED, clear_discovery_hash
+from .const import ATTR_DISCOVERY_HASH, CONF_PAYLOAD, CONF_QOS, CONF_TOPIC, DOMAIN
+from .discovery import MQTT_DISCOVERY_DONE
 from .mixins import (
-    CONF_CONNECTIONS,
-    CONF_IDENTIFIERS,
     MQTT_ENTITY_DEVICE_INFO_SCHEMA,
-    cleanup_device_registry,
-    device_info_from_config,
+    MqttDiscoveryDeviceUpdate,
+    send_discovery_done,
+    update_device,
 )
 
 _LOGGER = logging.getLogger(__name__)
@@ -89,6 +79,8 @@
 
 DEVICE_TRIGGERS = "mqtt_device_triggers"
 
+LOG_NAME = "Device trigger"
+
 
 @attr.s(slots=True)
 class TriggerInstance:
@@ -99,7 +91,7 @@ class TriggerInstance:
     trigger: Trigger = attr.ib()
     remove: CALLBACK_TYPE | None = attr.ib(default=None)
 
-    async def async_attach_trigger(self):
+    async def async_attach_trigger(self) -> None:
         """Attach MQTT trigger."""
         mqtt_config = {
             mqtt_trigger.CONF_PLATFORM: mqtt.DOMAIN,
@@ -132,14 +124,15 @@ class Trigger:
     hass: HomeAssistant = attr.ib()
     payload: str | None = attr.ib()
     qos: int | None = attr.ib()
-    remove_signal: Callable[[], None] | None = attr.ib()
     subtype: str = attr.ib()
     topic: str | None = attr.ib()
     type: str = attr.ib()
     value_template: str | None = attr.ib()
     trigger_instances: list[TriggerInstance] = attr.ib(factory=list)
 
-    async def add_trigger(self, action, automation_info):
+    async def add_trigger(
+        self, action: AutomationActionType, automation_info: AutomationTriggerInfo
+    ) -> Callable:
         """Add MQTT trigger."""
         instance = TriggerInstance(action, automation_info, self)
         self.trigger_instances.append(instance)
@@ -160,9 +153,8 @@ def async_remove() -> None:
 
         return async_remove
 
-    async def update_trigger(self, config, discovery_hash, remove_signal):
+    async def update_trigger(self, config: ConfigType) -> None:
         """Update MQTT device trigger."""
-        self.remove_signal = remove_signal
         self.type = config[CONF_TYPE]
         self.subtype = config[CONF_SUBTYPE]
         self.payload = config[CONF_PAYLOAD]
@@ -178,7 +170,7 @@ async def update_trigger(self, config, discovery_hash, remove_signal):
             for trig in self.trigger_instances:
                 await trig.async_attach_trigger()
 
-    def detach_trigger(self):
+    def detach_trigger(self) -> None:
         """Remove MQTT device trigger."""
         # Mark trigger as unknown
         self.topic = None
@@ -190,110 +182,110 @@ def detach_trigger(self):
                 trig.remove = None
 
 
-def _update_device(hass, config_entry, config):
-    """Update device registry."""
-    device_registry = dr.async_get(hass)
-    config_entry_id = config_entry.entry_id
-    device_info = device_info_from_config(config[CONF_DEVICE])
-
-    if config_entry_id is not None and device_info is not None:
-        device_info["config_entry_id"] = config_entry_id
-        device_registry.async_get_or_create(**device_info)
-
+class MqttDeviceTrigger(MqttDiscoveryDeviceUpdate):
+    """Setup a MQTT device trigger with auto discovery."""
+
+    def __init__(
+        self,
+        hass: HomeAssistant,
+        config: ConfigType,
+        device_id: str,
+        discovery_data: dict,
+        config_entry: ConfigEntry,
+    ) -> None:
+        """Initialize."""
+        self._config = config
+        self._config_entry = config_entry
+        self.device_id = device_id
+        self.discovery_data = discovery_data
+        self.hass = hass
+
+        MqttDiscoveryDeviceUpdate.__init__(
+            self,
+            hass,
+            discovery_data,
+            device_id,
+            config_entry,
+            LOG_NAME,
+        )
 
-async def async_setup_trigger(hass, config, config_entry, discovery_data):
-    """Set up the MQTT device trigger."""
-    config = TRIGGER_DISCOVERY_SCHEMA(config)
-    discovery_hash = discovery_data[ATTR_DISCOVERY_HASH]
-    discovery_id = discovery_hash[1]
-    remove_signal = None
+    async def async_setup(self) -> None:
+        """Initialize the device trigger."""
+        discovery_hash = self.discovery_data[ATTR_DISCOVERY_HASH]
+        discovery_id = discovery_hash[1]
+        if discovery_id not in self.hass.data.setdefault(DEVICE_TRIGGERS, {}):
+            self.hass.data[DEVICE_TRIGGERS][discovery_id] = Trigger(
+                hass=self.hass,
+                device_id=self.device_id,
+                discovery_data=self.discovery_data,
+                type=self._config[CONF_TYPE],
+                subtype=self._config[CONF_SUBTYPE],
+                topic=self._config[CONF_TOPIC],
+                payload=self._config[CONF_PAYLOAD],
+                qos=self._config[CONF_QOS],
+                value_template=self._config[CONF_VALUE_TEMPLATE],
+            )
+        else:
+            await self.hass.data[DEVICE_TRIGGERS][discovery_id].update_trigger(
+                self._config
+            )
+        debug_info.add_trigger_discovery_data(
+            self.hass, discovery_hash, self.discovery_data, self.device_id
+        )
 
-    async def discovery_update(payload):
-        """Handle discovery update."""
-        _LOGGER.info(
-            "Got update for trigger with hash: %s '%s'", discovery_hash, payload
+    async def async_update(self, discovery_data: dict) -> None:
+        """Handle MQTT device trigger discovery updates."""
+        discovery_hash = self.discovery_data[ATTR_DISCOVERY_HASH]
+        discovery_id = discovery_hash[1]
+        debug_info.update_trigger_discovery_data(
+            self.hass, discovery_hash, discovery_data
         )
-        if not payload:
-            # Empty payload: Remove trigger
+        config = TRIGGER_DISCOVERY_SCHEMA(discovery_data)
+        update_device(self.hass, self._config_entry, config)
+        device_trigger: Trigger = self.hass.data[DEVICE_TRIGGERS][discovery_id]
+        await device_trigger.update_trigger(config)
+
+    async def async_tear_down(self) -> None:
+        """Cleanup device trigger."""
+        discovery_hash = self.discovery_data[ATTR_DISCOVERY_HASH]
+        discovery_id = discovery_hash[1]
+        if discovery_id in self.hass.data[DEVICE_TRIGGERS]:
             _LOGGER.info("Removing trigger: %s", discovery_hash)
-            debug_info.remove_trigger_discovery_data(hass, discovery_hash)
-            if discovery_id in hass.data[DEVICE_TRIGGERS]:
-                device_trigger = hass.data[DEVICE_TRIGGERS][discovery_id]
-                device_trigger.detach_trigger()
-                clear_discovery_hash(hass, discovery_hash)
-                remove_signal()
-                await cleanup_device_registry(hass, device.id, config_entry.entry_id)
-        else:
-            # Non-empty payload: Update trigger
-            _LOGGER.info("Updating trigger: %s", discovery_hash)
-            debug_info.update_trigger_discovery_data(hass, discovery_hash, payload)
-            config = TRIGGER_DISCOVERY_SCHEMA(payload)
-            _update_device(hass, config_entry, config)
-            device_trigger = hass.data[DEVICE_TRIGGERS][discovery_id]
-            await device_trigger.update_trigger(config, discovery_hash, remove_signal)
-        async_dispatcher_send(hass, MQTT_DISCOVERY_DONE.format(discovery_hash), None)
-
-    remove_signal = async_dispatcher_connect(
-        hass, MQTT_DISCOVERY_UPDATED.format(discovery_hash), discovery_update
-    )
+            trigger: Trigger = self.hass.data[DEVICE_TRIGGERS][discovery_id]
+            trigger.detach_trigger()
+            debug_info.remove_trigger_discovery_data(self.hass, discovery_hash)
 
-    _update_device(hass, config_entry, config)
 
-    device_registry = dr.async_get(hass)
-    device = device_registry.async_get_device(
-        {(DOMAIN, id_) for id_ in config[CONF_DEVICE][CONF_IDENTIFIERS]},
-        {tuple(x) for x in config[CONF_DEVICE][CONF_CONNECTIONS]},
-    )
+async def async_setup_trigger(
+    hass, config: ConfigType, config_entry: ConfigEntry, discovery_data: dict
+) -> None:
+    """Set up the MQTT device trigger."""
+    config = TRIGGER_DISCOVERY_SCHEMA(config)
+    discovery_hash = discovery_data[ATTR_DISCOVERY_HASH]
 
-    if device is None:
+    if (device_id := update_device(hass, config_entry, config)) is None:
         async_dispatcher_send(hass, MQTT_DISCOVERY_DONE.format(discovery_hash), None)
         return
 
-    if DEVICE_TRIGGERS not in hass.data:
-        hass.data[DEVICE_TRIGGERS] = {}
-    if discovery_id not in hass.data[DEVICE_TRIGGERS]:
-        hass.data[DEVICE_TRIGGERS][discovery_id] = Trigger(
-            hass=hass,
-            device_id=device.id,
-            discovery_data=discovery_data,
-            type=config[CONF_TYPE],
-            subtype=config[CONF_SUBTYPE],
-            topic=config[CONF_TOPIC],
-            payload=config[CONF_PAYLOAD],
-            qos=config[CONF_QOS],
-            remove_signal=remove_signal,
-            value_template=config[CONF_VALUE_TEMPLATE],
-        )
-    else:
-        await hass.data[DEVICE_TRIGGERS][discovery_id].update_trigger(
-            config, discovery_hash, remove_signal
-        )
-    debug_info.add_trigger_discovery_data(
-        hass, discovery_hash, discovery_data, device.id
+    mqtt_device_trigger = MqttDeviceTrigger(
+        hass, config, device_id, discovery_data, config_entry
     )
-
-    async_dispatcher_send(hass, MQTT_DISCOVERY_DONE.format(discovery_hash), None)
+    await mqtt_device_trigger.async_setup()
+    send_discovery_done(hass, discovery_data)
 
 
-async def async_removed_from_device(hass: HomeAssistant, device_id: str):
+async def async_removed_from_device(hass: HomeAssistant, device_id: str) -> None:
     """Handle Mqtt removed from a device."""
     triggers = await async_get_triggers(hass, device_id)
     for trig in triggers:
-        device_trigger = hass.data[DEVICE_TRIGGERS].pop(trig[CONF_DISCOVERY_ID])
+        device_trigger: Trigger = hass.data[DEVICE_TRIGGERS].pop(
+            trig[CONF_DISCOVERY_ID]
+        )
         if device_trigger:
-            discovery_hash = device_trigger.discovery_data[ATTR_DISCOVERY_HASH]
-            discovery_topic = device_trigger.discovery_data[ATTR_DISCOVERY_TOPIC]
-
-            debug_info.remove_trigger_discovery_data(hass, discovery_hash)
             device_trigger.detach_trigger()
-            clear_discovery_hash(hass, discovery_hash)
-            device_trigger.remove_signal()
-            mqtt.publish(
-                hass,
-                discovery_topic,
-                "",
-                retain=True,
-            )
+            discovery_data = cast(dict, device_trigger.discovery_data)
+            discovery_hash = discovery_data[ATTR_DISCOVERY_HASH]
+            debug_info.remove_trigger_discovery_data(hass, discovery_hash)
 
 
 async def async_get_triggers(
@@ -328,8 +320,7 @@ async def async_attach_trigger(
     automation_info: AutomationTriggerInfo,
 ) -> CALLBACK_TYPE:
     """Attach a trigger."""
-    if DEVICE_TRIGGERS not in hass.data:
-        hass.data[DEVICE_TRIGGERS] = {}
+    hass.data.setdefault(DEVICE_TRIGGERS, {})
     device_id = config[CONF_DEVICE_ID]
     discovery_id = config[CONF_DISCOVERY_ID]
 
@@ -338,7 +329,6 @@ async def async_attach_trigger(
             hass=hass,
             device_id=device_id,
             discovery_data=None,
-            remove_signal=None,
             type=config[CONF_TYPE],
             subtype=config[CONF_SUBTYPE],
             topic=None,

homeassistant/components/mqtt/discovery.py

@@ -1,4 +1,6 @@
 """Support for MQTT discovery."""
+from __future__ import annotations
+
 import asyncio
 from collections import deque
 import functools
@@ -73,20 +75,22 @@
 TOPIC_BASE = "~"
 
 
-def clear_discovery_hash(hass, discovery_hash):
+class MQTTConfig(dict):
+    """Dummy class to allow adding attributes."""
+
+    discovery_data: dict
+
+
+def clear_discovery_hash(hass: HomeAssistant, discovery_hash: tuple) -> None:
     """Clear entry in ALREADY_DISCOVERED list."""
     del hass.data[ALREADY_DISCOVERED][discovery_hash]
 
 
-def set_discovery_hash(hass, discovery_hash):
+def set_discovery_hash(hass: HomeAssistant, discovery_hash: tuple):
     """Clear entry in ALREADY_DISCOVERED list."""
     hass.data[ALREADY_DISCOVERED][discovery_hash] = {}
 
 
-class MQTTConfig(dict):
-    """Dummy class to allow adding attributes."""
-
-
 async def async_start(  # noqa: C901
     hass: HomeAssistant, discovery_topic, config_entry=None
 ) -> None:
@@ -181,6 +185,7 @@ async def async_discovery_message_received(msg):
         await async_process_discovery_payload(component, discovery_id, payload)
 
     async def async_process_discovery_payload(component, discovery_id, payload):
+        """Process the payload of a new discovery."""
 
         _LOGGER.debug("Process discovery payload %s", payload)
         discovery_hash = (component, discovery_id)