diff --git a/homeassistant/components/shelly/binary_sensor.py b/homeassistant/components/shelly/binary_sensor.py index 449fc142218..1474906cacb 100644 --- a/homeassistant/components/shelly/binary_sensor.py +++ b/homeassistant/components/shelly/binary_sensor.py @@ -14,6 +14,7 @@ from homeassistant.const import STATE_ON, EntityCategory from homeassistant.core import HomeAssistant from homeassistant.helpers.entity_platform import AddEntitiesCallback from homeassistant.helpers.entity_registry import RegistryEntry +from homeassistant.helpers.restore_state import RestoreEntity from .const import CONF_SLEEP_PERIOD from .entity import ( @@ -290,11 +291,18 @@ class RpcBinarySensor(ShellyRpcAttributeEntity, BinarySensorEntity): return bool(self.attribute_value) -class BlockSleepingBinarySensor(ShellySleepingBlockAttributeEntity, BinarySensorEntity): +class BlockSleepingBinarySensor( + ShellySleepingBlockAttributeEntity, BinarySensorEntity, RestoreEntity +): """Represent a block sleeping binary sensor.""" entity_description: BlockBinarySensorDescription + async def async_added_to_hass(self) -> None: + """Handle entity which will be added.""" + await super().async_added_to_hass() + self.last_state = await self.async_get_last_state() + @property def is_on(self) -> bool | None: """Return true if sensor state is on.""" @@ -307,11 +315,18 @@ class BlockSleepingBinarySensor(ShellySleepingBlockAttributeEntity, BinarySensor return self.last_state.state == STATE_ON -class RpcSleepingBinarySensor(ShellySleepingRpcAttributeEntity, BinarySensorEntity): +class RpcSleepingBinarySensor( + ShellySleepingRpcAttributeEntity, BinarySensorEntity, RestoreEntity +): """Represent a RPC sleeping binary sensor entity.""" entity_description: RpcBinarySensorDescription + async def async_added_to_hass(self) -> None: + """Handle entity which will be added.""" + await super().async_added_to_hass() + self.last_state = await self.async_get_last_state() + @property def is_on(self) -> bool | None: """Return true if RPC sensor state is on.""" diff --git a/homeassistant/components/shelly/entity.py b/homeassistant/components/shelly/entity.py index 1f18a5f8e18..50d41899800 100644 --- a/homeassistant/components/shelly/entity.py +++ b/homeassistant/components/shelly/entity.py @@ -19,7 +19,6 @@ from homeassistant.helpers.entity_registry import ( async_entries_for_config_entry, async_get as er_async_get, ) -from homeassistant.helpers.restore_state import RestoreEntity from homeassistant.helpers.typing import StateType from homeassistant.helpers.update_coordinator import CoordinatorEntity @@ -552,7 +551,7 @@ class ShellyRpcAttributeEntity(ShellyRpcEntity, Entity): return self.entity_description.available(self.sub_status) -class ShellySleepingBlockAttributeEntity(ShellyBlockAttributeEntity, RestoreEntity): +class ShellySleepingBlockAttributeEntity(ShellyBlockAttributeEntity): """Represent a shelly sleeping block attribute entity.""" # pylint: disable=super-init-not-called @@ -589,11 +588,6 @@ class ShellySleepingBlockAttributeEntity(ShellyBlockAttributeEntity, RestoreEnti self._attr_unique_id = entry.unique_id self._attr_name = cast(str, entry.original_name) - async def async_added_to_hass(self) -> None: - """Handle entity which will be added.""" - await super().async_added_to_hass() - self.last_state = await self.async_get_last_state() - @callback def _update_callback(self) -> None: """Handle device update.""" @@ -629,7 +623,7 @@ class ShellySleepingBlockAttributeEntity(ShellyBlockAttributeEntity, RestoreEnti return -class ShellySleepingRpcAttributeEntity(ShellyRpcAttributeEntity, RestoreEntity): +class ShellySleepingRpcAttributeEntity(ShellyRpcAttributeEntity): """Helper class to represent a sleeping rpc attribute.""" entity_description: RpcEntityDescription @@ -665,8 +659,3 @@ class ShellySleepingRpcAttributeEntity(ShellyRpcAttributeEntity, RestoreEntity): ) elif entry is not None: self._attr_name = cast(str, entry.original_name) - - async def async_added_to_hass(self) -> None: - """Handle entity which will be added.""" - await super().async_added_to_hass() - self.last_state = await self.async_get_last_state() diff --git a/homeassistant/components/shelly/number.py b/homeassistant/components/shelly/number.py index a89c74f9e50..f2b6bedb443 100644 --- a/homeassistant/components/shelly/number.py +++ b/homeassistant/components/shelly/number.py @@ -1,15 +1,18 @@ """Number for Shelly.""" from __future__ import annotations +from collections.abc import Mapping from dataclasses import dataclass from typing import Any, Final, cast +from aioshelly.block_device import Block from aioshelly.exceptions import DeviceConnectionError, InvalidAuthError from homeassistant.components.number import ( - NumberEntity, NumberEntityDescription, + NumberExtraStoredData, NumberMode, + RestoreNumber, ) from homeassistant.config_entries import ConfigEntry from homeassistant.const import PERCENTAGE, EntityCategory @@ -19,6 +22,7 @@ from homeassistant.helpers.entity_platform import AddEntitiesCallback from homeassistant.helpers.entity_registry import RegistryEntry from .const import CONF_SLEEP_PERIOD, LOGGER +from .coordinator import ShellyBlockCoordinator from .entity import ( BlockEntityDescription, ShellySleepingBlockAttributeEntity, @@ -85,22 +89,39 @@ async def async_setup_entry( ) -# pylint: disable-next=hass-invalid-inheritance # needs fixing -class BlockSleepingNumber(ShellySleepingBlockAttributeEntity, NumberEntity): +class BlockSleepingNumber(ShellySleepingBlockAttributeEntity, RestoreNumber): """Represent a block sleeping number.""" entity_description: BlockNumberDescription + def __init__( + self, + coordinator: ShellyBlockCoordinator, + block: Block | None, + attribute: str, + description: BlockNumberDescription, + entry: RegistryEntry | None = None, + sensors: Mapping[tuple[str, str], BlockNumberDescription] | None = None, + ) -> None: + """Initialize the sleeping sensor.""" + self.restored_data: NumberExtraStoredData | None = None + super().__init__(coordinator, block, attribute, description, entry, sensors) + + async def async_added_to_hass(self) -> None: + """Handle entity which will be added.""" + await super().async_added_to_hass() + self.restored_data = await self.async_get_last_number_data() + @property def native_value(self) -> float | None: """Return value of number.""" if self.block is not None: return cast(float, self.attribute_value) - if self.last_state is None: + if self.restored_data is None: return None - return cast(float, self.last_state.state) + return cast(float, self.restored_data.native_value) async def async_set_native_value(self, value: float) -> None: """Set value.""" diff --git a/homeassistant/components/shelly/sensor.py b/homeassistant/components/shelly/sensor.py index 4a88157efc6..0260a540f0c 100644 --- a/homeassistant/components/shelly/sensor.py +++ b/homeassistant/components/shelly/sensor.py @@ -8,14 +8,15 @@ from typing import Final, cast from aioshelly.block_device import Block from homeassistant.components.sensor import ( + RestoreSensor, SensorDeviceClass, SensorEntity, SensorEntityDescription, + SensorExtraStoredData, SensorStateClass, ) from homeassistant.config_entries import ConfigEntry from homeassistant.const import ( - ATTR_UNIT_OF_MEASUREMENT, CONCENTRATION_PARTS_PER_MILLION, DEGREE, LIGHT_LUX, @@ -35,7 +36,7 @@ from homeassistant.helpers.entity_registry import RegistryEntry from homeassistant.helpers.typing import StateType from .const import CONF_SLEEP_PERIOD, SHAIR_MAX_WORK_HOURS -from .coordinator import ShellyBlockCoordinator +from .coordinator import ShellyBlockCoordinator, ShellyRpcCoordinator from .entity import ( BlockEntityDescription, RestEntityDescription, @@ -776,8 +777,7 @@ class RpcSensor(ShellyRpcAttributeEntity, SensorEntity): return self.attribute_value -# pylint: disable-next=hass-invalid-inheritance # needs fixing -class BlockSleepingSensor(ShellySleepingBlockAttributeEntity, SensorEntity): +class BlockSleepingSensor(ShellySleepingBlockAttributeEntity, RestoreSensor): """Represent a block sleeping sensor.""" entity_description: BlockSensorDescription @@ -793,6 +793,12 @@ class BlockSleepingSensor(ShellySleepingBlockAttributeEntity, SensorEntity): ) -> None: """Initialize the sleeping sensor.""" super().__init__(coordinator, block, attribute, description, entry, sensors) + self.restored_data: SensorExtraStoredData | None = None + + async def async_added_to_hass(self) -> None: + """Handle entity which will be added.""" + await super().async_added_to_hass() + self.restored_data = await self.async_get_last_sensor_data() @property def native_value(self) -> StateType: @@ -800,10 +806,10 @@ class BlockSleepingSensor(ShellySleepingBlockAttributeEntity, SensorEntity): if self.block is not None: return self.attribute_value - if self.last_state is None: + if self.restored_data is None: return None - return self.last_state.state + return cast(StateType, self.restored_data.native_value) @property def native_unit_of_measurement(self) -> str | None: @@ -811,28 +817,44 @@ class BlockSleepingSensor(ShellySleepingBlockAttributeEntity, SensorEntity): if self.block is not None: return self.entity_description.native_unit_of_measurement - if self.last_state is None: + if self.restored_data is None: return None - return self.last_state.attributes.get(ATTR_UNIT_OF_MEASUREMENT) + return self.restored_data.native_unit_of_measurement -# pylint: disable-next=hass-invalid-inheritance # needs fixing -class RpcSleepingSensor(ShellySleepingRpcAttributeEntity, SensorEntity): +class RpcSleepingSensor(ShellySleepingRpcAttributeEntity, RestoreSensor): """Represent a RPC sleeping sensor.""" entity_description: RpcSensorDescription + def __init__( + self, + coordinator: ShellyRpcCoordinator, + key: str, + attribute: str, + description: RpcEntityDescription, + entry: RegistryEntry | None = None, + ) -> None: + """Initialize the sleeping sensor.""" + super().__init__(coordinator, key, attribute, description, entry) + self.restored_data: SensorExtraStoredData | None = None + + async def async_added_to_hass(self) -> None: + """Handle entity which will be added.""" + await super().async_added_to_hass() + self.restored_data = await self.async_get_last_sensor_data() + @property def native_value(self) -> StateType: """Return value of sensor.""" if self.coordinator.device.initialized: return self.attribute_value - if self.last_state is None: + if self.restored_data is None: return None - return self.last_state.state + return cast(StateType, self.restored_data.native_value) @property def native_unit_of_measurement(self) -> str | None: diff --git a/homeassistant/components/shelly/update.py b/homeassistant/components/shelly/update.py index 47a39105ac4..3b2096f0c1a 100644 --- a/homeassistant/components/shelly/update.py +++ b/homeassistant/components/shelly/update.py @@ -21,6 +21,7 @@ from homeassistant.const import EntityCategory from homeassistant.core import HomeAssistant from homeassistant.exceptions import HomeAssistantError from homeassistant.helpers.entity_platform import AddEntitiesCallback +from homeassistant.helpers.restore_state import RestoreEntity from .const import CONF_SLEEP_PERIOD from .coordinator import ShellyBlockCoordinator, ShellyRpcCoordinator @@ -282,11 +283,18 @@ class RpcUpdateEntity(ShellyRpcAttributeEntity, UpdateEntity): LOGGER.debug("OTA update call successful") -class RpcSleepingUpdateEntity(ShellySleepingRpcAttributeEntity, UpdateEntity): +class RpcSleepingUpdateEntity( + ShellySleepingRpcAttributeEntity, UpdateEntity, RestoreEntity +): """Represent a RPC sleeping update entity.""" entity_description: RpcUpdateDescription + async def async_added_to_hass(self) -> None: + """Handle entity which will be added.""" + await super().async_added_to_hass() + self.last_state = await self.async_get_last_state() + @property def installed_version(self) -> str | None: """Version currently in use.""" diff --git a/tests/components/shelly/test_number.py b/tests/components/shelly/test_number.py index 57a8e801b92..403d2f2993d 100644 --- a/tests/components/shelly/test_number.py +++ b/tests/components/shelly/test_number.py @@ -17,7 +17,7 @@ from homeassistant.exceptions import HomeAssistantError from . import init_integration, register_device, register_entity -from tests.common import mock_restore_cache +from tests.common import mock_restore_cache_with_extra_data DEVICE_BLOCK_ID = 4 @@ -62,7 +62,14 @@ async def test_block_restored_number( entry, capabilities, ) - mock_restore_cache(hass, [State(entity_id, "40")]) + extra_data = { + "native_max_value": 100, + "native_min_value": 0, + "native_step": 1, + "native_unit_of_measurement": "%", + "native_value": "40", + } + mock_restore_cache_with_extra_data(hass, ((State(entity_id, ""), extra_data),)) monkeypatch.setattr(mock_block_device, "initialized", False) await hass.config_entries.async_setup(entry.entry_id) diff --git a/tests/components/shelly/test_sensor.py b/tests/components/shelly/test_sensor.py index 0b906d60079..d87460fb17d 100644 --- a/tests/components/shelly/test_sensor.py +++ b/tests/components/shelly/test_sensor.py @@ -20,7 +20,7 @@ from . import ( register_entity, ) -from tests.common import mock_restore_cache +from tests.common import mock_restore_cache_with_extra_data RELAY_BLOCK_ID = 0 SENSOR_BLOCK_ID = 3 @@ -137,7 +137,9 @@ async def test_block_restored_sleeping_sensor( entity_id = register_entity( hass, SENSOR_DOMAIN, "test_name_temperature", "sensor_0-temp", entry ) - mock_restore_cache(hass, [State(entity_id, "20.4")]) + extra_data = {"native_value": "20.4", "native_unit_of_measurement": "°C"} + + mock_restore_cache_with_extra_data(hass, ((State(entity_id, ""), extra_data),)) monkeypatch.setattr(mock_block_device, "initialized", False) await hass.config_entries.async_setup(entry.entry_id) await hass.async_block_till_done() @@ -216,7 +218,9 @@ async def test_block_not_matched_restored_sleeping_sensor( entity_id = register_entity( hass, SENSOR_DOMAIN, "test_name_temperature", "sensor_0-temp", entry ) - mock_restore_cache(hass, [State(entity_id, "20.4")]) + extra_data = {"native_value": "20.4", "native_unit_of_measurement": "°C"} + + mock_restore_cache_with_extra_data(hass, ((State(entity_id, ""), extra_data),)) monkeypatch.setattr(mock_block_device, "initialized", False) await hass.config_entries.async_setup(entry.entry_id) await hass.async_block_till_done() @@ -357,8 +361,9 @@ async def test_rpc_restored_sleeping_sensor( "temperature:0-temperature_0", entry, ) + extra_data = {"native_value": "21.0", "native_unit_of_measurement": "°C"} - mock_restore_cache(hass, [State(entity_id, "21.0")]) + mock_restore_cache_with_extra_data(hass, ((State(entity_id, ""), extra_data),)) monkeypatch.setattr(mock_rpc_device, "initialized", False) await hass.config_entries.async_setup(entry.entry_id)