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

This commit is contained in:
Erik Montnemery 2020-10-30 09:56:41 +01:00 committed by GitHub
parent b228ffc761
commit 62b5279d86
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
5 changed files with 265 additions and 15 deletions

View file

@ -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"
)

View file

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

View file

@ -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)
)

View file

@ -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:

View file

@ -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