From 88821b1d0ee1c6e585acd198456f506ad8a6f13c Mon Sep 17 00:00:00 2001 From: Maciej Bieniek Date: Tue, 22 Apr 2025 12:44:02 +0200 Subject: [PATCH] Use aioshelly methods with Shelly RPC number entities (#142482) --- homeassistant/components/shelly/entity.py | 40 ++++++++- homeassistant/components/shelly/number.py | 41 +++------- tests/components/shelly/test_number.py | 98 ++++++++++++++++++----- 3 files changed, 124 insertions(+), 55 deletions(-) diff --git a/homeassistant/components/shelly/entity.py b/homeassistant/components/shelly/entity.py index 9ed3f47b41a..377479ee81c 100644 --- a/homeassistant/components/shelly/entity.py +++ b/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 diff --git a/homeassistant/components/shelly/number.py b/homeassistant/components/shelly/number.py index c629eb4a57a..83606df5a4d 100644 --- a/homeassistant/components/shelly/number.py +++ b/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, diff --git a/tests/components/shelly/test_number.py b/tests/components/shelly/test_number.py index 41002917d86..8589d643b2b 100644 --- a/tests/components/shelly/test_number.py +++ b/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