diff --git a/homeassistant/components/roborock/switch.py b/homeassistant/components/roborock/switch.py index 6971ff9d900..d8ff50430cb 100644 --- a/homeassistant/components/roborock/switch.py +++ b/homeassistant/components/roborock/switch.py @@ -1,9 +1,11 @@ """Support for Roborock switch.""" +import asyncio from collections.abc import Callable, Coroutine from dataclasses import dataclass import logging from typing import Any +from roborock.exceptions import RoborockException from roborock.roborock_typing import RoborockCommand from homeassistant.components.switch import SwitchEntity, SwitchEntityDescription @@ -30,6 +32,8 @@ class RoborockSwitchDescriptionMixin: evaluate_value: Callable[[dict], bool] # Sets the status of the switch set_command: Callable[[RoborockEntity, bool], Coroutine[Any, Any, dict]] + # Check support of this feature + check_support: Callable[[RoborockDataUpdateCoordinator], Coroutine[Any, Any, dict]] @dataclass @@ -45,6 +49,9 @@ SWITCH_DESCRIPTIONS: list[RoborockSwitchDescription] = [ RoborockCommand.SET_CHILD_LOCK_STATUS, {"lock_status": 1 if value else 0} ), get_value=lambda data: data.send(RoborockCommand.GET_CHILD_LOCK_STATUS), + check_support=lambda data: data.api.send_command( + RoborockCommand.GET_CHILD_LOCK_STATUS + ), evaluate_value=lambda data: data["lock_status"] == 1, key="child_lock", translation_key="child_lock", @@ -56,6 +63,9 @@ SWITCH_DESCRIPTIONS: list[RoborockSwitchDescription] = [ RoborockCommand.SET_FLOW_LED_STATUS, {"status": 1 if value else 0} ), get_value=lambda data: data.send(RoborockCommand.GET_FLOW_LED_STATUS), + check_support=lambda data: data.api.send_command( + RoborockCommand.GET_FLOW_LED_STATUS + ), evaluate_value=lambda data: data["status"] == 1, key="status_indicator", translation_key="status_indicator", @@ -75,16 +85,38 @@ async def async_setup_entry( coordinators: dict[str, RoborockDataUpdateCoordinator] = hass.data[DOMAIN][ config_entry.entry_id ] - async_add_entities( - ( - RoborockSwitchEntity( - f"{description.key}_{slugify(device_id)}", - coordinator, - description, - ) - for device_id, coordinator in coordinators.items() - for description in SWITCH_DESCRIPTIONS + possible_entities: list[ + tuple[str, RoborockDataUpdateCoordinator, RoborockSwitchDescription] + ] = [ + (device_id, coordinator, description) + for device_id, coordinator in coordinators.items() + for description in SWITCH_DESCRIPTIONS + ] + # We need to check if this function is supported by the device. + results = await asyncio.gather( + *( + description.check_support(coordinator) + for _, coordinator, description in possible_entities ), + return_exceptions=True, + ) + valid_entities: list[RoborockSwitchEntity] = [] + for posible_entity, result in zip(possible_entities, results): + if isinstance(result, Exception): + if not isinstance(result, RoborockException): + raise result + _LOGGER.debug("Not adding entity because of %s", result) + else: + valid_entities.append( + RoborockSwitchEntity( + f"{posible_entity[2].key}_{slugify(posible_entity[0])}", + posible_entity[1], + posible_entity[2], + result, + ) + ) + async_add_entities( + valid_entities, True, ) @@ -99,10 +131,12 @@ class RoborockSwitchEntity(RoborockEntity, SwitchEntity): unique_id: str, coordinator: RoborockDataUpdateCoordinator, entity_description: RoborockSwitchDescription, + initial_value: bool, ) -> None: """Create a switch entity.""" self.entity_description = entity_description super().__init__(unique_id, coordinator.device_info, coordinator.api) + self._attr_is_on = initial_value async def async_turn_off(self, **kwargs: Any) -> None: """Turn off the switch."""