Use aioshelly methods with Shelly RPC number entities (#142482)

homeassistant/components/shelly/entity.py

@@ -2,9 +2,10 @@
 
 from __future__ import annotations
 
-from collections.abc import Callable, Mapping
+from collections.abc import Awaitable, Callable, Coroutine, Mapping
 from dataclasses import dataclass
-from typing import Any, cast
+from functools import wraps
+from typing import Any, Concatenate, cast
 
 from aioshelly.block_device import Block
 from aioshelly.exceptions import DeviceConnectionError, InvalidAuthError, RpcCallError
@@ -707,3 +708,38 @@ def get_entity_class(
         return description.entity_class
 
     return sensor_class
+
+
+def rpc_call[_T: ShellyRpcEntity, **_P](
+    func: Callable[Concatenate[_T, _P], Awaitable[None]],
+) -> Callable[Concatenate[_T, _P], Coroutine[Any, Any, None]]:
+    """Catch rpc_call exceptions."""
+
+    @wraps(func)
+    async def cmd_wrapper(self: _T, *args: _P.args, **kwargs: _P.kwargs) -> None:
+        """Wrap all command methods."""
+        try:
+            await func(self, *args, **kwargs)
+        except DeviceConnectionError as err:
+            self.coordinator.last_update_success = False
+            raise HomeAssistantError(
+                translation_domain=DOMAIN,
+                translation_key="device_communication_action_error",
+                translation_placeholders={
+                    "entity": self.entity_id,
+                    "device": self.coordinator.name,
+                },
+            ) from err
+        except RpcCallError as err:
+            raise HomeAssistantError(
+                translation_domain=DOMAIN,
+                translation_key="rpc_call_action_error",
+                translation_placeholders={
+                    "entity": self.entity_id,
+                    "device": self.coordinator.name,
+                },
+            ) from err
+        except InvalidAuthError:
+            await self.coordinator.async_shutdown_device_and_start_reauth()
+
+    return cmd_wrapper

homeassistant/components/shelly/number.py

@@ -7,7 +7,7 @@ from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any, Final, cast
 
 from aioshelly.block_device import Block
-from aioshelly.const import BLU_TRV_TIMEOUT, RPC_GENERATIONS
+from aioshelly.const import RPC_GENERATIONS
 from aioshelly.exceptions import DeviceConnectionError, InvalidAuthError
 
 from homeassistant.components.number import (
@@ -34,6 +34,7 @@ from .entity import (
     ShellySleepingBlockAttributeEntity,
     async_setup_entry_attribute_entities,
     async_setup_entry_rpc,
+    rpc_call,
 )
 from .utils import (
     async_remove_orphaned_entities,
@@ -59,7 +60,6 @@ class RpcNumberDescription(RpcEntityDescription, NumberEntityDescription):
     step_fn: Callable[[dict], float] | None = None
     mode_fn: Callable[[dict], NumberMode] | None = None
     method: str
-    method_params_fn: Callable[[int, float], dict]
 
 
 class RpcNumber(ShellyRpcAttributeEntity, NumberEntity):
@@ -98,15 +98,16 @@ class RpcNumber(ShellyRpcAttributeEntity, NumberEntity):
 
         return self.attribute_value
 
+    @rpc_call
     async def async_set_native_value(self, value: float) -> None:
         """Change the value."""
+        method = getattr(self.coordinator.device, self.entity_description.method)
+
         if TYPE_CHECKING:
             assert isinstance(self._id, int)
+            assert method is not None
 
-        await self.call_rpc(
-            self.entity_description.method,
-            self.entity_description.method_params_fn(self._id, value),
-        )
+        await method(self._id, value)
 
 
 class RpcBluTrvNumber(RpcNumber):
@@ -127,17 +128,6 @@ class RpcBluTrvNumber(RpcNumber):
             connections={(CONNECTION_BLUETOOTH, ble_addr)}
         )
 
-    async def async_set_native_value(self, value: float) -> None:
-        """Change the value."""
-        if TYPE_CHECKING:
-            assert isinstance(self._id, int)
-
-        await self.call_rpc(
-            self.entity_description.method,
-            self.entity_description.method_params_fn(self._id, value),
-            timeout=BLU_TRV_TIMEOUT,
-        )
-
 
 class RpcBluTrvExtTempNumber(RpcBluTrvNumber):
     """Represent a RPC BluTrv External Temperature number."""
@@ -187,12 +177,7 @@ RPC_NUMBERS: Final = {
         mode=NumberMode.BOX,
         entity_category=EntityCategory.CONFIG,
         native_unit_of_measurement=UnitOfTemperature.CELSIUS,
-        method="BluTRV.Call",
-        method_params_fn=lambda idx, value: {
-            "id": idx,
-            "method": "Trv.SetExternalTemperature",
-            "params": {"id": 0, "t_C": value},
-        },
+        method="blu_trv_set_external_temperature",
         entity_class=RpcBluTrvExtTempNumber,
     ),
     "number": RpcNumberDescription(
@@ -209,8 +194,7 @@ RPC_NUMBERS: Final = {
         unit=lambda config: config["meta"]["ui"]["unit"]
         if config["meta"]["ui"]["unit"]
         else None,
-        method="Number.Set",
-        method_params_fn=lambda idx, value: {"id": idx, "value": value},
+        method="number_set",
     ),
     "valve_position": RpcNumberDescription(
         key="blutrv",
@@ -222,12 +206,7 @@ RPC_NUMBERS: Final = {
         native_step=1,
         mode=NumberMode.SLIDER,
         native_unit_of_measurement=PERCENTAGE,
-        method="BluTRV.Call",
-        method_params_fn=lambda idx, value: {
-            "id": idx,
-            "method": "Trv.SetPosition",
-            "params": {"id": 0, "pos": int(value)},
-        },
+        method="blu_trv_set_valve_position",
         removal_condition=lambda config, _status, key: config[key].get("enable", True)
         is True,
         entity_class=RpcBluTrvNumber,

tests/components/shelly/test_number.py

@@ -3,8 +3,8 @@
 from copy import deepcopy
 from unittest.mock import AsyncMock, Mock
 
-from aioshelly.const import BLU_TRV_TIMEOUT, MODEL_BLU_GATEWAY_G3
-from aioshelly.exceptions import DeviceConnectionError, InvalidAuthError
+from aioshelly.const import MODEL_BLU_GATEWAY_G3
+from aioshelly.exceptions import DeviceConnectionError, InvalidAuthError, RpcCallError
 import pytest
 from syrupy import SnapshotAssertion
 
@@ -334,6 +334,8 @@ async def test_rpc_device_virtual_number(
         blocking=True,
     )
     mock_rpc_device.mock_update()
+    mock_rpc_device.number_set.assert_called_once_with(203, 56.7)
+
     assert (state := hass.states.get(entity_id))
     assert state.state == "56.7"
 
@@ -446,15 +448,7 @@ async def test_blu_trv_ext_temp_set_value(
         blocking=True,
     )
     mock_blu_trv.mock_update()
-    mock_blu_trv.call_rpc.assert_called_once_with(
-        "BluTRV.Call",
-        {
-            "id": 200,
-            "method": "Trv.SetExternalTemperature",
-            "params": {"id": 0, "t_C": 22.2},
-        },
-        BLU_TRV_TIMEOUT,
-    )
+    mock_blu_trv.blu_trv_set_external_temperature.assert_called_once_with(200, 22.2)
 
     assert (state := hass.states.get(entity_id))
     assert state.state == "22.2"
@@ -487,17 +481,77 @@ async def test_blu_trv_valve_pos_set_value(
         blocking=True,
     )
     mock_blu_trv.mock_update()
-    mock_blu_trv.call_rpc.assert_called_once_with(
-        "BluTRV.Call",
-        {
-            "id": 200,
-            "method": "Trv.SetPosition",
-            "params": {"id": 0, "pos": 20},
-        },
-        BLU_TRV_TIMEOUT,
-    )
-    # device only accepts int for 'pos' value
-    assert isinstance(mock_blu_trv.call_rpc.call_args[0][1]["params"]["pos"], int)
+    mock_blu_trv.blu_trv_set_valve_position.assert_called_once_with(200, 20.0)
 
     assert (state := hass.states.get(entity_id))
     assert state.state == "20"
+
+
+@pytest.mark.parametrize(
+    ("exception", "error"),
+    [
+        (
+            DeviceConnectionError,
+            "Device communication error occurred while calling action for number.trv_name_external_temperature of Test name",
+        ),
+        (
+            RpcCallError(999),
+            "RPC call error occurred while calling action for number.trv_name_external_temperature of Test name",
+        ),
+    ],
+)
+async def test_blu_trv_number_exc(
+    hass: HomeAssistant,
+    mock_blu_trv: Mock,
+    exception: Exception,
+    error: str,
+) -> None:
+    """Test RPC/BLU TRV number with exception."""
+    await init_integration(hass, 3, model=MODEL_BLU_GATEWAY_G3)
+
+    mock_blu_trv.blu_trv_set_external_temperature.side_effect = exception
+
+    with pytest.raises(HomeAssistantError, match=error):
+        await hass.services.async_call(
+            NUMBER_DOMAIN,
+            SERVICE_SET_VALUE,
+            {
+                ATTR_ENTITY_ID: f"{NUMBER_DOMAIN}.trv_name_external_temperature",
+                ATTR_VALUE: 20.0,
+            },
+            blocking=True,
+        )
+
+
+async def test_blu_trv_number_reauth_error(
+    hass: HomeAssistant,
+    mock_blu_trv: Mock,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """Test RPC/BLU TRV number with authentication error."""
+    entry = await init_integration(hass, 3, model=MODEL_BLU_GATEWAY_G3)
+
+    mock_blu_trv.blu_trv_set_external_temperature.side_effect = InvalidAuthError
+
+    await hass.services.async_call(
+        NUMBER_DOMAIN,
+        SERVICE_SET_VALUE,
+        {
+            ATTR_ENTITY_ID: f"{NUMBER_DOMAIN}.trv_name_external_temperature",
+            ATTR_VALUE: 20.0,
+        },
+        blocking=True,
+    )
+
+    assert entry.state is ConfigEntryState.LOADED
+
+    flows = hass.config_entries.flow.async_progress()
+    assert len(flows) == 1
+
+    flow = flows[0]
+    assert flow.get("step_id") == "reauth_confirm"
+    assert flow.get("handler") == DOMAIN
+
+    assert "context" in flow
+    assert flow["context"].get("source") == SOURCE_REAUTH
+    assert flow["context"].get("entry_id") == entry.entry_id