Cleanup entity and device registry on MQTT discovery removal (#32693)

* Cleanup entity and device registry on MQTT discovery removal.

* Review comments

homeassistant/components/mqtt/__init__.py

@@ -37,7 +37,6 @@ from homeassistant.exceptions import (
     Unauthorized,
 )
 from homeassistant.helpers import config_validation as cv, event, template
-from homeassistant.helpers.device_registry import async_get_registry as get_dev_reg
 from homeassistant.helpers.dispatcher import async_dispatcher_connect
 from homeassistant.helpers.entity import Entity
 from homeassistant.helpers.typing import ConfigType, HomeAssistantType, ServiceDataType
@@ -1157,6 +1156,23 @@ class MqttAvailability(Entity):
         return availability_topic is None or self._available
 
 
+async def cleanup_device_registry(hass, device_id):
+    """Remove device registry entry if there are no entities or triggers."""
+    # Local import to avoid circular dependencies
+    from . import device_trigger
+
+    device_registry = await hass.helpers.device_registry.async_get_registry()
+    entity_registry = await hass.helpers.entity_registry.async_get_registry()
+    if (
+        device_id
+        and not hass.helpers.entity_registry.async_entries_for_device(
+            entity_registry, device_id
+        )
+        and not await device_trigger.async_get_triggers(hass, device_id)
+    ):
+        device_registry.async_remove_device(device_id)
+
+
 class MqttDiscoveryUpdate(Entity):
     """Mixin used to handle updated discovery message."""
 
@@ -1173,8 +1189,18 @@ class MqttDiscoveryUpdate(Entity):
             self._discovery_data[ATTR_DISCOVERY_HASH] if self._discovery_data else None
         )
 
+        async def async_remove_from_registry(self) -> None:
+            """Remove entity from entity registry."""
+            entity_registry = (
+                await self.hass.helpers.entity_registry.async_get_registry()
+            )
+            if entity_registry.async_is_registered(self.entity_id):
+                entity_entry = entity_registry.async_get(self.entity_id)
+                entity_registry.async_remove(self.entity_id)
+                await cleanup_device_registry(self.hass, entity_entry.device_id)
+
         @callback
