diff --git a/homeassistant/components/shelly/switch.py b/homeassistant/components/shelly/switch.py index 9b34b2e079b..41826706945 100644 --- a/homeassistant/components/shelly/switch.py +++ b/homeassistant/components/shelly/switch.py @@ -2,6 +2,7 @@ from __future__ import annotations +from collections.abc import Callable from dataclasses import dataclass from typing import Any, cast @@ -29,7 +30,7 @@ from .entity import ( ShellyRpcEntity, ShellySleepingBlockAttributeEntity, async_setup_entry_attribute_entities, - async_setup_rpc_attribute_entities, + async_setup_entry_rpc, ) from .utils import ( async_remove_orphaned_entities, @@ -60,18 +61,32 @@ MOTION_SWITCH = BlockSwitchDescription( class RpcSwitchDescription(RpcEntityDescription, SwitchEntityDescription): """Class to describe a RPC virtual switch.""" + is_on: Callable[[dict[str, Any]], bool] + method_on: str + method_off: str + method_params_fn: Callable[[int | None, bool], dict] -RPC_VIRTUAL_SWITCH = RpcSwitchDescription( - key="boolean", - sub_key="value", -) -RPC_SCRIPT_SWITCH = RpcSwitchDescription( - key="script", - sub_key="running", - entity_registry_enabled_default=False, - entity_category=EntityCategory.CONFIG, -) +RPC_SWITCHES = { + "boolean": RpcSwitchDescription( + key="boolean", + sub_key="value", + is_on=lambda status: bool(status["value"]), + method_on="Boolean.Set", + method_off="Boolean.Set", + method_params_fn=lambda id, value: {"id": id, "value": value}, + ), + "script": RpcSwitchDescription( + key="script", + sub_key="running", + is_on=lambda status: bool(status["running"]), + method_on="Script.Start", + method_off="Script.Stop", + method_params_fn=lambda id, _: {"id": id}, + entity_registry_enabled_default=False, + entity_category=EntityCategory.CONFIG, + ), +} async def async_setup_entry( @@ -174,20 +189,8 @@ def async_setup_rpc_entry( unique_id = f"{coordinator.mac}-switch:{id_}" async_remove_shelly_entity(hass, "light", unique_id) - async_setup_rpc_attribute_entities( - hass, - config_entry, - async_add_entities, - {"boolean": RPC_VIRTUAL_SWITCH}, - RpcVirtualSwitch, - ) - - async_setup_rpc_attribute_entities( - hass, - config_entry, - async_add_entities, - {"script": RPC_SCRIPT_SWITCH}, - RpcScriptSwitch, + async_setup_entry_rpc( + hass, config_entry, async_add_entities, RPC_SWITCHES, RpcSwitch ) # the user can remove virtual components from the device configuration, so we need @@ -324,8 +327,8 @@ class RpcRelaySwitch(ShellyRpcEntity, SwitchEntity): await self.call_rpc("Switch.Set", {"id": self._id, "on": False}) -class RpcVirtualSwitch(ShellyRpcAttributeEntity, SwitchEntity): - """Entity that controls a virtual boolean component on RPC based Shelly devices.""" +class RpcSwitch(ShellyRpcAttributeEntity, SwitchEntity): + """Entity that controls a switch on RPC based Shelly devices.""" entity_description: RpcSwitchDescription _attr_has_entity_name = True @@ -333,32 +336,18 @@ class RpcVirtualSwitch(ShellyRpcAttributeEntity, SwitchEntity): @property def is_on(self) -> bool: """If switch is on.""" - return bool(self.attribute_value) + return self.entity_description.is_on(self.status) async def async_turn_on(self, **kwargs: Any) -> None: """Turn on relay.""" - await self.call_rpc("Boolean.Set", {"id": self._id, "value": True}) + await self.call_rpc( + self.entity_description.method_on, + self.entity_description.method_params_fn(self._id, True), + ) async def async_turn_off(self, **kwargs: Any) -> None: """Turn off relay.""" - await self.call_rpc("Boolean.Set", {"id": self._id, "value": False}) - - -class RpcScriptSwitch(ShellyRpcAttributeEntity, SwitchEntity): - """Entity that controls a script component on RPC based Shelly devices.""" - - entity_description: RpcSwitchDescription - _attr_has_entity_name = True - - @property - def is_on(self) -> bool: - """If switch is on.""" - return bool(self.status["running"]) - - async def async_turn_on(self, **kwargs: Any) -> None: - """Turn on relay.""" - await self.call_rpc("Script.Start", {"id": self._id}) - - async def async_turn_off(self, **kwargs: Any) -> None: - """Turn off relay.""" - await self.call_rpc("Script.Stop", {"id": self._id}) + await self.call_rpc( + self.entity_description.method_off, + self.entity_description.method_params_fn(self._id, False), + )