Add MQTT WS command to remove device (#31989)

* Add MQTT WS command to remove device

* Review comments, fix test

* Fix tests

tests/components/mqtt/test_discovery.py

@@ -3,6 +3,8 @@ from pathlib import Path
 import re
 from unittest.mock import patch
 
+import pytest
+
 from homeassistant.components import mqtt
 from homeassistant.components.mqtt.abbreviations import (
     ABBREVIATIONS,
@@ -11,7 +13,25 @@ from homeassistant.components.mqtt.abbreviations import (
 from homeassistant.components.mqtt.discovery import ALREADY_DISCOVERED, async_start
 from homeassistant.const import STATE_OFF, STATE_ON
 
-from tests.common import MockConfigEntry, async_fire_mqtt_message, mock_coro
+from tests.common import (
+    MockConfigEntry,
+    async_fire_mqtt_message,
+    mock_coro,
+    mock_device_registry,
+    mock_registry,
+)
+
+
+@pytest.fixture
+def device_reg(hass):
+    """Return an empty, loaded, registry."""
+    return mock_device_registry(hass)
+
+
+@pytest.fixture
+def entity_reg(hass):
+    """Return an empty, loaded, registry."""
+    return mock_registry(hass)
 
 
 async def test_subscribing_config_topic(hass, mqtt_mock):
@@ -213,6 +233,114 @@ async def test_non_duplicate_discovery(hass, mqtt_mock, caplog):
     assert "Component has already been discovered: binary_sensor bla" in caplog.text
 
 
+async def test_removal(hass, mqtt_mock, caplog):
+    """Test removal of component through empty discovery message."""
+    entry = MockConfigEntry(domain=mqtt.DOMAIN)
+
+    await async_start(hass, "homeassistant", {}, entry)
+
+    async_fire_mqtt_message(
+        hass, "homeassistant/binary_sensor/bla/config", '{ "name": "Beer" }'
+    )
+    await hass.async_block_till_done()
+    state = hass.states.get("binary_sensor.beer")
+    assert state is not None
+
+    async_fire_mqtt_message(hass, "homeassistant/binary_sensor/bla/config", "")
+    await hass.async_block_till_done()
+    state = hass.states.get("binary_sensor.beer")
+    assert state is None
+
+
+async def test_rediscover(hass, mqtt_mock, caplog):
+    """Test rediscover of removed component."""
+    entry = MockConfigEntry(domain=mqtt.DOMAIN)
+
+    await async_start(hass, "homeassistant", {}, entry)
+
+    async_fire_mqtt_message(
+        hass, "homeassistant/binary_sensor/bla/config", '{ "name": "Beer" }'
+    )
+    await hass.async_block_till_done()
+    state = hass.states.get("binary_sensor.beer")
+    assert state is not None
+
+    async_fire_mqtt_message(hass, "homeassistant/binary_sensor/bla/config", "")
+    await hass.async_block_till_done()
+    state = hass.states.get("binary_sensor.beer")
+    assert state is None
+
+    async_fire_mqtt_message(
+        hass, "homeassistant/binary_sensor/bla/config", '{ "name": "Beer" }'
+    )
+    await hass.async_block_till_done()
+    state = hass.states.get("binary_sensor.beer")
+    assert state is not None
+
+
+async def test_duplicate_removal(hass, mqtt_mock, caplog):
+    """Test for a non duplicate component."""
+    entry = MockConfigEntry(domain=mqtt.DOMAIN)
+
+    await async_start(hass, "homeassistant", {}, entry)
+
+    async_fire_mqtt_message(
+        hass, "homeassistant/binary_sensor/bla/config", '{ "name": "Beer" }'
+    )
+    await hass.async_block_till_done()
+    async_fire_mqtt_message(hass, "homeassistant/binary_sensor/bla/config", "")
+    await hass.async_block_till_done()
+    assert "Component has already been discovered: binary_sensor bla" in caplog.text
+    caplog.clear()
+    async_fire_mqtt_message(hass, "homeassistant/binary_sensor/bla/config", "")
+    await hass.async_block_till_done()
+
+    assert "Component has already been discovered: binary_sensor bla" not in caplog.text
+
+
+async def test_cleanup_device(hass, device_reg, entity_reg, mqtt_mock):
+    """Test discvered device is cleaned up when removed from registry."""
+    config_entry = MockConfigEntry(domain=mqtt.DOMAIN)
+    config_entry.add_to_hass(hass)
+    await async_start(hass, "homeassistant", {}, config_entry)
+
+    data = (
+        '{ "device":{"identifiers":["0AFFD2"]},'
+        '  "state_topic": "foobar/sensor",'
+        '  "unique_id": "unique" }'
+    )
+
+    async_fire_mqtt_message(hass, "homeassistant/sensor/bla/config", data)
+    await hass.async_block_till_done()
+
+    # Verify device and registry entries are created
+    device_entry = device_reg.async_get_device({("mqtt", "0AFFD2")}, set())
+    assert device_entry is not None
+    entity_entry = entity_reg.async_get("sensor.mqtt_sensor")
+    assert entity_entry is not None
+
+    state = hass.states.get("sensor.mqtt_sensor")
+    assert state is not None
+
+    device_reg.async_remove_device(device_entry.id)
+    await hass.async_block_till_done()
+
+    # Verify device and registry entries are cleared
+    device_entry = device_reg.async_get_device({("mqtt", "0AFFD2")}, set())
+    assert device_entry is None
+    entity_entry = entity_reg.async_get("sensor.mqtt_sensor")
+    assert entity_entry is None
+
+    # Verify state is removed
+    state = hass.states.get("sensor.mqtt_sensor")
+    assert state is None
+
+    # Verify retained discovery topic has been cleared
+    mqtt_mock.async_publish.assert_called_once_with(
+        "homeassistant/sensor/bla/config", "", 0, True
+    )
+
+
 async def test_discovery_expansion(hass, mqtt_mock, caplog):
     """Test expansion of abbreviated discovery payload."""
     entry = MockConfigEntry(domain=mqtt.DOMAIN)