Add WS command to remove a Tasmota device (#42266)

homeassistant/components/tasmota/__init__.py

@@ -11,23 +11,32 @@ from hatasmota.const import (
 )
 from hatasmota.discovery import clear_discovery_topic
 from hatasmota.mqtt import TasmotaMQTTClient
+import voluptuous as vol
 
-from homeassistant.components import mqtt
+from homeassistant.components import mqtt, websocket_api
 from homeassistant.components.mqtt.subscription import (
     async_subscribe_topics,
     async_unsubscribe_topics,
 )
 from homeassistant.core import callback
-from homeassistant.helpers.device_registry import CONNECTION_NETWORK_MAC
+from homeassistant.helpers.device_registry import (
+    CONNECTION_NETWORK_MAC,
+    EVENT_DEVICE_REGISTRY_UPDATED,
+    async_entries_for_config_entry,
+)
 from homeassistant.helpers.typing import HomeAssistantType
 
 from . import device_automation, discovery
-from .const import CONF_DISCOVERY_PREFIX, DATA_REMOVE_DISCOVER_COMPONENT, PLATFORMS
+from .const import (
+    CONF_DISCOVERY_PREFIX,
+    DATA_REMOVE_DISCOVER_COMPONENT,
+    DATA_UNSUB,
+    DOMAIN,
+    PLATFORMS,
+)
 
 _LOGGER = logging.getLogger(__name__)
 
-DEVICE_MACS = "tasmota_devices"
-
 
 async def async_setup(hass: HomeAssistantType, config: dict):
     """Set up the Tasmota component."""
@@ -36,7 +45,8 @@ async def async_setup(hass: HomeAssistantType, config: dict):
 
 async def async_setup_entry(hass, entry):
     """Set up Tasmota from a config entry."""
-    hass.data[DEVICE_MACS] = {}
+    websocket_api.async_register_command(hass, websocket_remove_device)
+    hass.data[DATA_UNSUB] = []
 
     def _publish(*args, **kwds):
         mqtt.async_publish(hass, *args, **kwds)
@@ -59,6 +69,25 @@ async def async_setup_entry(hass, entry):
         """Discover and add a Tasmota device."""
         async_setup_device(hass, mac, config, entry, tasmota_mqtt, device_registry)
 
+    async def async_device_removed(event):
+        """Handle the removal of a device."""
+        device_registry = await hass.helpers.device_registry.async_get_registry()
+        if event.data["action"] != "remove":
+            return
+
+        device = device_registry.deleted_devices[event.data["device_id"]]
+
+        if entry.entry_id not in device.config_entries:
+            return
+
+        macs = [c[1] for c in device.connections if c[0] == CONNECTION_NETWORK_MAC]
+        for mac in macs:
+            clear_discovery_topic(mac, entry.data[CONF_DISCOVERY_PREFIX], tasmota_mqtt)
+
+    hass.data[DATA_UNSUB].append(
+        hass.bus.async_listen(EVENT_DEVICE_REGISTRY_UPDATED, async_device_removed)
+    )
+
     async def start_platforms():
         await device_automation.async_setup_entry(hass, entry)
         await asyncio.gather(
@@ -94,11 +123,20 @@ async def async_unload_entry(hass, entry):
 
     # disable discovery
     await discovery.async_stop(hass)
-    hass.data.pop(DEVICE_MACS)
+
+    # cleanup subscriptions
+    for unsub in hass.data[DATA_UNSUB]:
+        unsub()
     hass.data.pop(DATA_REMOVE_DISCOVER_COMPONENT.format("device_automation"))()
     for component in PLATFORMS:
         hass.data.pop(DATA_REMOVE_DISCOVER_COMPONENT.format(component))()
 
+    # deattach device triggers
+    device_registry = await hass.helpers.device_registry.async_get_registry()
+    devices = async_entries_for_config_entry(device_registry, entry.entry_id)
+    for device in devices:
+        await device_automation.async_remove_automations(hass, device.id)
+
     return True
 
 
@@ -126,8 +164,7 @@ def _update_device(hass, config_entry, config, device_registry):
         "config_entry_id": config_entry_id,
     }
     _LOGGER.debug("Adding or updating tasmota device %s", config[CONF_MAC])
-    device = device_registry.async_get_or_create(**device_info)
+    device_registry.async_get_or_create(**device_info)
-    hass.data[DEVICE_MACS][device.id] = config[CONF_MAC]
 
 
 def async_setup_device(hass, mac, config, config_entry, tasmota_mqtt, device_registry):
@@ -136,3 +173,32 @@ def async_setup_device(hass, mac, config, config_entry, tasmota_mqtt, device_reg
         _remove_device(hass, config_entry, mac, tasmota_mqtt, device_registry)
     else:
         _update_device(hass, config_entry, config, device_registry)
+
+
+@websocket_api.websocket_command(
+    {vol.Required("type"): "tasmota/device/remove", vol.Required("device_id"): str}
+)
+@websocket_api.async_response
+async def websocket_remove_device(hass, connection, msg):
+    """Delete device."""
+    device_id = msg["device_id"]
+    dev_registry = await hass.helpers.device_registry.async_get_registry()
+
+    device = dev_registry.async_get(device_id)
+    if not device:
+        connection.send_error(
+            msg["id"], websocket_api.const.ERR_NOT_FOUND, "Device not found"
+        )
+        return
+
+    for config_entry in device.config_entries:
+        config_entry = hass.config_entries.async_get_entry(config_entry)
+        # Only delete the device if it belongs to a Tasmota device entry
+        if config_entry.domain == DOMAIN:
+            dev_registry.async_remove_device(device_id)
+            connection.send_message(websocket_api.result_message(msg["id"]))
+            return
+
+    connection.send_error(
+        msg["id"], websocket_api.const.ERR_NOT_FOUND, "Non Tasmota device"
+    )

homeassistant/components/tasmota/const.py

@@ -2,6 +2,7 @@
 CONF_DISCOVERY_PREFIX = "discovery_prefix"
 
 DATA_REMOVE_DISCOVER_COMPONENT = "tasmota_discover_{}"
+DATA_UNSUB = "tasmota_subscriptions"
 
 DEFAULT_PREFIX = "tasmota/discovery"
 

homeassistant/components/tasmota/device_automation.py

@@ -6,10 +6,15 @@ from homeassistant.helpers.device_registry import EVENT_DEVICE_REGISTRY_UPDATED
 from homeassistant.helpers.dispatcher import async_dispatcher_connect
 
 from . import device_trigger
-from .const import DATA_REMOVE_DISCOVER_COMPONENT
+from .const import DATA_REMOVE_DISCOVER_COMPONENT, DATA_UNSUB
 from .discovery import TASMOTA_DISCOVERY_ENTITY_NEW
 
 
+async def async_remove_automations(hass, device_id):
+    """Remove automations for a Tasmota device."""
+    await device_trigger.async_remove_triggers(hass, device_id)
+
+
 async def async_setup_entry(hass, config_entry):
     """Set up Tasmota device automation dynamically through discovery."""
 
@@ -17,7 +22,7 @@ async def async_setup_entry(hass, config_entry):
         """Handle the removal of a device."""
         if event.data["action"] != "remove":
             return
-        await device_trigger.async_device_removed(hass, event.data["device_id"])
+        await async_remove_automations(hass, event.data["device_id"])
 
     async def async_discover(tasmota_automation, discovery_hash):
         """Discover and add a Tasmota device automation."""
@@ -33,4 +38,6 @@ async def async_setup_entry(hass, config_entry):
         TASMOTA_DISCOVERY_ENTITY_NEW.format("device_automation", "tasmota"),
         async_discover,
     )
-    hass.bus.async_listen(EVENT_DEVICE_REGISTRY_UPDATED, async_device_removed)
+    hass.data[DATA_UNSUB].append(
+        hass.bus.async_listen(EVENT_DEVICE_REGISTRY_UPDATED, async_device_removed)
+    )

homeassistant/components/tasmota/device_trigger.py

@@ -224,8 +224,8 @@ async def async_setup_trigger(hass, tasmota_trigger, config_entry, discovery_has
     await device_trigger.arm_tasmota_trigger()
 
 
-async def async_device_removed(hass: HomeAssistant, device_id: str):
+async def async_remove_triggers(hass: HomeAssistant, device_id: str):
-    """Handle the removal of a Tasmota device - cleanup any device triggers."""
+    """Cleanup any device triggers for a Tasmota device."""
     triggers = await async_get_triggers(hass, device_id)
     for trig in triggers:
         device_trigger = hass.data[DEVICE_TRIGGERS].pop(trig[CONF_DISCOVERY_ID])
@@ -239,7 +239,7 @@ async def async_device_removed(hass: HomeAssistant, device_id: str):
 
 
 async def async_get_triggers(hass: HomeAssistant, device_id: str) -> List[dict]:
-    """List device triggers for Tasmota devices."""
+    """List device triggers for a Tasmota device."""
     triggers = []
 
     if DEVICE_TRIGGERS not in hass.data:

tests/components/tasmota/test_init.py

@@ -0,0 +1,176 @@
+"""The tests for the Tasmota binary sensor platform."""
+import copy
+import json
+
+from homeassistant.components import websocket_api
+from homeassistant.components.tasmota.const import DEFAULT_PREFIX
+
+from .test_common import DEFAULT_CONFIG
+
+from tests.async_mock import call
+from tests.common import MockConfigEntry, async_fire_mqtt_message
+
+
+async def test_device_remove(
+    hass, mqtt_mock, caplog, device_reg, entity_reg, setup_tasmota
+):
+    """Test removing a discovered device through device registry."""
+    config = copy.deepcopy(DEFAULT_CONFIG)
+    mac = config["mac"]
+
+    async_fire_mqtt_message(hass, f"{DEFAULT_PREFIX}/{mac}/config", json.dumps(config))
+    await hass.async_block_till_done()
+
+    # Verify device entry is created
+    device_entry = device_reg.async_get_device(set(), {("mac", mac)})
+    assert device_entry is not None
+
+    device_reg.async_remove_device(device_entry.id)
+    await hass.async_block_till_done()
+
+    # Verify device entry is removed
+    device_entry = device_reg.async_get_device(set(), {("mac", mac)})
+    assert device_entry is None
+
+    # Verify retained discovery topic has been cleared
+    mqtt_mock.async_publish.assert_has_calls(
+        [
+            call(f"tasmota/discovery/{mac}/config", "", 0, True),
+            call(f"tasmota/discovery/{mac}/sensors", "", 0, True),
+        ],
+        any_order=True,
+    )
+
+
+async def test_device_remove_non_tasmota_device(
+    hass, device_reg, hass_ws_client, mqtt_mock, setup_tasmota
+):
+    """Test removing a non Tasmota device through device registry."""
+    config_entry = MockConfigEntry(domain="test")
+    config_entry.add_to_hass(hass)
+
+    mac = "12:34:56:AB:CD:EF"
+    device_entry = device_reg.async_get_or_create(
+        config_entry_id=config_entry.entry_id,
+        connections={("mac", mac)},
+    )
+    assert device_entry is not None
+
+    device_reg.async_remove_device(device_entry.id)
+    await hass.async_block_till_done()
+
+    # Verify device entry is removed
+    device_entry = device_reg.async_get_device(set(), {("mac", mac)})
+    assert device_entry is None
+
+    # Verify no Tasmota discovery message was sent
+    mqtt_mock.async_publish.assert_not_called()
+
+
+async def test_device_remove_stale_tasmota_device(
+    hass, device_reg, hass_ws_client, mqtt_mock, setup_tasmota
+):
+    """Test removing a stale (undiscovered) Tasmota device through device registry."""
+    config_entry = hass.config_entries.async_entries("tasmota")[0]
+
+    mac = "12:34:56:AB:CD:EF"
+    device_entry = device_reg.async_get_or_create(
+        config_entry_id=config_entry.entry_id,
+        connections={("mac", mac)},
+    )
+    assert device_entry is not None
+
+    device_reg.async_remove_device(device_entry.id)
+    await hass.async_block_till_done()
+
+    # Verify device entry is removed
+    device_entry = device_reg.async_get_device(set(), {("mac", mac)})
+    assert device_entry is None
+
+    # Verify retained discovery topic has been cleared
+    mac = mac.replace(":", "")
+    mqtt_mock.async_publish.assert_has_calls(
+        [
+            call(f"tasmota/discovery/{mac}/config", "", 0, True),
+            call(f"tasmota/discovery/{mac}/sensors", "", 0, True),
+        ],
+        any_order=True,
+    )
+
+
+async def test_tasmota_ws_remove_discovered_device(
+    hass, device_reg, entity_reg, hass_ws_client, mqtt_mock, setup_tasmota
+):
+    """Test Tasmota websocket device removal."""
+    config = copy.deepcopy(DEFAULT_CONFIG)
+    mac = config["mac"]
+
+    async_fire_mqtt_message(hass, f"{DEFAULT_PREFIX}/{mac}/config", json.dumps(config))
+    await hass.async_block_till_done()
+
+    # Verify device entry is created
+    device_entry = device_reg.async_get_device(set(), {("mac", mac)})
+    assert device_entry is not None
+
+    client = await hass_ws_client(hass)
+    await client.send_json(
+        {"id": 5, "type": "tasmota/device/remove", "device_id": device_entry.id}
+    )
+    response = await client.receive_json()
+    assert response["success"]
+
+    # Verify device entry is cleared
+    device_entry = device_reg.async_get_device(set(), {("mac", mac)})
+    assert device_entry is None
+
+
+async def test_tasmota_ws_remove_discovered_device_twice(
+    hass, device_reg, hass_ws_client, mqtt_mock, setup_tasmota
+):
+    """Test Tasmota websocket device removal."""
+    config = copy.deepcopy(DEFAULT_CONFIG)
+    mac = config["mac"]
+
+    async_fire_mqtt_message(hass, f"{DEFAULT_PREFIX}/{mac}/config", json.dumps(config))
+    await hass.async_block_till_done()
+
+    # Verify device entry is created
+    device_entry = device_reg.async_get_device(set(), {("mac", mac)})
+    assert device_entry is not None
+
+    client = await hass_ws_client(hass)
+    await client.send_json(
+        {"id": 5, "type": "tasmota/device/remove", "device_id": device_entry.id}
+    )
+    response = await client.receive_json()
+    assert response["success"]
+
+    await client.send_json(
+        {"id": 6, "type": "tasmota/device/remove", "device_id": device_entry.id}
+    )
+    response = await client.receive_json()
+    assert not response["success"]
+    assert response["error"]["code"] == websocket_api.const.ERR_NOT_FOUND
+    assert response["error"]["message"] == "Device not found"
+
+
+async def test_tasmota_ws_remove_non_tasmota_device(
+    hass, device_reg, hass_ws_client, mqtt_mock, setup_tasmota
+):
+    """Test Tasmota websocket device removal of device belonging to other domain."""
+    config_entry = MockConfigEntry(domain="test")
+    config_entry.add_to_hass(hass)
+
+    device_entry = device_reg.async_get_or_create(
+        config_entry_id=config_entry.entry_id,
+        connections={("mac", "12:34:56:AB:CD:EF")},
+    )
+    assert device_entry is not None
+
+    client = await hass_ws_client(hass)
+    await client.send_json(
+        {"id": 5, "type": "tasmota/device/remove", "device_id": device_entry.id}
+    )
+    response = await client.receive_json()
+    assert not response["success"]
+    assert response["error"]["code"] == websocket_api.const.ERR_NOT_FOUND