Add device action support for ZHA (#26903)

* start implementing device actions

* rename file

* cleanup and add tests

* fix docstrings

* sort imports

homeassistant/components/zha/.translations/en.json

@@ -18,6 +18,10 @@
     "title": "ZHA"
   },
   "device_automation": {
+    "action_type": {
+      "squawk": "Squawk",
+      "warn": "Warn"
+    },
     "trigger_subtype": {
       "both_buttons": "Both buttons",
       "button_1": "First button",

homeassistant/components/zha/core/helpers.py

@@ -10,7 +10,14 @@ import logging
 
 from homeassistant.core import callback
 
-from .const import CLUSTER_TYPE_IN, CLUSTER_TYPE_OUT, DEFAULT_BAUDRATE, RadioType
+from .const import (
+    CLUSTER_TYPE_IN,
+    CLUSTER_TYPE_OUT,
+    DATA_ZHA,
+    DATA_ZHA_GATEWAY,
+    DEFAULT_BAUDRATE,
+    RadioType,
+)
 from .registries import BINDABLE_CLUSTERS
 
 _LOGGER = logging.getLogger(__name__)
@@ -132,6 +139,16 @@ def async_is_bindable_target(source_zha_device, target_zha_device):
     return False
 
 
+async def async_get_zha_device(hass, device_id):
+    """Get a ZHA device for the given device registry id."""
+    device_registry = await hass.helpers.device_registry.async_get_registry()
+    registry_device = device_registry.async_get(device_id)
+    zha_gateway = hass.data[DATA_ZHA][DATA_ZHA_GATEWAY]
+    ieee_address = list(list(registry_device.identifiers)[0])[1]
+    ieee = convert_ieee(ieee_address)
+    return zha_gateway.devices[ieee]
+
+
 class LogMixin:
     """Log helper."""
 

homeassistant/components/zha/device_action.py

@@ -0,0 +1,92 @@
+"""Provides device actions for ZHA devices."""
+from typing import List
+
+import voluptuous as vol
+
+from homeassistant.const import CONF_DEVICE_ID, CONF_DOMAIN, CONF_TYPE
+from homeassistant.core import Context, HomeAssistant
+from homeassistant.helpers import config_validation as cv, service
+from homeassistant.helpers.typing import ConfigType, TemplateVarsType
+
+from . import DOMAIN
+from .api import SERVICE_WARNING_DEVICE_SQUAWK, SERVICE_WARNING_DEVICE_WARN
+from .core.const import CHANNEL_IAS_WD
+from .core.helpers import async_get_zha_device
+
+ACTION_SQUAWK = "squawk"
+ACTION_WARN = "warn"
+ATTR_DATA = "data"
+ATTR_IEEE = "ieee"
+CONF_ZHA_ACTION_TYPE = "zha_action_type"
+ZHA_ACTION_TYPE_SERVICE_CALL = "service_call"
+
+ACTION_SCHEMA = cv.DEVICE_ACTION_BASE_SCHEMA.extend(
+    {vol.Required(CONF_DOMAIN): DOMAIN, vol.Required(CONF_TYPE): str}
+)
+
+DEVICE_ACTIONS = {
+    CHANNEL_IAS_WD: [
+        {CONF_TYPE: ACTION_SQUAWK, CONF_DOMAIN: DOMAIN},
+        {CONF_TYPE: ACTION_WARN, CONF_DOMAIN: DOMAIN},
+    ]
+}
+
+DEVICE_ACTION_TYPES = {
+    ACTION_SQUAWK: ZHA_ACTION_TYPE_SERVICE_CALL,
+    ACTION_WARN: ZHA_ACTION_TYPE_SERVICE_CALL,
+}
+
+SERVICE_NAMES = {
+    ACTION_SQUAWK: SERVICE_WARNING_DEVICE_SQUAWK,
+    ACTION_WARN: SERVICE_WARNING_DEVICE_WARN,
+}
+
+
+async def async_call_action_from_config(
+    hass: HomeAssistant,
+    config: ConfigType,
+    variables: TemplateVarsType,
+    context: Context,
+) -> None:
+    """Perform an action based on configuration."""
+    config = ACTION_SCHEMA(config)
+    await ZHA_ACTION_TYPES[DEVICE_ACTION_TYPES[config[CONF_TYPE]]](
+        hass, config, variables, context
+    )
+
+
+async def async_get_actions(hass: HomeAssistant, device_id: str) -> List[dict]:
+    """List device actions."""
+    zha_device = await async_get_zha_device(hass, device_id)
+    actions = [
+        action
+        for channel in DEVICE_ACTIONS
+        for action in DEVICE_ACTIONS[channel]
+        if channel in zha_device.cluster_channels
+    ]
+    for action in actions:
+        action[CONF_DEVICE_ID] = device_id
+    return actions
+
+
+async def _execute_service_based_action(
+    hass: HomeAssistant,
+    config: ACTION_SCHEMA,
+    variables: TemplateVarsType,
+    context: Context,
+) -> None:
+    action_type = config[CONF_TYPE]
+    service_name = SERVICE_NAMES[action_type]
+    zha_device = await async_get_zha_device(hass, config[CONF_DEVICE_ID])
+
+    service_action = {
+        service.CONF_SERVICE: "{}.{}".format(DOMAIN, service_name),
+        ATTR_DATA: {ATTR_IEEE: str(zha_device.ieee)},
+    }
+
+    await service.async_call_from_config(
+        hass, service_action, blocking=True, variables=variables, context=context
+    )
+
+
+ZHA_ACTION_TYPES = {ZHA_ACTION_TYPE_SERVICE_CALL: _execute_service_based_action}

homeassistant/components/zha/device_trigger.py

@@ -9,8 +9,7 @@ from homeassistant.const import CONF_DEVICE_ID, CONF_DOMAIN, CONF_PLATFORM, CONF
 from homeassistant.components.device_automation import TRIGGER_BASE_SCHEMA
 
 from . import DOMAIN
-from .core.const import DATA_ZHA, DATA_ZHA_GATEWAY
+from .core.helpers import async_get_zha_device
-from .core.helpers import convert_ieee
 
 CONF_SUBTYPE = "subtype"
 DEVICE = "device"
@@ -26,7 +25,7 @@ async def async_attach_trigger(hass, config, action, automation_info):
     """Listen for state changes based on configuration."""
     config = TRIGGER_SCHEMA(config)
     trigger = (config[CONF_TYPE], config[CONF_SUBTYPE])
-    zha_device = await _async_get_zha_device(hass, config[CONF_DEVICE_ID])
+    zha_device = await async_get_zha_device(hass, config[CONF_DEVICE_ID])
 
     if (
         zha_device.device_automation_triggers is None
@@ -52,7 +51,7 @@ async def async_get_triggers(hass, device_id):
     Make sure the device supports device automations and
     if it does return the trigger list.
     """
-    zha_device = await _async_get_zha_device(hass, device_id)
+    zha_device = await async_get_zha_device(hass, device_id)
 
     if not zha_device.device_automation_triggers:
         return
@@ -70,15 +69,3 @@ async def async_get_triggers(hass, device_id):
         )
 
     return triggers
-
-
-async def _async_get_zha_device(hass, device_id):
-    device_registry = await hass.helpers.device_registry.async_get_registry()
-    registry_device = device_registry.async_get(device_id)
-    zha_gateway = hass.data[DATA_ZHA][DATA_ZHA_GATEWAY]
-    ieee_address = list(list(registry_device.identifiers)[0])[1]
-    ieee = convert_ieee(ieee_address)
-    zha_device = zha_gateway.devices[ieee]
-    if not zha_device:
-        raise InvalidDeviceAutomationConfig
-    return zha_device

homeassistant/components/zha/strings.json

@@ -18,6 +18,10 @@
     }
   },
   "device_automation": {
+    "action_type": {
+      "squawk": "Squawk",
+      "warn": "Warn"
+    },
     "trigger_type": {
       "remote_button_short_press": "\"{subtype}\" button pressed",
       "remote_button_short_release": "\"{subtype}\" button released",

tests/components/zha/test_device_action.py

@@ -0,0 +1,133 @@
+"""The test for zha device automation actions."""
+from unittest.mock import patch
+
+import pytest
+
+import homeassistant.components.automation as automation
+from homeassistant.components.device_automation import (
+    _async_get_device_automations as async_get_device_automations,
+)
+from homeassistant.components.zha import DOMAIN
+from homeassistant.components.zha.core.const import CHANNEL_ON_OFF
+from homeassistant.helpers.device_registry import async_get_registry
+from homeassistant.setup import async_setup_component
+
+from .common import async_enable_traffic, async_init_zigpy_device
+
+from tests.common import async_mock_service, mock_coro
+
+SHORT_PRESS = "remote_button_short_press"
+COMMAND = "command"
+COMMAND_SINGLE = "single"
+
+
+@pytest.fixture
+def calls(hass):
+    """Track calls to a mock serivce."""
+    return async_mock_service(hass, "zha", "warning_device_warn")
+
+
+async def test_get_actions(hass, config_entry, zha_gateway):
+    """Test we get the expected actions from a zha device."""
+    from zigpy.zcl.clusters.general import Basic
+    from zigpy.zcl.clusters.security import IasZone, IasWd
+
+    # create zigpy device
+    zigpy_device = await async_init_zigpy_device(
+        hass,
+        [Basic.cluster_id, IasZone.cluster_id, IasWd.cluster_id],
+        [],
+        None,
+        zha_gateway,
+    )
+
+    await hass.config_entries.async_forward_entry_setup(config_entry, DOMAIN)
+    await hass.async_block_till_done()
+    hass.config_entries._entries.append(config_entry)
+
+    zha_device = zha_gateway.get_device(zigpy_device.ieee)
+    ieee_address = str(zha_device.ieee)
+
+    ha_device_registry = await async_get_registry(hass)
+    reg_device = ha_device_registry.async_get_device({(DOMAIN, ieee_address)}, set())
+
+    actions = await async_get_device_automations(hass, "action", reg_device.id)
+
+    expected_actions = [
+        {"domain": DOMAIN, "type": "squawk", "device_id": reg_device.id},
+        {"domain": DOMAIN, "type": "warn", "device_id": reg_device.id},
+    ]
+
+    assert actions == expected_actions
+
+
+async def test_action(hass, config_entry, zha_gateway, calls):
+    """Test for executing a zha device action."""
+
+    from zigpy.zcl.clusters.general import Basic, OnOff
+    from zigpy.zcl.clusters.security import IasZone, IasWd
+    from zigpy.zcl.foundation import Status
+
+    # create zigpy device
+    zigpy_device = await async_init_zigpy_device(
+        hass,
+        [Basic.cluster_id, IasZone.cluster_id, IasWd.cluster_id],
+        [OnOff.cluster_id],
+        None,
+        zha_gateway,
+    )
+
+    zigpy_device.device_automation_triggers = {
+        (SHORT_PRESS, SHORT_PRESS): {COMMAND: COMMAND_SINGLE}
+    }
+
+    await hass.config_entries.async_forward_entry_setup(config_entry, "switch")
+    await hass.async_block_till_done()
+
+    hass.config_entries._entries.append(config_entry)
+
+    zha_device = zha_gateway.get_device(zigpy_device.ieee)
+    ieee_address = str(zha_device.ieee)
+
+    ha_device_registry = await async_get_registry(hass)
+    reg_device = ha_device_registry.async_get_device({(DOMAIN, ieee_address)}, set())
+
+    # allow traffic to flow through the gateway and device
+    await async_enable_traffic(hass, zha_gateway, [zha_device])
+
+    with patch(
+        "zigpy.zcl.Cluster.request", return_value=mock_coro([0x00, Status.SUCCESS])
+    ):
+        assert await async_setup_component(
+            hass,
+            automation.DOMAIN,
+            {
+                automation.DOMAIN: [
+                    {
+                        "trigger": {
+                            "device_id": reg_device.id,
+                            "domain": "zha",
+                            "platform": "device",
+                            "type": SHORT_PRESS,
+                            "subtype": SHORT_PRESS,
+                        },
+                        "action": {
+                            "domain": DOMAIN,
+                            "device_id": reg_device.id,
+                            "type": "warn",
+                        },
+                    }
+                ]
+            },
+        )
+
+        await hass.async_block_till_done()
+
+        on_off_channel = zha_device.cluster_channels[CHANNEL_ON_OFF]
+        on_off_channel.zha_send_event(on_off_channel.cluster, COMMAND_SINGLE, [])
+        await hass.async_block_till_done()
+
+        assert len(calls) == 1
+        assert calls[0].domain == DOMAIN
+        assert calls[0].service == "warning_device_warn"
+        assert calls[0].data["ieee"] == ieee_address

tests/components/zha/test_device_trigger.py

@@ -1,4 +1,4 @@
-"""ZHA device automation tests."""
+"""ZHA device automation trigger tests."""
 from unittest.mock import patch
 
 import pytest