Fix Shelly Gen1 entity description restore

homeassistant/components/shelly/binary_sensor.py

@@ -15,7 +15,6 @@
 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 @@ class RestBinarySensorDescription(RestEntityDescription, BinarySensorEntityDescr
 }
 
 
-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(
- sensor_class(coordinator, None, attribute, description, entry, sensors)
- )
+ if description := sensors.get((block_type, attribute)):
+ entities.append(
+ sensor_class(coordinator, None, attribute, description, entry)
+ )
 
  if not entities:
  return
@@ -444,7 +442,7 @@ def available(self) -> bool:
  """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 @@ def __init__(
  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 @@ def __init__(
  @callback
  def _update_callback(self) -> None:
  """Handle device update."""
- if (
- self.block is not None
- or not self.coordinator.device.initialized
- or self.sensors is None
- ):
+ if self.block is not None or not self.coordinator.device.initialized:
  super()._update_callback()
  return
 
@@ -607,13 +599,7 @@ def _update_callback(self) -> None:
  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 @@ class BlockNumberDescription(BlockEntityDescription, NumberEntityDescription):
 }
 
 
-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 @@ def __init__(
  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.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 @@ class RestSensorDescription(RestEntityDescription, SensorEntityDescription):
 }
 
 
-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 @@ def __init__(
  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 @@
  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()