-        def discovery_callback(payload):
+        async def discovery_callback(payload):
             """Handle discovery update."""
             _LOGGER.info(
                 "Got update for entity with hash: %s '%s'", discovery_hash, payload,
@@ -1182,13 +1208,13 @@ class MqttDiscoveryUpdate(Entity):
             if not payload:
                 # Empty payload: Remove component
                 _LOGGER.info("Removing component: %s", self.entity_id)
-                self.hass.async_create_task(self.async_remove())
-                clear_discovery_hash(self.hass, discovery_hash)
-                self._remove_signal()
+                self._cleanup_on_remove()
+                await async_remove_from_registry(self)
+                await self.async_remove()
             elif self._discovery_update:
                 # Non-empty payload: Notify component
                 _LOGGER.info("Updating component: %s", self.entity_id)
-                self.hass.async_create_task(self._discovery_update(payload))
+                await self._discovery_update(payload)
 
         if discovery_hash:
             self._remove_signal = async_dispatcher_connect(
@@ -1199,15 +1225,25 @@ class MqttDiscoveryUpdate(Entity):
 
     async def async_removed_from_registry(self) -> None:
         """Clear retained discovery topic in broker."""
+        if self._discovery_data:
             discovery_topic = self._discovery_data[ATTR_DISCOVERY_TOPIC]
             publish(
                 self.hass, discovery_topic, "", retain=True,
             )
 
     async def async_will_remove_from_hass(self) -> None:
-        """Stop listening to signal."""
+        """Stop listening to signal and cleanup discovery data.."""
+        self._cleanup_on_remove()
+
+    def _cleanup_on_remove(self) -> None:
+        """Stop listening to signal and cleanup discovery data."""
+        if self._discovery_data:
+            clear_discovery_hash(self.hass, self._discovery_data[ATTR_DISCOVERY_HASH])
+            self._discovery_data = None
+
         if self._remove_signal:
             self._remove_signal()
+            self._remove_signal = None
 
 
 def device_info_from_config(config):
@@ -1270,7 +1306,7 @@ class MqttEntityDeviceInfo(Entity):
 async def websocket_remove_device(hass, connection, msg):
     """Delete device."""
     device_id = msg["device_id"]
-    dev_registry = await get_dev_reg(hass)
+    dev_registry = await hass.helpers.device_registry.async_get_registry()
 
     device = dev_registry.async_get(device_id)
     if not device:

homeassistant/components/mqtt/device_trigger.py

@@ -25,6 +25,7 @@ from . import (
     CONF_PAYLOAD,
     CONF_QOS,
     DOMAIN,
+    cleanup_device_registry,
 )
 from .discovery import MQTT_DISCOVERY_UPDATED, clear_discovery_hash
 
@@ -187,6 +188,7 @@ async def async_setup_trigger(hass, config, config_entry, discovery_data):
                 device_trigger.detach_trigger()
                 clear_discovery_hash(hass, discovery_hash)
                 remove_signal()
+                await cleanup_device_registry(hass, device.id)
         else:
             # Non-empty payload: Update trigger
             _LOGGER.info("Updating trigger: %s", discovery_hash)

tests/components/mqtt/common.py

@@ -192,6 +192,30 @@ async def help_test_entity_device_info_with_identifier(hass, mqtt_mock, domain,
     assert device.sw_version == "0.1-beta"
 
 
+async def help_test_entity_device_info_remove(hass, mqtt_mock, domain, config):
+    """Test device registry remove."""
+    entry = MockConfigEntry(domain=mqtt.DOMAIN)
+    entry.add_to_hass(hass)
+    await async_start(hass, "homeassistant", {}, entry)
+    dev_registry = await hass.helpers.device_registry.async_get_registry()
+    ent_registry = await hass.helpers.entity_registry.async_get_registry()
+
+    data = json.dumps(config)
+    async_fire_mqtt_message(hass, f"homeassistant/{domain}/bla/config", data)
+    await hass.async_block_till_done()
+
+    device = dev_registry.async_get_device({("mqtt", "helloworld")}, set())
+    assert device is not None
+    assert ent_registry.async_get_entity_id(domain, mqtt.DOMAIN, "veryunique")
+
+    async_fire_mqtt_message(hass, f"homeassistant/{domain}/bla/config", "")
+    await hass.async_block_till_done()
+
+    device = dev_registry.async_get_device({("mqtt", "helloworld")}, set())
+    assert device is None
+    assert not ent_registry.async_get_entity_id(domain, mqtt.DOMAIN, "veryunique")
+
+
 async def help_test_entity_device_info_update(hass, mqtt_mock, domain, config):
     """Test device registry update.
 

tests/components/mqtt/test_alarm_control_panel.py

@@ -19,6 +19,7 @@ from .common import (
     help_test_discovery_removal,
     help_test_discovery_update,
     help_test_discovery_update_attr,
+    help_test_entity_device_info_remove,
     help_test_entity_device_info_update,
     help_test_entity_device_info_with_identifier,
     help_test_entity_id_update,
@@ -509,6 +510,21 @@ async def test_entity_device_info_update(hass, mqtt_mock):
     )
 
 
+async def test_entity_device_info_remove(hass, mqtt_mock):
+    """Test device registry remove."""
+    config = {
+        "platform": "mqtt",
+        "name": "Test 1",
+        "state_topic": "test-topic",
+        "command_topic": "test-command-topic",
+        "device": {"identifiers": ["helloworld"]},
+        "unique_id": "veryunique",
+    }
+    await help_test_entity_device_info_remove(
+        hass, mqtt_mock, alarm_control_panel.DOMAIN, config
+    )
+
+
 async def test_entity_id_update(hass, mqtt_mock):
     """Test MQTT subscriptions are managed when entity_id is updated."""
     config = {

tests/components/mqtt/test_binary_sensor.py

@@ -20,6 +20,7 @@ from .common import (
     help_test_discovery_removal,
     help_test_discovery_update,
     help_test_discovery_update_attr,
+    help_test_entity_device_info_remove,
     help_test_entity_device_info_update,
     help_test_entity_device_info_with_identifier,
     help_test_entity_id_update,
@@ -556,6 +557,20 @@ async def test_entity_device_info_update(hass, mqtt_mock):
     )
 
 
+async def test_entity_device_info_remove(hass, mqtt_mock):
+    """Test device registry remove."""
+    config = {
+        "platform": "mqtt",
+        "name": "Test 1",
+        "state_topic": "test-topic",
+        "device": {"identifiers": ["helloworld"]},
+        "unique_id": "veryunique",
+    }
+    await help_test_entity_device_info_remove(
+        hass, mqtt_mock, binary_sensor.DOMAIN, config
+    )
+
+
 async def test_entity_id_update(hass, mqtt_mock):
     """Test MQTT subscriptions are managed when entity_id is updated."""
     config = {

tests/components/mqtt/test_climate.py

@@ -30,6 +30,7 @@ from .common import (
     help_test_discovery_removal,
     help_test_discovery_update,
     help_test_discovery_update_attr,
+    help_test_entity_device_info_remove,
     help_test_entity_device_info_update,
     help_test_entity_device_info_with_identifier,
     help_test_entity_id_update,
@@ -893,6 +894,19 @@ async def test_entity_device_info_update(hass, mqtt_mock):
     )
 
 
+async def test_entity_device_info_remove(hass, mqtt_mock):
+    """Test device registry remove."""
+    config = {
+        "platform": "mqtt",
+        "name": "Test 1",
+        "power_state_topic": "test-topic",
+        "power_command_topic": "test-command-topic",
+        "device": {"identifiers": ["helloworld"]},
+        "unique_id": "veryunique",
+    }
+    await help_test_entity_device_info_remove(hass, mqtt_mock, CLIMATE_DOMAIN, config)
+
+
 async def test_entity_id_update(hass, mqtt_mock):
     """Test MQTT subscriptions are managed when entity_id is updated."""
     config = {

tests/components/mqtt/test_cover.py

@@ -28,6 +28,7 @@ from .common import (
     help_test_discovery_removal,
     help_test_discovery_update,
     help_test_discovery_update_attr,
+    help_test_entity_device_info_remove,
     help_test_entity_device_info_update,
     help_test_entity_device_info_with_identifier,
     help_test_entity_id_update,
@@ -1819,6 +1820,19 @@ async def test_entity_device_info_update(hass, mqtt_mock):
     )
 
 
+async def test_entity_device_info_remove(hass, mqtt_mock):
+    """Test device registry remove."""
+    config = {
+        "platform": "mqtt",
+        "name": "Test 1",
+        "state_topic": "test-topic",
+        "command_topic": "test-command-topic",
+        "device": {"identifiers": ["helloworld"]},
+        "unique_id": "veryunique",
+    }
+    await help_test_entity_device_info_remove(hass, mqtt_mock, cover.DOMAIN, config)
+
+
 async def test_entity_id_update(hass, mqtt_mock):
     """Test MQTT subscriptions are managed when entity_id is updated."""
     config = {

tests/components/mqtt/test_device_trigger.py

@@ -233,8 +233,8 @@ async def test_update_remove_triggers(hass, device_reg, entity_reg, mqtt_mock):
     async_fire_mqtt_message(hass, "homeassistant/device_automation/bla/config", "")
     await hass.async_block_till_done()
 
-    triggers = await async_get_device_automations(hass, "trigger", device_entry.id)
-    assert_lists_same(triggers, [])
+    device_entry = device_reg.async_get_device({("mqtt", "0AFFD2")}, set())
+    assert device_entry is None
 
 
 async def test_if_fires_on_mqtt_message(hass, device_reg, calls, mqtt_mock):
@@ -833,8 +833,8 @@ async def test_entity_device_info_update(hass, mqtt_mock):
     assert device.name == "Milk"
 
 
-async def test_cleanup_device(hass, device_reg, entity_reg, mqtt_mock):
-    """Test discovered device is cleaned up when removed from registry."""
+async def test_cleanup_trigger(hass, device_reg, entity_reg, mqtt_mock):
+    """Test trigger discovery topic is cleaned when device is removed from registry."""
     config_entry = MockConfigEntry(domain=DOMAIN)
     config_entry.add_to_hass(hass)
     await async_start(hass, "homeassistant", {}, config_entry)
@@ -863,10 +863,212 @@ async def test_cleanup_device(hass, device_reg, entity_reg, mqtt_mock):
     await hass.async_block_till_done()
 
     # Verify device registry entry is cleared
-    device_entry = device_reg.async_get_device({("mqtt", "0AFFD2")}, set())
+    device_entry = device_reg.async_get_device({("mqtt", "helloworld")}, set())
     assert device_entry is None
 
     # Verify retained discovery topic has been cleared
     mqtt_mock.async_publish.assert_called_once_with(
         "homeassistant/device_automation/bla/config", "", 0, True
     )
+
+
+async def test_cleanup_device(hass, device_reg, entity_reg, mqtt_mock):
+    """Test removal from device registry when trigger is removed."""
+    config_entry = MockConfigEntry(domain=DOMAIN)
+    config_entry.add_to_hass(hass)
+    await async_start(hass, "homeassistant", {}, config_entry)
+
+    config = {
+        "automation_type": "trigger",
+        "topic": "test-topic",
+        "type": "foo",
+        "subtype": "bar",
+        "device": {"identifiers": ["helloworld"]},
+    }
+
+    data = json.dumps(config)
+    async_fire_mqtt_message(hass, "homeassistant/device_automation/bla/config", data)
+    await hass.async_block_till_done()
+
+    # Verify device registry entry is created
+    device_entry = device_reg.async_get_device({("mqtt", "helloworld")}, set())
+    assert device_entry is not None
+
+    triggers = await async_get_device_automations(hass, "trigger", device_entry.id)
+    assert triggers[0]["type"] == "foo"
+
+    async_fire_mqtt_message(hass, "homeassistant/device_automation/bla/config", "")
+    await hass.async_block_till_done()
+
+    # Verify device registry entry is cleared
+    device_entry = device_reg.async_get_device({("mqtt", "helloworld")}, set())
+    assert device_entry is None
+
+
+async def test_cleanup_device_several_triggers(hass, device_reg, entity_reg, mqtt_mock):
+    """Test removal from device registry when the last trigger is removed."""
+    config_entry = MockConfigEntry(domain=DOMAIN)
+    config_entry.add_to_hass(hass)
+    await async_start(hass, "homeassistant", {}, config_entry)
+
+    config1 = {
+        "automation_type": "trigger",
+        "topic": "test-topic",
+        "type": "foo",
+        "subtype": "bar",
+        "device": {"identifiers": ["helloworld"]},
+    }
+
+    config2 = {
+        "automation_type": "trigger",
+        "topic": "test-topic",
+        "type": "foo2",
+        "subtype": "bar",
+        "device": {"identifiers": ["helloworld"]},
+    }
+
+    data1 = json.dumps(config1)
+    data2 = json.dumps(config2)
+    async_fire_mqtt_message(hass, "homeassistant/device_automation/bla1/config", data1)
+    await hass.async_block_till_done()
+    async_fire_mqtt_message(hass, "homeassistant/device_automation/bla2/config", data2)
+    await hass.async_block_till_done()
+
+    # Verify device registry entry is created
+    device_entry = device_reg.async_get_device({("mqtt", "helloworld")}, set())
+    assert device_entry is not None
+
+    triggers = await async_get_device_automations(hass, "trigger", device_entry.id)
+    assert len(triggers) == 2
+    assert triggers[0]["type"] == "foo"
+    assert triggers[1]["type"] == "foo2"
+
+    async_fire_mqtt_message(hass, "homeassistant/device_automation/bla1/config", "")
+    await hass.async_block_till_done()
+
+    # Verify device registry entry is not cleared
+    device_entry = device_reg.async_get_device({("mqtt", "helloworld")}, set())
+    assert device_entry is not None
+
+    triggers = await async_get_device_automations(hass, "trigger", device_entry.id)
+    assert len(triggers) == 1
+    assert triggers[0]["type"] == "foo2"
+
+    async_fire_mqtt_message(hass, "homeassistant/device_automation/bla2/config", "")
+    await hass.async_block_till_done()
+
+    # Verify device registry entry is cleared
+    device_entry = device_reg.async_get_device({("mqtt", "helloworld")}, set())
+    assert device_entry is None
+
+
+async def test_cleanup_device_with_entity1(hass, device_reg, entity_reg, mqtt_mock):
+    """Test removal from device registry for device with entity.
+
+    Trigger removed first, then entity.
+    """
+    config_entry = MockConfigEntry(domain=DOMAIN)
+    config_entry.add_to_hass(hass)
+    await async_start(hass, "homeassistant", {}, config_entry)
+
+    config1 = {
+        "automation_type": "trigger",
+        "topic": "test-topic",
+        "type": "foo",
+        "subtype": "bar",
+        "device": {"identifiers": ["helloworld"]},
+    }
+
+    config2 = {
+        "name": "test_binary_sensor",
+        "state_topic": "test-topic",
+        "device": {"identifiers": ["helloworld"]},
+        "unique_id": "veryunique",
+    }
+
+    data1 = json.dumps(config1)
+    data2 = json.dumps(config2)
+    async_fire_mqtt_message(hass, "homeassistant/device_automation/bla1/config", data1)
+    await hass.async_block_till_done()
+    async_fire_mqtt_message(hass, "homeassistant/binary_sensor/bla2/config", data2)
+    await hass.async_block_till_done()
+
+    # Verify device registry entry is created
+    device_entry = device_reg.async_get_device({("mqtt", "helloworld")}, set())
+    assert device_entry is not None
+
+    triggers = await async_get_device_automations(hass, "trigger", device_entry.id)
+    assert len(triggers) == 3  # 2 binary_sensor triggers + device trigger
+
+    async_fire_mqtt_message(hass, "homeassistant/device_automation/bla1/config", "")
+    await hass.async_block_till_done()
+
+    # Verify device registry entry is not cleared
+    device_entry = device_reg.async_get_device({("mqtt", "helloworld")}, set())
+    assert device_entry is not None
+
+    triggers = await async_get_device_automations(hass, "trigger", device_entry.id)
+    assert len(triggers) == 2  # 2 binary_sensor triggers
+
+    async_fire_mqtt_message(hass, "homeassistant/binary_sensor/bla2/config", "")
+    await hass.async_block_till_done()
+
+    # Verify device registry entry is cleared
+    device_entry = device_reg.async_get_device({("mqtt", "helloworld")}, set())
+    assert device_entry is None
+
+
+async def test_cleanup_device_with_entity2(hass, device_reg, entity_reg, mqtt_mock):
+    """Test removal from device registry for device with entity.
+
+    Entity removed first, then trigger.
+    """
+    config_entry = MockConfigEntry(domain=DOMAIN)
+    config_entry.add_to_hass(hass)
+    await async_start(hass, "homeassistant", {}, config_entry)
+
+    config1 = {
+        "automation_type": "trigger",
+        "topic": "test-topic",
+        "type": "foo",
+        "subtype": "bar",
+        "device": {"identifiers": ["helloworld"]},
+    }
+
+    config2 = {
+        "name": "test_binary_sensor",
+        "state_topic": "test-topic",
+        "device": {"identifiers": ["helloworld"]},
+        "unique_id": "veryunique",
+    }
+
+    data1 = json.dumps(config1)
+    data2 = json.dumps(config2)
+    async_fire_mqtt_message(hass, "homeassistant/device_automation/bla1/config", data1)
+    await hass.async_block_till_done()
+    async_fire_mqtt_message(hass, "homeassistant/binary_sensor/bla2/config", data2)
+    await hass.async_block_till_done()
+
+    # Verify device registry entry is created
+    device_entry = device_reg.async_get_device({("mqtt", "helloworld")}, set())
+    assert device_entry is not None
+
+    triggers = await async_get_device_automations(hass, "trigger", device_entry.id)
+    assert len(triggers) == 3  # 2 binary_sensor triggers + device trigger
+
+    async_fire_mqtt_message(hass, "homeassistant/binary_sensor/bla2/config", "")
+    await hass.async_block_till_done()
+
+    # Verify device registry entry is not cleared
+    device_entry = device_reg.async_get_device({("mqtt", "helloworld")}, set())
+    assert device_entry is not None
+
+    triggers = await async_get_device_automations(hass, "trigger", device_entry.id)
+    assert len(triggers) == 1  # device trigger
+
+    async_fire_mqtt_message(hass, "homeassistant/device_automation/bla1/config", "")
+    await hass.async_block_till_done()
+
+    # Verify device registry entry is cleared
+    device_entry = device_reg.async_get_device({("mqtt", "helloworld")}, set())
+    assert device_entry is None

tests/components/mqtt/test_fan.py

@@ -13,6 +13,7 @@ from .common import (
     help_test_discovery_removal,
     help_test_discovery_update,
     help_test_discovery_update_attr,
+    help_test_entity_device_info_remove,
     help_test_entity_device_info_update,
     help_test_entity_device_info_with_identifier,
     help_test_entity_id_update,
@@ -563,6 +564,19 @@ async def test_entity_device_info_update(hass, mqtt_mock):
     )
 
 
+async def test_entity_device_info_remove(hass, mqtt_mock):
+    """Test device registry remove."""
+    config = {
+        "platform": "mqtt",
+        "name": "Test 1",
+        "state_topic": "test-topic",
+        "command_topic": "test-command-topic",
+        "device": {"identifiers": ["helloworld"]},
+        "unique_id": "veryunique",
+    }
+    await help_test_entity_device_info_remove(hass, mqtt_mock, fan.DOMAIN, config)
+
+
 async def test_entity_id_update(hass, mqtt_mock):
     """Test MQTT subscriptions are managed when entity_id is updated."""
     config = {

tests/components/mqtt/test_legacy_vacuum.py

@@ -30,6 +30,7 @@ from .common import (
     help_test_discovery_removal,
     help_test_discovery_update,
     help_test_discovery_update_attr,
+    help_test_entity_device_info_remove,
     help_test_entity_device_info_update,
     help_test_entity_device_info_with_identifier,
     help_test_entity_id_update,
@@ -643,6 +644,18 @@ async def test_entity_device_info_update(hass, mqtt_mock):
     )
 
 
+async def test_entity_device_info_remove(hass, mqtt_mock):
+    """Test device registry remove."""
+    config = {
+        "platform": "mqtt",
+        "name": "Test 1",
+        "command_topic": "test-command-topic",
+        "device": {"identifiers": ["helloworld"]},
+        "unique_id": "veryunique",
+    }
+    await help_test_entity_device_info_remove(hass, mqtt_mock, vacuum.DOMAIN, config)
+
+
 async def test_entity_id_update(hass, mqtt_mock):
     """Test MQTT subscriptions are managed when entity_id is updated."""
     config = {

tests/components/mqtt/test_light.py

@@ -172,6 +172,7 @@ from .common import (
     help_test_discovery_removal,
     help_test_discovery_update,
     help_test_discovery_update_attr,
+    help_test_entity_device_info_remove,
     help_test_entity_device_info_update,
     help_test_entity_device_info_with_identifier,
     help_test_entity_id_update,
@@ -1226,6 +1227,19 @@ async def test_entity_device_info_update(hass, mqtt_mock):
     )
 
 
+async def test_entity_device_info_remove(hass, mqtt_mock):
+    """Test device registry remove."""
+    config = {
+        "platform": "mqtt",
+        "name": "Test 1",
+        "state_topic": "test-topic",
+        "command_topic": "test-command-topic",
+        "device": {"identifiers": ["helloworld"]},
+        "unique_id": "veryunique",
+    }
+    await help_test_entity_device_info_remove(hass, mqtt_mock, light.DOMAIN, config)
+
+
 async def test_entity_id_update(hass, mqtt_mock):
     """Test MQTT subscriptions are managed when entity_id is updated."""
     config = {

tests/components/mqtt/test_light_json.py

@@ -108,6 +108,7 @@ from .common import (
     help_test_discovery_removal,
     help_test_discovery_update,
     help_test_discovery_update_attr,
+    help_test_entity_device_info_remove,
     help_test_entity_device_info_update,
     help_test_entity_device_info_with_identifier,
     help_test_entity_id_update,
@@ -1080,6 +1081,20 @@ async def test_entity_device_info_update(hass, mqtt_mock):
     )
 
 
+async def test_entity_device_info_remove(hass, mqtt_mock):
+    """Test device registry remove."""
+    config = {
+        "platform": "mqtt",
+        "name": "Test 1",
+        "schema": "json",
+        "state_topic": "test-topic",
+        "command_topic": "test-command-topic",
+        "device": {"identifiers": ["helloworld"]},
+        "unique_id": "veryunique",
+    }
+    await help_test_entity_device_info_remove(hass, mqtt_mock, light.DOMAIN, config)
+
+
 async def test_entity_id_update(hass, mqtt_mock):
     """Test MQTT subscriptions are managed when entity_id is updated."""
     config = {

tests/components/mqtt/test_light_template.py

@@ -44,6 +44,7 @@ from .common import (
     help_test_discovery_removal,
     help_test_discovery_update,
     help_test_discovery_update_attr,
+    help_test_entity_device_info_remove,
     help_test_entity_device_info_update,
     help_test_entity_device_info_with_identifier,
     help_test_entity_id_update,
@@ -704,6 +705,19 @@ async def test_entity_device_info_update(hass, mqtt_mock):
     )
 
 
+async def test_entity_device_info_remove(hass, mqtt_mock):
+    """Test device registry remove."""
+    config = {
+        "platform": "mqtt",
+        "name": "Test 1",
+        "state_topic": "test-topic",
+        "command_topic": "test-command-topic",
+        "device": {"identifiers": ["helloworld"]},
+        "unique_id": "veryunique",
+    }
+    await help_test_entity_device_info_remove(hass, mqtt_mock, light.DOMAIN, config)
+
+
 async def test_entity_id_update(hass, mqtt_mock):
     """Test MQTT subscriptions are managed when entity_id is updated."""
     config = {

tests/components/mqtt/test_lock.py

@@ -13,6 +13,7 @@ from .common import (
     help_test_discovery_removal,
     help_test_discovery_update,
     help_test_discovery_update_attr,
+    help_test_entity_device_info_remove,
     help_test_entity_device_info_update,
     help_test_entity_device_info_with_identifier,
     help_test_entity_id_update,
@@ -442,6 +443,19 @@ async def test_entity_device_info_update(hass, mqtt_mock):
     )
 
 
+async def test_entity_device_info_remove(hass, mqtt_mock):
+    """Test device registry remove."""
+    config = {
+        "platform": "mqtt",
+        "name": "Test 1",
+        "state_topic": "test-topic",
+        "command_topic": "test-command-topic",
+        "device": {"identifiers": ["helloworld"]},
+        "unique_id": "veryunique",
+    }
+    await help_test_entity_device_info_remove(hass, mqtt_mock, lock.DOMAIN, config)
+
+
 async def test_entity_id_update(hass, mqtt_mock):
     """Test MQTT subscriptions are managed when entity_id is updated."""
     config = {

tests/components/mqtt/test_sensor.py

@@ -16,6 +16,7 @@ from .common import (
     help_test_discovery_removal,
     help_test_discovery_update,
     help_test_discovery_update_attr,
+    help_test_entity_device_info_remove,
     help_test_entity_device_info_update,
     help_test_entity_device_info_with_identifier,
     help_test_entity_id_update,
@@ -559,6 +560,18 @@ async def test_entity_device_info_update(hass, mqtt_mock):
     )
 
 
+async def test_entity_device_info_remove(hass, mqtt_mock):
+    """Test device registry remove."""
+    config = {
+        "platform": "mqtt",
+        "name": "Test 1",
+        "state_topic": "test-topic",
+        "device": {"identifiers": ["helloworld"]},
+        "unique_id": "veryunique",
+    }
+    await help_test_entity_device_info_remove(hass, mqtt_mock, sensor.DOMAIN, config)
+
+
 async def test_entity_id_update(hass, mqtt_mock):
     """Test MQTT subscriptions are managed when entity_id is updated."""
     config = {

tests/components/mqtt/test_state_vacuum.py

@@ -36,6 +36,7 @@ from .common import (
     help_test_discovery_removal,
     help_test_discovery_update,
     help_test_discovery_update_attr,
+    help_test_entity_device_info_remove,
     help_test_entity_device_info_update,
     help_test_entity_device_info_with_identifier,
     help_test_entity_id_update,
@@ -475,6 +476,20 @@ async def test_entity_device_info_update(hass, mqtt_mock):
     )
 
 
+async def test_entity_device_info_remove(hass, mqtt_mock):
+    """Test device registry remove."""
+    config = {
+        "platform": "mqtt",
+        "schema": "state",
+        "name": "Test 1",
+        "state_topic": "test-topic",
+        "command_topic": "test-command-topic",
+        "device": {"identifiers": ["helloworld"]},
+        "unique_id": "veryunique",
+    }
+    await help_test_entity_device_info_remove(hass, mqtt_mock, vacuum.DOMAIN, config)
+
+
 async def test_entity_id_update(hass, mqtt_mock):
     """Test MQTT subscriptions are managed when entity_id is updated."""
     config = {

tests/components/mqtt/test_switch.py

@@ -17,6 +17,7 @@ from .common import (
     help_test_discovery_removal,
     help_test_discovery_update,
     help_test_discovery_update_attr,
+    help_test_entity_device_info_remove,
     help_test_entity_device_info_update,
     help_test_entity_device_info_with_identifier,
     help_test_entity_id_update,
@@ -407,6 +408,19 @@ async def test_entity_device_info_update(hass, mqtt_mock):
     )
 
 
+async def test_entity_device_info_remove(hass, mqtt_mock):
+    """Test device registry remove."""
+    config = {
+        "platform": "mqtt",
+        "name": "Test 1",
+        "state_topic": "test-topic",
+        "command_topic": "test-command-topic",
+        "device": {"identifiers": ["helloworld"]},
+        "unique_id": "veryunique",
+    }
+    await help_test_entity_device_info_remove(hass, mqtt_mock, switch.DOMAIN, config)
+
+
 async def test_entity_id_update(hass, mqtt_mock):
     """Test MQTT subscriptions are managed when entity_id is updated."""
     config = {