Use aioshelly methods for switches (#153746)

homeassistant/components/shelly/switch.py

@@ -4,7 +4,7 @@ from __future__ import annotations
 
 from collections.abc import Callable
 from dataclasses import dataclass
-from typing import Any, cast
+from typing import TYPE_CHECKING, Any, cast
 
 from aioshelly.block_device import Block
 from aioshelly.const import RPC_GENERATIONS
@@ -37,6 +37,7 @@ from .entity import (
     ShellySleepingBlockAttributeEntity,
     async_setup_entry_attribute_entities,
     async_setup_entry_rpc,
+    rpc_call,
 )
 from .utils import (
     async_remove_orphaned_entities,
@@ -78,7 +79,7 @@ class RpcSwitchDescription(RpcEntityDescription, SwitchEntityDescription):
     is_on: Callable[[dict[str, Any]], bool]
     method_on: str
     method_off: str
-    method_params_fn: Callable[[int | None, bool], dict]
+    method_params_fn: Callable[[int | None, bool], tuple]
 
 
 RPC_RELAY_SWITCHES = {
@@ -87,9 +88,9 @@ RPC_RELAY_SWITCHES = {
         sub_key="output",
         removal_condition=is_rpc_exclude_from_relay,
         is_on=lambda status: bool(status["output"]),
-        method_on="Switch.Set",
-        method_off="Switch.Set",
-        method_params_fn=lambda id, value: {"id": id, "on": value},
+        method_on="switch_set",
+        method_off="switch_set",
+        method_params_fn=lambda id, value: (id, value),
     ),
 }
 
@@ -101,9 +102,9 @@ RPC_SWITCHES = {
             config, key, SWITCH_PLATFORM
         ),
         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},
+        method_on="boolean_set",
+        method_off="boolean_set",
+        method_params_fn=lambda id, value: (id, value),
         role="generic",
     ),
     "boolean_anti_freeze": RpcSwitchDescription(
@@ -111,9 +112,9 @@ RPC_SWITCHES = {
         sub_key="value",
         entity_registry_enabled_default=False,
         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},
+        method_on="boolean_set",
+        method_off="boolean_set",
+        method_params_fn=lambda id, value: (id, value),
         role="anti_freeze",
         models={MODEL_LINKEDGO_ST802_THERMOSTAT, MODEL_LINKEDGO_ST1820_THERMOSTAT},
     ),
@@ -121,9 +122,9 @@ RPC_SWITCHES = {
         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},
+        method_on="boolean_set",
+        method_off="boolean_set",
+        method_params_fn=lambda id, value: (id, value),
         role="child_lock",
         models={MODEL_LINKEDGO_ST1820_THERMOSTAT},
     ),
@@ -132,9 +133,9 @@ RPC_SWITCHES = {
         sub_key="value",
         entity_registry_enabled_default=False,
         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},
+        method_on="boolean_set",
+        method_off="boolean_set",
+        method_params_fn=lambda id, value: (id, value),
         role="enable",
         models={MODEL_LINKEDGO_ST802_THERMOSTAT, MODEL_LINKEDGO_ST1820_THERMOSTAT},
     ),
@@ -142,9 +143,9 @@ RPC_SWITCHES = {
         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},
+        method_on="boolean_set",
+        method_off="boolean_set",
+        method_params_fn=lambda id, value: (id, value),
         role="start_charging",
         models={MODEL_TOP_EV_CHARGER_EVE01},
     ),
@@ -153,9 +154,9 @@ RPC_SWITCHES = {
         sub_key="value",
         entity_registry_enabled_default=False,
         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},
+        method_on="boolean_set",
+        method_off="boolean_set",
+        method_params_fn=lambda id, value: (id, value),
         role="state",
         models={MODEL_NEO_WATER_VALVE},
     ),
@@ -163,9 +164,9 @@ RPC_SWITCHES = {
         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},
+        method_on="boolean_set",
+        method_off="boolean_set",
+        method_params_fn=lambda id, value: (id, value),
         role="zone0",
         models={MODEL_FRANKEVER_IRRIGATION_CONTROLLER},
     ),
@@ -173,9 +174,9 @@ RPC_SWITCHES = {
         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},
+        method_on="boolean_set",
+        method_off="boolean_set",
+        method_params_fn=lambda id, value: (id, value),
         role="zone1",
         models={MODEL_FRANKEVER_IRRIGATION_CONTROLLER},
     ),
@@ -183,9 +184,9 @@ RPC_SWITCHES = {
         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},
+        method_on="boolean_set",
+        method_off="boolean_set",
+        method_params_fn=lambda id, value: (id, value),
         role="zone2",
         models={MODEL_FRANKEVER_IRRIGATION_CONTROLLER},
     ),
@@ -193,9 +194,9 @@ RPC_SWITCHES = {
         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},
+        method_on="boolean_set",
+        method_off="boolean_set",
+        method_params_fn=lambda id, value: (id, value),
         role="zone3",
         models={MODEL_FRANKEVER_IRRIGATION_CONTROLLER},
     ),
@@ -203,9 +204,9 @@ RPC_SWITCHES = {
         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},
+        method_on="boolean_set",
+        method_off="boolean_set",
+        method_params_fn=lambda id, value: (id, value),
         role="zone4",
         models={MODEL_FRANKEVER_IRRIGATION_CONTROLLER},
     ),
