Fix Shelly Gen1 entity description restore (#108052)

* Fix Shelly Gen1 entity description restore

* Update tests/components/shelly/test_sensor.py

Co-authored-by: J. Nick Koston <nick@koston.org>

---------

Co-authored-by: J. Nick Koston <nick@koston.org>

homeassistant/components/shelly/binary_sensor.py

@@ -15,7 +15,6 @@ from homeassistant.config_entries import ConfigEntry
 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
@@ -210,16 +209,6 @@ RPC_SENSORS: Final = {
 }
 
 
-def _build_block_description(entry: RegistryEntry) -> BlockBinarySensorDescription:
-    """Build description when restoring block attribute entities."""
-    return BlockBinarySensorDescription(
-        key="",
-        name="",
-        icon=entry.original_icon,
-        device_class=entry.original_device_class,
-    )
-
-
 async def async_setup_entry(
     hass: HomeAssistant,
     config_entry: ConfigEntry,
@@ -248,7 +237,6 @@ async def async_setup_entry(
             async_add_entities,
             SENSORS,
             BlockSleepingBinarySensor,
-            _build_block_description,
         )
     else:
         async_setup_entry_attribute_entities(
@@ -257,7 +245,6 @@ async def async_setup_entry(
             async_add_entities,
             SENSORS,
             BlockBinarySensor,
-            _build_block_description,
         )
         async_setup_entry_rest(
             hass,

homeassistant/components/shelly/entity.py

@@ -39,7 +39,6 @@ def async_setup_entry_attribute_entities(
     async_add_entities: AddEntitiesCallback,
     sensors: Mapping[tuple[str, str], BlockEntityDescription],
     sensor_class: Callable,
-    description_class: Callable[[RegistryEntry], BlockEntityDescription],
 ) -> None:
     """Set up entities for attributes."""
     coordinator = get_entry_data(hass)[config_entry.entry_id].block
@@ -56,7 +55,6 @@ def async_setup_entry_attribute_entities(
             coordinator,
             sensors,
             sensor_class,
-            description_class,
         )
 
 
@@ -113,7 +111,6 @@ def async_restore_block_attribute_entities(
     coordinator: ShellyBlockCoordinator,
     sensors: Mapping[tuple[str, str], BlockEntityDescription],
     sensor_class: Callable,
-    description_class: Callable[[RegistryEntry], BlockEntityDescription],
 ) -> None:
     """Restore block attributes entities."""
     entities = []
@@ -128,11 +125,12 @@ def async_restore_block_attribute_entities(
             continue
 
         attribute = entry.unique_id.split("-")[-1]
-        description = description_class(entry)
+        block_type = entry.unique_id.split("-")[-2].split("_")[0]
 
-        entities.append(
+        if description := sensors.get((block_type, attribute)):
-            sensor_class(coordinator, None, attribute, description, entry, sensors)
+            entities.append(
-        )
+                sensor_class(coordinator, None, attribute, description, entry)
+            )
 
     if not entities:
         return
@@ -444,7 +442,7 @@ class ShellyBlockAttributeEntity(ShellyBlockEntity, Entity):
         """Available."""
         available = super().available
 
-        if not available or not self.entity_description.available:
+        if not available or not self.entity_description.available or self.block is None:
             return available
 
         return self.entity_description.available(self.block)
@@ -559,10 +557,8 @@ class ShellySleepingBlockAttributeEntity(ShellyBlockAttributeEntity):
         attribute: str,
         description: BlockEntityDescription,
         entry: RegistryEntry | None = None,
-        sensors: Mapping[tuple[str, str], BlockEntityDescription] | None = None,
     ) -> None:
         """Initialize the sleeping sensor."""
-        self.sensors = sensors
         self.last_state: State | None = None
         self.coordinator = coordinator
         self.attribute = attribute
@@ -587,11 +583,7 @@ class ShellySleepingBlockAttributeEntity(ShellyBlockAttributeEntity):
     @callback
     def _update_callback(self) -> None:
         """Handle device update."""
-        if (
+        if self.block is not None or not self.coordinator.device.initialized:
-            self.block is not None
-            or not self.coordinator.device.initialized
-            or self.sensors is None
-        ):
             super()._update_callback()
             return
 
@@ -607,13 +599,7 @@ class ShellySleepingBlockAttributeEntity(ShellyBlockAttributeEntity):
                 if sensor_id != entity_sensor:
                     continue
 
-                description = self.sensors.get((block.type, sensor_id))
-                if description is None:
-                    continue
-
                 self.block = block
-                self.entity_description = description
-
                 LOGGER.debug("Entity %s attached to block", self.name)
                 super()._update_callback()
                 return

homeassistant/components/shelly/number.py

@@ -1,7 +1,6 @@
 """Number for Shelly."""
 from __future__ import annotations
 
-from collections.abc import Mapping
 from dataclasses import dataclass
 from typing import Any, Final, cast
 
@@ -56,22 +55,6 @@ NUMBERS: Final = {
 }
 
 
-def _build_block_description(entry: RegistryEntry) -> BlockNumberDescription:
-    """Build description when restoring block attribute entities."""
-    assert entry.capabilities
-    return BlockNumberDescription(
-        key="",
-        name="",
-        icon=entry.original_icon,
-        native_unit_of_measurement=entry.unit_of_measurement,
-        device_class=entry.original_device_class,
-        native_min_value=cast(float, entry.capabilities.get("min")),
-        native_max_value=cast(float, entry.capabilities.get("max")),
-        native_step=cast(float, entry.capabilities.get("step")),
-        mode=cast(NumberMode, entry.capabilities.get("mode")),
-    )
-
-
 async def async_setup_entry(
     hass: HomeAssistant,
     config_entry: ConfigEntry,
@@ -85,7 +68,6 @@ async def async_setup_entry(
             async_add_entities,
             NUMBERS,
             BlockSleepingNumber,
-            _build_block_description,
         )
 
 
@@ -101,11 +83,10 @@ class BlockSleepingNumber(ShellySleepingBlockAttributeEntity, RestoreNumber):
         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)
+        super().__init__(coordinator, block, attribute, description, entry)
 
     async def async_added_to_hass(self) -> None:
         """Handle entity which will be added."""

homeassistant/components/shelly/sensor.py

@@ -1,7 +1,6 @@
 """Sensor for Shelly."""
 from __future__ import annotations
 
-from collections.abc import Mapping
 from dataclasses import dataclass
 from typing import Final, cast
 
@@ -36,7 +35,6 @@ from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.entity_registry import RegistryEntry
 from homeassistant.helpers.typing import StateType
-from homeassistant.util.enum import try_parse_enum
 
 from .const import CONF_SLEEP_PERIOD, SHAIR_MAX_WORK_HOURS
 from .coordinator import ShellyBlockCoordinator, ShellyRpcCoordinator
@@ -963,17 +961,6 @@ RPC_SENSORS: Final = {
 }
 
 
-def _build_block_description(entry: RegistryEntry) -> BlockSensorDescription:
-    """Build description when restoring block attribute entities."""
-    return BlockSensorDescription(
-        key="",
-        name="",
-        icon=entry.original_icon,
-        native_unit_of_measurement=entry.unit_of_measurement,
-        device_class=try_parse_enum(SensorDeviceClass, entry.original_device_class),
-    )
-
-
 async def async_setup_entry(
     hass: HomeAssistant,
     config_entry: ConfigEntry,
@@ -1002,7 +989,6 @@ async def async_setup_entry(
             async_add_entities,
             SENSORS,
             BlockSleepingSensor,
-            _build_block_description,
         )
     else:
         async_setup_entry_attribute_entities(
@@ -1011,7 +997,6 @@ async def async_setup_entry(
             async_add_entities,
             SENSORS,
             BlockSensor,
-            _build_block_description,
         )
         async_setup_entry_rest(
             hass, config_entry, async_add_entities, REST_SENSORS, RestSensor
@@ -1075,10 +1060,9 @@ class BlockSleepingSensor(ShellySleepingBlockAttributeEntity, RestoreSensor):
         attribute: str,
         description: BlockSensorDescription,
         entry: RegistryEntry | None = None,
-        sensors: Mapping[tuple[str, str], BlockSensorDescription] | None = None,
     ) -> None:
         """Initialize the sleeping sensor."""
-        super().__init__(coordinator, block, attribute, description, entry, sensors)
+        super().__init__(coordinator, block, attribute, description, entry)
         self.restored_data: SensorExtraStoredData | None = None
 
     async def async_added_to_hass(self) -> None:

tests/components/shelly/test_sensor.py

@@ -6,9 +6,15 @@ from homeassistant.components.homeassistant import (
     DOMAIN as HA_DOMAIN,
     SERVICE_UPDATE_ENTITY,
 )
-from homeassistant.components.sensor import DOMAIN as SENSOR_DOMAIN
+from homeassistant.components.sensor import (
+    ATTR_STATE_CLASS,
+    DOMAIN as SENSOR_DOMAIN,
+    SensorDeviceClass,
+    SensorStateClass,
+)
 from homeassistant.components.shelly.const import DOMAIN
 from homeassistant.const import (
+    ATTR_DEVICE_CLASS,
     ATTR_ENTITY_ID,
     ATTR_UNIT_OF_MEASUREMENT,
     PERCENTAGE,
@@ -153,7 +159,11 @@ async def test_block_restored_sleeping_sensor(
     await hass.config_entries.async_setup(entry.entry_id)
     await hass.async_block_till_done()
 
-    assert hass.states.get(entity_id).state == "20.4"
+    state = hass.states.get(entity_id)
+    assert state
+    assert state.state == "20.4"
+    assert state.attributes[ATTR_STATE_CLASS] == SensorStateClass.MEASUREMENT
+    assert state.attributes[ATTR_DEVICE_CLASS] == SensorDeviceClass.TEMPERATURE
 
     # Make device online
     monkeypatch.setattr(mock_block_device, "initialized", True)
@@ -237,7 +247,9 @@ async def test_block_not_matched_restored_sleeping_sensor(
     assert hass.states.get(entity_id).state == "20.4"
 
     # Make device online
-    monkeypatch.setattr(mock_block_device.blocks[SENSOR_BLOCK_ID], "type", "other_type")
+    monkeypatch.setattr(
+        mock_block_device.blocks[SENSOR_BLOCK_ID], "description", "other_desc"
+    )
     monkeypatch.setattr(mock_block_device, "initialized", True)
     mock_block_device.mock_update()
     await hass.async_block_till_done()