Add calibrate button for Shelly BLU TRV (#140578)

* Initial commit

* Refactor

* Call async_add_entities() once

* Type

* Cleaning

* `supported` is not needed here

* Add error handling

* Add test

* Fix name

* Change class name

* Change method name

* Move BLU_TRV_TIMEOUT

* Fix BLU_TRV_TIMEOUT import

* Coverage

* Use test snapshots

* Support error translations

* Fix tests

* Introduce ShellyBaseButton class

* Rename press_method to _press_method

* Improve exception strings
This commit is contained in:
Maciej Bieniek 2025-03-21 10:14:20 +01:00 committed by GitHub
parent bce7fcc3c6
commit 2785688f57
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
9 changed files with 426 additions and 29 deletions

View File

@ -2,12 +2,13 @@
from __future__ import annotations
from collections.abc import Callable, Coroutine
from collections.abc import Callable
from dataclasses import dataclass
from functools import partial
from typing import TYPE_CHECKING, Any, Final
from aioshelly.const import RPC_GENERATIONS
from aioshelly.const import BLU_TRV_IDENTIFIER, MODEL_BLU_GATEWAY, RPC_GENERATIONS
from aioshelly.exceptions import DeviceConnectionError, InvalidAuthError, RpcCallError
from homeassistant.components.button import (
ButtonDeviceClass,
@ -16,15 +17,20 @@ from homeassistant.components.button import (
)
from homeassistant.const import EntityCategory
from homeassistant.core import HomeAssistant, callback
from homeassistant.exceptions import HomeAssistantError
from homeassistant.helpers import entity_registry as er
from homeassistant.helpers.device_registry import CONNECTION_NETWORK_MAC, DeviceInfo
from homeassistant.helpers.device_registry import (
CONNECTION_BLUETOOTH,
CONNECTION_NETWORK_MAC,
DeviceInfo,
)
from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
from homeassistant.helpers.update_coordinator import CoordinatorEntity
from homeassistant.util import slugify
from .const import LOGGER, SHELLY_GAS_MODELS
from .const import DOMAIN, LOGGER, SHELLY_GAS_MODELS
from .coordinator import ShellyBlockCoordinator, ShellyConfigEntry, ShellyRpcCoordinator
from .utils import get_device_entry_gen
from .utils import get_device_entry_gen, get_rpc_key_ids
@dataclass(frozen=True, kw_only=True)
@ -33,7 +39,7 @@ class ShellyButtonDescription[
](ButtonEntityDescription):
"""Class to describe a Button entity."""
press_action: Callable[[_ShellyCoordinatorT], Coroutine[Any, Any, None]]
press_action: str
supported: Callable[[_ShellyCoordinatorT], bool] = lambda _: True
@ -44,14 +50,14 @@ BUTTONS: Final[list[ShellyButtonDescription[Any]]] = [
name="Reboot",
device_class=ButtonDeviceClass.RESTART,
entity_category=EntityCategory.CONFIG,
press_action=lambda coordinator: coordinator.device.trigger_reboot(),
press_action="trigger_reboot",
),
ShellyButtonDescription[ShellyBlockCoordinator](
key="self_test",
name="Self test",
translation_key="self_test",
entity_category=EntityCategory.DIAGNOSTIC,
press_action=lambda coordinator: coordinator.device.trigger_shelly_gas_self_test(),
press_action="trigger_shelly_gas_self_test",
supported=lambda coordinator: coordinator.device.model in SHELLY_GAS_MODELS,
),
ShellyButtonDescription[ShellyBlockCoordinator](
@ -59,7 +65,7 @@ BUTTONS: Final[list[ShellyButtonDescription[Any]]] = [
name="Mute",
translation_key="mute",
entity_category=EntityCategory.CONFIG,
press_action=lambda coordinator: coordinator.device.trigger_shelly_gas_mute(),
press_action="trigger_shelly_gas_mute",
supported=lambda coordinator: coordinator.device.model in SHELLY_GAS_MODELS,
),
ShellyButtonDescription[ShellyBlockCoordinator](
@ -67,11 +73,22 @@ BUTTONS: Final[list[ShellyButtonDescription[Any]]] = [
name="Unmute",
translation_key="unmute",
entity_category=EntityCategory.CONFIG,
press_action=lambda coordinator: coordinator.device.trigger_shelly_gas_unmute(),
press_action="trigger_shelly_gas_unmute",
supported=lambda coordinator: coordinator.device.model in SHELLY_GAS_MODELS,
),
]
BLU_TRV_BUTTONS: Final[list[ShellyButtonDescription]] = [
ShellyButtonDescription[ShellyRpcCoordinator](
key="calibrate",
name="Calibrate",
translation_key="calibrate",
entity_category=EntityCategory.CONFIG,
press_action="trigger_blu_trv_calibration",
supported=lambda coordinator: coordinator.device.model == MODEL_BLU_GATEWAY,
),
]
@callback
def async_migrate_unique_ids(
@ -123,14 +140,28 @@ async def async_setup_entry(
hass, config_entry.entry_id, partial(async_migrate_unique_ids, coordinator)
)
async_add_entities(
entities: list[ShellyButton | ShellyBluTrvButton] = []
entities.extend(
ShellyButton(coordinator, button)
for button in BUTTONS
if button.supported(coordinator)
)
if blutrv_key_ids := get_rpc_key_ids(coordinator.device.status, BLU_TRV_IDENTIFIER):
if TYPE_CHECKING:
assert isinstance(coordinator, ShellyRpcCoordinator)
class ShellyButton(
entities.extend(
ShellyBluTrvButton(coordinator, button, id_)
for id_ in blutrv_key_ids
for button in BLU_TRV_BUTTONS
)
async_add_entities(entities)
class ShellyBaseButton(
CoordinatorEntity[ShellyRpcCoordinator | ShellyBlockCoordinator], ButtonEntity
):
"""Defines a Shelly base button."""
@ -148,14 +179,100 @@ class ShellyButton(
) -> None:
"""Initialize Shelly button."""
super().__init__(coordinator)
self.entity_description = description
async def async_press(self) -> None:
"""Triggers the Shelly button press service."""
try:
await self._press_method()
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.device.name,
"error": repr(err),
},
) 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.device.name,
"error": repr(err),
},
) from err
except InvalidAuthError:
await self.coordinator.async_shutdown_device_and_start_reauth()
async def _press_method(self) -> None:
"""Press method."""
raise NotImplementedError
class ShellyButton(ShellyBaseButton):
"""Defines a Shelly button."""
def __init__(
self,
coordinator: ShellyRpcCoordinator | ShellyBlockCoordinator,
description: ShellyButtonDescription[
ShellyRpcCoordinator | ShellyBlockCoordinator
],
) -> None:
"""Initialize Shelly button."""
super().__init__(coordinator, description)
self._attr_name = f"{coordinator.device.name} {description.name}"
self._attr_unique_id = f"{coordinator.mac}_{description.key}"
self._attr_device_info = DeviceInfo(
connections={(CONNECTION_NETWORK_MAC, coordinator.mac)}
)
async def async_press(self) -> None:
"""Triggers the Shelly button press service."""
await self.entity_description.press_action(self.coordinator)
async def _press_method(self) -> None:
"""Press method."""
method = getattr(self.coordinator.device, self.entity_description.press_action)
if TYPE_CHECKING:
assert method is not None
await method()
class ShellyBluTrvButton(ShellyBaseButton):
"""Represent a Shelly BLU TRV button."""
def __init__(
self,
coordinator: ShellyRpcCoordinator,
description: ShellyButtonDescription,
id_: int,
) -> None:
"""Initialize."""
super().__init__(coordinator, description)
ble_addr: str = coordinator.device.config[f"{BLU_TRV_IDENTIFIER}:{id_}"]["addr"]
device_name = (
coordinator.device.config[f"{BLU_TRV_IDENTIFIER}:{id_}"]["name"]
or f"shellyblutrv-{ble_addr.replace(':', '')}"
)
self._attr_name = f"{device_name} {description.name}"
self._attr_unique_id = f"{ble_addr}_{description.key}"
self._attr_device_info = DeviceInfo(
connections={(CONNECTION_BLUETOOTH, ble_addr)}
)
self._id = id_
async def _press_method(self) -> None:
"""Press method."""
method = getattr(self.coordinator.device, self.entity_description.press_action)
if TYPE_CHECKING:
assert method is not None
await method(self._id)

View File

@ -7,7 +7,12 @@ from dataclasses import asdict, dataclass
from typing import Any, cast
from aioshelly.block_device import Block
from aioshelly.const import BLU_TRV_IDENTIFIER, BLU_TRV_MODEL_NAME, RPC_GENERATIONS
from aioshelly.const import (
BLU_TRV_IDENTIFIER,
BLU_TRV_MODEL_NAME,
BLU_TRV_TIMEOUT,
RPC_GENERATIONS,
)
from aioshelly.exceptions import DeviceConnectionError, InvalidAuthError
from homeassistant.components.climate import (
@ -36,7 +41,6 @@ from homeassistant.util.unit_system import US_CUSTOMARY_SYSTEM
from .const import (
BLU_TRV_TEMPERATURE_SETTINGS,
BLU_TRV_TIMEOUT,
DOMAIN,
LOGGER,
NOT_CALIBRATED_ISSUE_ID,

View File

@ -271,9 +271,6 @@ API_WS_URL = "/api/shelly/ws"
COMPONENT_ID_PATTERN = re.compile(r"[a-z\d]+:\d+")
# value confirmed by Shelly team
BLU_TRV_TIMEOUT = 60
ROLE_TO_DEVICE_CLASS_MAP = {
"current_humidity": SensorDeviceClass.HUMIDITY,
"current_temperature": SensorDeviceClass.TEMPERATURE,

View File

@ -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 RPC_GENERATIONS
from aioshelly.const import BLU_TRV_TIMEOUT, RPC_GENERATIONS
from aioshelly.exceptions import DeviceConnectionError, InvalidAuthError
from homeassistant.components.number import (
@ -25,7 +25,7 @@ from homeassistant.helpers.device_registry import CONNECTION_BLUETOOTH, DeviceIn
from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
from homeassistant.helpers.entity_registry import RegistryEntry
from .const import BLU_TRV_TIMEOUT, CONF_SLEEP_PERIOD, LOGGER, VIRTUAL_NUMBER_MODE_MAP
from .const import CONF_SLEEP_PERIOD, LOGGER, VIRTUAL_NUMBER_MODE_MAP
from .coordinator import ShellyBlockCoordinator, ShellyConfigEntry, ShellyRpcCoordinator
from .entity import (
BlockEntityDescription,

View File

@ -203,6 +203,14 @@
}
}
},
"exceptions": {
"device_communication_action_error": {
"message": "Device communication error occurred while calling the entity {entity} action for {device} device: {error}"
},
"rpc_call_action_error": {
"message": "RPC call error occurred while calling the entity {entity} action for {device} device: {error}"
}
},
"issues": {
"device_not_calibrated": {
"title": "Shelly device {device_name} is not calibrated",

View File

@ -0,0 +1,96 @@
# serializer version: 1
# name: test_rpc_blu_trv_button[button.trv_name_calibrate-entry]
EntityRegistryEntrySnapshot({
'aliases': set({
}),
'area_id': None,
'capabilities': None,
'config_entry_id': <ANY>,
'config_subentry_id': <ANY>,
'device_class': None,
'device_id': <ANY>,
'disabled_by': None,
'domain': 'button',
'entity_category': <EntityCategory.CONFIG: 'config'>,
'entity_id': 'button.trv_name_calibrate',
'has_entity_name': False,
'hidden_by': None,
'icon': None,
'id': <ANY>,
'labels': set({
}),
'name': None,
'options': dict({
}),
'original_device_class': None,
'original_icon': None,
'original_name': 'TRV-Name Calibrate',
'platform': 'shelly',
'previous_unique_id': None,
'supported_features': 0,
'translation_key': 'calibrate',
'unique_id': 'f8:44:77:25:f0:dd_calibrate',
'unit_of_measurement': None,
})
# ---
# name: test_rpc_blu_trv_button[button.trv_name_calibrate-state]
StateSnapshot({
'attributes': ReadOnlyDict({
'friendly_name': 'TRV-Name Calibrate',
}),
'context': <ANY>,
'entity_id': 'button.trv_name_calibrate',
'last_changed': <ANY>,
'last_reported': <ANY>,
'last_updated': <ANY>,
'state': 'unknown',
})
# ---
# name: test_rpc_button[button.test_name_reboot-entry]
EntityRegistryEntrySnapshot({
'aliases': set({
}),
'area_id': None,
'capabilities': None,
'config_entry_id': <ANY>,
'config_subentry_id': <ANY>,
'device_class': None,
'device_id': <ANY>,
'disabled_by': None,
'domain': 'button',
'entity_category': <EntityCategory.CONFIG: 'config'>,
'entity_id': 'button.test_name_reboot',
'has_entity_name': False,
'hidden_by': None,
'icon': None,
'id': <ANY>,
'labels': set({
}),
'name': None,
'options': dict({
}),
'original_device_class': <ButtonDeviceClass.RESTART: 'restart'>,
'original_icon': None,
'original_name': 'Test name Reboot',
'platform': 'shelly',
'previous_unique_id': None,
'supported_features': 0,
'translation_key': None,
'unique_id': '123456789ABC_reboot',
'unit_of_measurement': None,
})
# ---
# name: test_rpc_button[button.test_name_reboot-state]
StateSnapshot({
'attributes': ReadOnlyDict({
'device_class': 'restart',
'friendly_name': 'Test name Reboot',
}),
'context': <ANY>,
'entity_id': 'button.test_name_reboot',
'last_changed': <ANY>,
'last_reported': <ANY>,
'last_updated': <ANY>,
'state': 'unknown',
})
# ---

View File

@ -2,12 +2,17 @@
from unittest.mock import Mock
from aioshelly.const import MODEL_BLU_GATEWAY_G3
from aioshelly.exceptions import DeviceConnectionError, InvalidAuthError, RpcCallError
import pytest
from syrupy import SnapshotAssertion
from homeassistant.components.button import DOMAIN as BUTTON_DOMAIN, SERVICE_PRESS
from homeassistant.components.shelly.const import DOMAIN
from homeassistant.config_entries import SOURCE_REAUTH, ConfigEntryState
from homeassistant.const import ATTR_ENTITY_ID, STATE_UNKNOWN
from homeassistant.core import HomeAssistant
from homeassistant.exceptions import HomeAssistantError
from homeassistant.helpers.entity_registry import EntityRegistry
from . import init_integration
@ -38,7 +43,10 @@ async def test_block_button(
async def test_rpc_button(
hass: HomeAssistant, mock_rpc_device: Mock, entity_registry: EntityRegistry
hass: HomeAssistant,
mock_rpc_device: Mock,
entity_registry: EntityRegistry,
snapshot: SnapshotAssertion,
) -> None:
"""Test rpc device OTA button."""
await init_integration(hass, 2)
@ -46,11 +54,11 @@ async def test_rpc_button(
entity_id = "button.test_name_reboot"
# reboot button
assert hass.states.get(entity_id).state == STATE_UNKNOWN
state = hass.states.get(entity_id)
assert state == snapshot(name=f"{entity_id}-state")
entry = entity_registry.async_get(entity_id)
assert entry
assert entry.unique_id == "123456789ABC_reboot"
assert entry == snapshot(name=f"{entity_id}-entry")
await hass.services.async_call(
BUTTON_DOMAIN,
@ -61,6 +69,68 @@ async def test_rpc_button(
assert mock_rpc_device.trigger_reboot.call_count == 1
@pytest.mark.parametrize(
("exception", "error"),
[
(
DeviceConnectionError,
"Device communication error occurred while calling the entity button.test_name_reboot action for Test name device",
),
(
RpcCallError(999),
"RPC call error occurred while calling the entity button.test_name_reboot action for Test name device",
),
],
)
async def test_rpc_button_exc(
hass: HomeAssistant,
mock_rpc_device: Mock,
exception: Exception,
error: str,
) -> None:
"""Test RPC button with exception."""
await init_integration(hass, 2)
mock_rpc_device.trigger_reboot.side_effect = exception
with pytest.raises(HomeAssistantError, match=error):
await hass.services.async_call(
BUTTON_DOMAIN,
SERVICE_PRESS,
{ATTR_ENTITY_ID: "button.test_name_reboot"},
blocking=True,
)
async def test_rpc_button_reauth_error(
hass: HomeAssistant, mock_rpc_device: Mock
) -> None:
"""Test rpc device OTA button with authentication error."""
entry = await init_integration(hass, 2)
mock_rpc_device.trigger_reboot.side_effect = InvalidAuthError
await hass.services.async_call(
BUTTON_DOMAIN,
SERVICE_PRESS,
{ATTR_ENTITY_ID: "button.test_name_reboot"},
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
@pytest.mark.parametrize(
("gen", "old_unique_id", "new_unique_id", "migration"),
[
@ -104,3 +174,107 @@ async def test_migrate_unique_id(
bool("Migrating unique_id for button.test_name_reboot" in caplog.text)
== migration
)
async def test_rpc_blu_trv_button(
hass: HomeAssistant,
mock_blu_trv: Mock,
entity_registry: EntityRegistry,
monkeypatch: pytest.MonkeyPatch,
snapshot: SnapshotAssertion,
) -> None:
"""Test RPC BLU TRV button."""
monkeypatch.delitem(mock_blu_trv.status, "script:1")
monkeypatch.delitem(mock_blu_trv.status, "script:2")
monkeypatch.delitem(mock_blu_trv.status, "script:3")
await init_integration(hass, 3, model=MODEL_BLU_GATEWAY_G3)
entity_id = "button.trv_name_calibrate"
state = hass.states.get(entity_id)
assert state == snapshot(name=f"{entity_id}-state")
entry = entity_registry.async_get(entity_id)
assert entry == snapshot(name=f"{entity_id}-entry")
await hass.services.async_call(
BUTTON_DOMAIN,
SERVICE_PRESS,
{ATTR_ENTITY_ID: entity_id},
blocking=True,
)
assert mock_blu_trv.trigger_blu_trv_calibration.call_count == 1
@pytest.mark.parametrize(
("exception", "error"),
[
(
DeviceConnectionError,
"Device communication error occurred while calling the entity button.trv_name_calibrate action for Test name device",
),
(
RpcCallError(999),
"RPC call error occurred while calling the entity button.trv_name_calibrate action for Test name device",
),
],
)
async def test_rpc_blu_trv_button_exc(
hass: HomeAssistant,
mock_blu_trv: Mock,
monkeypatch: pytest.MonkeyPatch,
exception: Exception,
error: str,
) -> None:
"""Test RPC BLU TRV button with exception."""
monkeypatch.delitem(mock_blu_trv.status, "script:1")
monkeypatch.delitem(mock_blu_trv.status, "script:2")
monkeypatch.delitem(mock_blu_trv.status, "script:3")
await init_integration(hass, 3, model=MODEL_BLU_GATEWAY_G3)
mock_blu_trv.trigger_blu_trv_calibration.side_effect = exception
with pytest.raises(HomeAssistantError, match=error):
await hass.services.async_call(
BUTTON_DOMAIN,
SERVICE_PRESS,
{ATTR_ENTITY_ID: "button.trv_name_calibrate"},
blocking=True,
)
async def test_rpc_blu_trv_button_auth_error(
hass: HomeAssistant,
mock_blu_trv: Mock,
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""Test RPC BLU TRV button with authentication error."""
monkeypatch.delitem(mock_blu_trv.status, "script:1")
monkeypatch.delitem(mock_blu_trv.status, "script:2")
monkeypatch.delitem(mock_blu_trv.status, "script:3")
entry = await init_integration(hass, 3, model=MODEL_BLU_GATEWAY_G3)
mock_blu_trv.trigger_blu_trv_calibration.side_effect = InvalidAuthError
await hass.services.async_call(
BUTTON_DOMAIN,
SERVICE_PRESS,
{ATTR_ENTITY_ID: "button.trv_name_calibrate"},
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

View File

@ -5,6 +5,7 @@ from unittest.mock import AsyncMock, Mock, PropertyMock
from aioshelly.const import (
BLU_TRV_IDENTIFIER,
BLU_TRV_TIMEOUT,
MODEL_BLU_GATEWAY_G3,
MODEL_VALVE,
MODEL_WALL_DISPLAY,
@ -27,7 +28,7 @@ from homeassistant.components.climate import (
HVACAction,
HVACMode,
)
from homeassistant.components.shelly.const import BLU_TRV_TIMEOUT, DOMAIN
from homeassistant.components.shelly.const import DOMAIN
from homeassistant.components.switch import DOMAIN as SWITCH_DOMAIN
from homeassistant.config_entries import SOURCE_REAUTH, ConfigEntryState
from homeassistant.const import (

View File

@ -3,7 +3,7 @@
from copy import deepcopy
from unittest.mock import AsyncMock, Mock
from aioshelly.const import MODEL_BLU_GATEWAY_G3
from aioshelly.const import BLU_TRV_TIMEOUT, MODEL_BLU_GATEWAY_G3
from aioshelly.exceptions import DeviceConnectionError, InvalidAuthError
import pytest
from syrupy import SnapshotAssertion
@ -18,7 +18,7 @@ from homeassistant.components.number import (
SERVICE_SET_VALUE,
NumberMode,
)
from homeassistant.components.shelly.const import BLU_TRV_TIMEOUT, DOMAIN
from homeassistant.components.shelly.const import DOMAIN
from homeassistant.config_entries import SOURCE_REAUTH, ConfigEntryState
from homeassistant.const import ATTR_ENTITY_ID, ATTR_UNIT_OF_MEASUREMENT, STATE_UNKNOWN
from homeassistant.core import HomeAssistant, State