Don't inherit SensorEntity/NumberEntity and RestoreEntity in Shelly integration (#93531)

* Use RestoreNumber and Restore Sensor for block entities

* Update tests

* Use RestoreSensor for RPC entities

* Fix test for number platform

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."""

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

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."""

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:

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."""

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)

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)