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

homeassistant/components/shelly/button.py

@@ -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)

homeassistant/components/shelly/climate.py

@@ -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,

homeassistant/components/shelly/const.py

@@ -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,

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 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,

homeassistant/components/shelly/strings.json

@@ -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",

tests/components/shelly/snapshots/test_button.ambr

@@ -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',
+  })
+# ---

tests/components/shelly/test_button.py

@@ -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

tests/components/shelly/test_climate.py

@@ -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 (

tests/components/shelly/test_number.py

@@ -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