Start mocking out MQTT support

Author: wbyoung

custom_components/lampie/__init__.py

@@ -45,10 +45,13 @@ async def async_setup_entry(hass: HomeAssistant, entry: LampieConfigEntry) -> bo
     _LOGGER.debug("setup %s with config:%s", entry.title, entry.data)
 
     if DOMAIN not in hass.data:
-        hass.data[DOMAIN] = LampieOrchestrator(hass)
+        orchestrator = LampieOrchestrator(hass)
+        hass.data[DOMAIN] = orchestrator
+
+        await orchestrator.setup()
 
     coordinator = LampieUpdateCoordinator(hass, entry)
-    orchestrator: LampieOrchestrator = hass.data[DOMAIN]
+    orchestrator = hass.data[DOMAIN]
     orchestrator.add_coordinator(coordinator)
     entry.runtime_data = LampieConfigEntryRuntimeData(
         orchestrator=orchestrator,

custom_components/lampie/orchestrator.py

@@ -6,20 +6,26 @@
 from contextlib import suppress
 from dataclasses import replace
 import datetime as dt
-from enum import IntEnum
+from enum import Enum, IntEnum, auto
 from functools import partial
 import logging
 import re
 from typing import TYPE_CHECKING, Any, Final, NamedTuple, Protocol, Unpack
 
+from homeassistant.components import mqtt
 from homeassistant.components.script import DOMAIN as SCRIPT_DOMAIN
 from homeassistant.core import CALLBACK_TYPE, Event, HomeAssistant, callback
-from homeassistant.helpers import entity_registry as er, event as evt
+from homeassistant.helpers import (
+    device_registry as dr,
+    entity_registry as er,
+    event as evt,
+)
 from homeassistant.helpers.device import (
     async_device_info_to_link_from_entity,
     async_entity_id_to_device_id,
 )
 from homeassistant.util import dt as dt_util
+from homeassistant.util.json import json_loads_object
 
 from .const import (
     CONF_END_ACTION,
@@ -33,6 +39,7 @@
     DeviceId,
     Effect,
     ExpirationInfo,
+    Integration,
     LampieNotificationInfo,
     LampieNotificationOptionsDict,
     LampieSwitchInfo,
@@ -48,12 +55,19 @@
     from .coordinator import LampieUpdateCoordinator
 
 type ZHAEventData = dict[str, Any]
+type MQTTDeviceName = str
 
 _LOGGER = logging.getLogger(__name__)
 
+MQTT_DOMAIN: Final = "mqtt"
 ZHA_DOMAIN: Final = "zha"
 ALREADY_EXPIRED: Final = 0
 
+SWITCH_INTEGRATIONS = {
+    ZHA_DOMAIN: Integration.ZHA,
+    MQTT_DOMAIN: Integration.Z2M,
+}
+
 FIRMWARE_SECONDS_MAX = dt.timedelta(seconds=60).total_seconds()
 FIRMWARE_MINUTES_MAX = dt.timedelta(minutes=60).total_seconds()
 FIRMWARE_HOURS_MAX = dt.timedelta(hours=134).total_seconds()
@@ -86,12 +100,26 @@
     "button_6_double",
 }
 
+Z2M_COMMAND_MAP = {
+    "config_double": "button_3_double",
+}
+
 
 class _LEDMode(IntEnum):
     ALL = 1
     INDIVIDUAL = 3
 
 
+class _SwitchKeyType(Enum):
+    DEVICE_ID = auto()
+    MQTT_NAME = auto()
+
+
+class _SwitchKey[T](NamedTuple):
+    type: _SwitchKeyType
+    identifier: T
+
+
 class _StartScriptResult(NamedTuple):
     led_config: tuple[LEDConfig, ...] | None
     block_activation: bool
@@ -102,6 +130,10 @@ class _EndScriptResult(NamedTuple):
     block_next: bool
 
 
+class _UnknownIntegrationError(Exception):
+    pass
+
+
 class _LampieUnmanagedSwitchCoordinator:
     def async_update_listeners(self) -> None:
         pass
@@ -130,7 +162,8 @@ def __init__(self, hass: HomeAssistant) -> None:
         self._coordinators: dict[Slug, LampieUpdateCoordinator] = {}
         self._notifications: dict[Slug, LampieNotificationInfo] = {}
         self._switches: dict[SwitchId, LampieSwitchInfo] = {}
-        self._device_switches: dict[DeviceId, SwitchId] = {}
+        self._switch_ids: dict[_SwitchKey[DeviceId | MQTTDeviceName], SwitchId] = {}
+        self._cancel_mqtt_listener: CALLBACK_TYPE | None = None
         self._cancel_zha_listener: CALLBACK_TYPE = hass.bus.async_listen(
             "zha_event",
             self._handle_zha_event,
@@ -145,9 +178,20 @@ def remove_coordinator(self, coordinator: LampieUpdateCoordinator) -> None:
         self._coordinators.pop(coordinator.slug)
         self._update_references()
 
+    async def setup(self) -> None:
+        self._cancel_mqtt_listener = await mqtt.async_subscribe(
+            self._hass,
+            "zigbee2mqtt/+",
+            self._handle_z2m_message,
+        )
+
     def teardown(self) -> bool:
         if len(self._coordinators) == 0:
             self._cancel_zha_listener()
+
+            if self._cancel_mqtt_listener:
+                self._cancel_mqtt_listener()
+
             for key, expiration in (
                 *((key, info.expiration) for key, info in self._notifications.items()),
                 *((key, info.expiration) for key, info in self._switches.items()),
@@ -159,11 +203,17 @@ def teardown(self) -> bool:
     def switch_info(self, switch_id: SwitchId) -> LampieSwitchInfo:
         if switch_id not in self._switches:
             entity_registry = er.async_get(self._hass)
+            device_registry = dr.async_get(self._hass)
             device_id = async_entity_id_to_device_id(self._hass, switch_id)
+            device = device_registry.async_get(device_id)
+            integration = self._switch_integration(device) if device else None
             entity_entries = er.async_entries_for_device(entity_registry, device_id)
             local_protetction_id = None
             disable_clear_notification_id = None
 
+            if not integration:
+                raise _UnknownIntegrationError
+
             for entity_entry in entity_entries:
                 if entity_entry.unique_id.endswith("-local_protection"):
                     local_protetction_id = entity_entry.entity_id
@@ -172,8 +222,20 @@ def switch_info(self, switch_id: SwitchId) -> LampieSwitchInfo:
                 ):
                     disable_clear_notification_id = entity_entry.entity_id
 
-            self._device_switches[device_id] = switch_id
+            if (
+                device
+                and (mqtt_device_name := self._mqtt_device_name(device)) is not None
+            ):
+                self._switch_ids[
+                    _SwitchKey(_SwitchKeyType.MQTT_NAME, mqtt_device_name)
+                ] = switch_id
+
+            self._switch_ids[_SwitchKey(_SwitchKeyType.DEVICE_ID, device_id)] = (
+                switch_id
+            )
+
             self._switches[switch_id] = LampieSwitchInfo(
+                integration=integration,
                 led_config=(),
                 led_config_source=LEDConfigSource(None),
                 local_protetction_id=local_protetction_id,
@@ -712,9 +774,8 @@ async def _issue_switch_commands(
             "all" if led_mode == _LEDMode.ALL else ", ".join(updated_leds),
         )
 
-    @classmethod
-    def _switch_command_led_params(cls, led: LEDConfig) -> dict[str, Any]:
-        firmware_duration = cls._firmware_duration(led.duration)
+    def _switch_command_led_params(self, led: LEDConfig) -> dict[str, Any]:
+        firmware_duration = self._firmware_duration(led.duration)
 
         return {
             "led_color": int(led.color),
@@ -727,17 +788,24 @@ def _switch_command_led_params(cls, led: LEDConfig) -> dict[str, Any]:
             else firmware_duration,
         }
 
-    @classmethod
-    def _firmware_duration(cls, seconds: int | None) -> int | None:
+    def _firmware_duration(self, seconds: int | None) -> int | None:
         """Convert a timeframe to a duration supported by the switch firmware.
 
+        Note: Any usage of MQTT switches in the system will result in the
+        suspension of all firmware durations, and this method will always return
+        `None`. At the moment Zigbee2MQTT does not support the 0x24 command
+        which is used for `led_effect_complete`, so there is no way to know the
+        notification ended.
+
         Args:
             seconds: The duration as a number of seconds.
 
         Returns:
             The duration parameter value (0-255) if it can be handled by the
                 firmware or None if it cannot be.
         """
+        if any(key.type == _SwitchKeyType.MQTT_NAME for key in self._switch_ids):
+            return None
         if seconds is None or seconds == ALREADY_EXPIRED:
             return None
         if seconds <= FIRMWARE_SECONDS_MAX:
@@ -753,17 +821,40 @@ def _filter_zha_events(
         self,
         event_data: ZHAEventData,
     ) -> bool:
+        switch_key = _SwitchKey(_SwitchKeyType.DEVICE_ID, event_data["device_id"])
         return (
-            event_data["device_id"] in self._device_switches
+            switch_key in self._switch_ids
             and event_data["command"] in DISMISSAL_COMMANDS
         )
 
     @callback
     async def _handle_zha_event(self, event: Event[ZHAEventData]) -> None:
+        await self._handle_generic_event(
+            command=event.data["command"],
+            device_id=event.data["device_id"],
+        )
+
+    @callback
+    async def _handle_z2m_message(self, message: mqtt.ReceiveMessage) -> None:
+        command_path_base, device_name = message.topic.split("/", 1)
+        switch_key = _SwitchKey(_SwitchKeyType.MQTT_NAME, device_name)
+        device_id = self._switch_ids.get(switch_key)
+
+        if command_path_base == "zigbee2mqtt" and device_id is not None:
+            payload = json_loads_object(message.payload)
+            action = payload.get("action")
+            command = Z2M_COMMAND_MAP.get(action)
+
+            if command:
+                await self._handle_generic_event(
+                    command=command,
+                    device_id=device_id,
+                )
+
+    async def _handle_generic_event(self, command: str, device_id: str) -> None:
         hass = self._hass
-        command = event.data["command"]
-        device_id = event.data["device_id"]
-        switch_id = self._device_switches[device_id]
+        switch_key = _SwitchKey(_SwitchKeyType.DEVICE_ID, device_id)
+        switch_id = self._switch_ids[switch_key]
         from_state = self.switch_info(switch_id)
         led_config_source = from_state.led_config_source
         led_config = [*from_state.led_config]
@@ -1090,8 +1181,17 @@ def _update_references(self) -> None:
             )
 
             for switch_id in switch_ids:
+                try:
+                    switch_info = self.switch_info(switch_id)
+                except _UnknownIntegrationError:
+                    _LOGGER.exception(
+                        "ignoring switch %s: could not to a valid integration",
+                        switch_id,
+                    )
+                    continue
+
                 priorities = switch_priorities.get(switch_id) or [slug]
-                expected = [*self.switch_info(switch_id).priorities]
+                expected = [*switch_info.priorities]
 
                 if switch_id in processed_switches and expected != priorities:
                     _LOGGER.warning(
@@ -1109,13 +1209,34 @@ def _update_references(self) -> None:
                     continue
 
                 self._switches[switch_id] = replace(
-                    self.switch_info(switch_id),
+                    switch_info,
                     priorities=tuple(priorities),
                 )
                 processed_switches.add(switch_id)
 
             processed_slugs.append(slug)
 
+    @classmethod
+    def _switch_integration(cls, device: dr.DeviceEnttry) -> Integration | None:
+        id_tuple = next(iter(device.identifiers))
+        domain = id_tuple[0]
+        return SWITCH_INTEGRATIONS.get(domain)
+
+    @classmethod
+    def _mqtt_device_id(cls, device: dr.DeviceEnttry) -> str | None:
+        id_tuple = next(iter(device.identifiers))
+
+        if id_tuple[0] != MQTT_DOMAIN:
+            return None
+
+        identifier: str = id_tuple[1]
+
+        return identifier.split("_", 1)[1]
+
+    @classmethod
+    def _mqtt_device_name(cls, device: dr.DeviceEnttry) -> str | None:
+        return device.name if cls._mqtt_device_id(device) is not None else None
+
 
 def _all_clear(led_config: Sequence[LEDConfig]) -> bool:
     return all(item.effect == Effect.CLEAR for item in led_config)

custom_components/lampie/services.yaml

@@ -16,7 +16,6 @@ override:
         entity:
           multiple: true
           filter:
-            integration: zha
             domain:
               - light
               - fan

custom_components/lampie/types.py

@@ -219,6 +219,13 @@ class InvalidColor(Exception):
     index: int | None = None
 
 
+class Integration(StrEnum):
+    """Switch integration type."""
+
+    ZHA = auto()
+    Z2M = auto()
+
+
 @dataclass(frozen=True)
 class ExpirationInfo:
     """Storage of expiration info."""
@@ -255,6 +262,7 @@ class LampieSwitchInfo:
     disable_clear_notification_id: EntityId | None = None
     priorities: tuple[Slug, ...] = field(default_factory=tuple)
     expiration: ExpirationInfo = field(default_factory=ExpirationInfo)
+    integration: Integration = Integration.ZHA
 
 
 class LampieSwitchOptionsDict(TypedDict):

tests/conftest.py

@@ -1,6 +1,8 @@
 """Fixtures for testing."""
 
+from collections.abc import Generator
 import logging
+from unittest.mock import AsyncMock, Mock, patch
 
 from freezegun.api import FrozenDateTimeFactory
 from homeassistant.core import HomeAssistant
@@ -52,6 +54,19 @@ def auto_enable_custom_integrations(enable_custom_integrations):
     return
 
 
+@pytest.fixture(name="mqtt_subscribe", autouse=True)
+def auto_patch_mqtt_async_subscribe() -> Generator[AsyncMock]:
+    """Patch mqtt.async_subscribe."""
+
+    unsub = Mock()
+
+    with patch(
+        "homeassistant.components.mqtt.async_subscribe", return_value=unsub
+    ) as mock_subscribe:
+        mock_subscribe._unsub = unsub
+        yield mock_subscribe
+
+
 @pytest.fixture
 def snapshot(snapshot: SnapshotAssertion) -> SnapshotAssertion:
     """Return snapshot assertion fixture with the Home Assistant extension."""

tests/snapshots/test_diagnostics.ambr

@@ -93,6 +93,7 @@
           'cancel_listener': None,
           'started_at': None,
         }),
+        'integration': 'zha',
         'led_config': list([
           dict({
             'brightness': 100.0,