@@ -213,9 +214,9 @@ RPC_SWITCHES = {
         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},
+        method_on="boolean_set",
+        method_off="boolean_set",
+        method_params_fn=lambda id, value: (id, value),
         role="zone5",
         models={MODEL_FRANKEVER_IRRIGATION_CONTROLLER},
     ),
@@ -223,9 +224,9 @@ RPC_SWITCHES = {
         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},
+        method_on="script_start",
+        method_off="script_stop",
+        method_params_fn=lambda id, _: (id,),
         entity_registry_enabled_default=False,
         entity_category=EntityCategory.CONFIG,
     ),
@@ -422,19 +423,27 @@ class RpcSwitch(ShellyRpcAttributeEntity, SwitchEntity):
         """If switch is on."""
         return self.entity_description.is_on(self.status)
 
+    @rpc_call
     async def async_turn_on(self, **kwargs: Any) -> None:
-        """Turn on relay."""
-        await self.call_rpc(
-            self.entity_description.method_on,
-            self.entity_description.method_params_fn(self._id, True),
-        )
+        """Turn on switch."""
+        method = getattr(self.coordinator.device, self.entity_description.method_on)
 
+        if TYPE_CHECKING:
+            assert method is not None
+
+        params = self.entity_description.method_params_fn(self._id, True)
+        await method(*params)
+
+    @rpc_call
     async def async_turn_off(self, **kwargs: Any) -> None:
-        """Turn off relay."""
-        await self.call_rpc(
-            self.entity_description.method_off,
-            self.entity_description.method_params_fn(self._id, False),
-        )
+        """Turn off switch."""
+        method = getattr(self.coordinator.device, self.entity_description.method_off)
+
+        if TYPE_CHECKING:
+            assert method is not None
+
+        params = self.entity_description.method_params_fn(self._id, False)
+        await method(*params)
 
 
 class RpcRelaySwitch(RpcSwitch):

tests/components/shelly/test_switch.py

@@ -404,6 +404,7 @@ async def test_rpc_device_services(
     )
     assert (state := hass.states.get(entity_id))
     assert state.state == STATE_ON
+    mock_rpc_device.switch_set.assert_called_once_with(0, True)
 
     monkeypatch.setitem(mock_rpc_device.status["switch:0"], "output", False)
     await hass.services.async_call(
@@ -415,6 +416,7 @@ async def test_rpc_device_services(
     mock_rpc_device.mock_update()
     assert (state := hass.states.get(entity_id))
     assert state.state == STATE_OFF
+    mock_rpc_device.switch_set.assert_called_with(0, False)
 
 
 async def test_rpc_device_unique_ids(
@@ -507,7 +509,7 @@ async def test_rpc_set_state_errors(
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
     """Test RPC device set state connection/call errors."""
-    monkeypatch.setattr(mock_rpc_device, "call_rpc", AsyncMock(side_effect=exc))
+    mock_rpc_device.switch_set.side_effect = exc
     monkeypatch.delitem(mock_rpc_device.status, "cover:0")
     monkeypatch.setitem(mock_rpc_device.status["sys"], "relay_in_thermostat", False)
     await init_integration(hass, 2)
@@ -525,11 +527,7 @@ async def test_rpc_auth_error(
     hass: HomeAssistant, mock_rpc_device: Mock, monkeypatch: pytest.MonkeyPatch
 ) -> None:
     """Test RPC device set state authentication error."""
-    monkeypatch.setattr(
-        mock_rpc_device,
-        "call_rpc",
-        AsyncMock(side_effect=InvalidAuthError),
-    )
+    mock_rpc_device.switch_set.side_effect = InvalidAuthError
     monkeypatch.delitem(mock_rpc_device.status, "cover:0")
     monkeypatch.setitem(mock_rpc_device.status["sys"], "relay_in_thermostat", False)
     entry = await init_integration(hass, 2)
@@ -657,6 +655,7 @@ async def test_rpc_device_virtual_switch(
     mock_rpc_device.mock_update()
     assert (state := hass.states.get(entity_id))
     assert state.state == STATE_OFF
+    mock_rpc_device.boolean_set.assert_called_once_with(200, False)
 
     monkeypatch.setitem(mock_rpc_device.status["boolean:200"], "value", True)
     await hass.services.async_call(
@@ -668,6 +667,7 @@ async def test_rpc_device_virtual_switch(
     mock_rpc_device.mock_update()
     assert (state := hass.states.get(entity_id))
     assert state.state == STATE_ON
+    mock_rpc_device.boolean_set.assert_called_with(200, True)
 
 
 @pytest.mark.usefixtures("disable_async_remove_shelly_rpc_entities")
@@ -815,6 +815,7 @@ async def test_rpc_device_script_switch(
 
     assert (state := hass.states.get(entity_id))
     assert state.state == STATE_OFF
+    mock_rpc_device.script_stop.assert_called_once_with(1)
 
     monkeypatch.setitem(mock_rpc_device.status[key], "running", True)
     await hass.services.async_call(
@@ -827,3 +828,4 @@ async def test_rpc_device_script_switch(
 
     assert (state := hass.states.get(entity_id))
     assert state.state == STATE_ON
+    mock_rpc_device.script_start.assert_called_once_with(1)