From 62b5279d863fe39d66282b5e6799be11a8e4a314 Mon Sep 17 00:00:00 2001 From: Erik Montnemery Date: Fri, 30 Oct 2020 09:56:41 +0100 Subject: [PATCH] Add WS command to remove a Tasmota device (#42266) --- homeassistant/components/tasmota/__init__.py | 84 ++++++++- homeassistant/components/tasmota/const.py | 1 + .../components/tasmota/device_automation.py | 13 +- .../components/tasmota/device_trigger.py | 6 +- tests/components/tasmota/test_init.py | 176 ++++++++++++++++++ 5 files changed, 265 insertions(+), 15 deletions(-) create mode 100644 tests/components/tasmota/test_init.py diff --git a/homeassistant/components/tasmota/__init__.py b/homeassistant/components/tasmota/__init__.py index a82d95474cc..c0ebae7695e 100644 --- a/homeassistant/components/tasmota/__init__.py +++ b/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) - hass.data[DEVICE_MACS][device.id] = config[CONF_MAC] + device_registry.async_get_or_create(**device_info) 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" + ) diff --git a/homeassistant/components/tasmota/const.py b/homeassistant/components/tasmota/const.py index 7cddeb7a603..0f4dfde1646 100644 --- a/homeassistant/components/tasmota/const.py +++ b/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" diff --git a/homeassistant/components/tasmota/device_automation.py b/homeassistant/components/tasmota/device_automation.py index e921a186fea..aab0064bb96 100644 --- a/homeassistant/components/tasmota/device_automation.py +++ b/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) + ) diff --git a/homeassistant/components/tasmota/device_trigger.py b/homeassistant/components/tasmota/device_trigger.py index 9db7ca492af..e7dad0885a0 100644 --- a/homeassistant/components/tasmota/device_trigger.py +++ b/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): - """Handle the removal of a Tasmota device - cleanup any device triggers.""" +async def async_remove_triggers(hass: HomeAssistant, device_id: str): + """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: diff --git a/tests/components/tasmota/test_init.py b/tests/components/tasmota/test_init.py new file mode 100644 index 00000000000..39d99872741 --- /dev/null +++ b/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