From 28e18ce7bff19484a8513b2a7ac5a9427d5335cf Mon Sep 17 00:00:00 2001 From: Shay Levy Date: Mon, 15 Jan 2024 20:53:56 +0200 Subject: [PATCH] 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 --------- Co-authored-by: J. Nick Koston --- .../components/shelly/binary_sensor.py | 13 --------- homeassistant/components/shelly/entity.py | 28 +++++-------------- homeassistant/components/shelly/number.py | 21 +------------- homeassistant/components/shelly/sensor.py | 18 +----------- tests/components/shelly/test_sensor.py | 18 ++++++++++-- 5 files changed, 24 insertions(+), 74 deletions(-) diff --git a/homeassistant/components/shelly/binary_sensor.py b/homeassistant/components/shelly/binary_sensor.py index b07747f298e..4ad51e5cc0f 100644 --- a/homeassistant/components/shelly/binary_sensor.py +++ b/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, diff --git a/homeassistant/components/shelly/entity.py b/homeassistant/components/shelly/entity.py index 796402c8bba..3132f1f571e 100644 --- a/homeassistant/components/shelly/entity.py +++ b/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 @@ 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 ( - 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 @@ 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 diff --git a/homeassistant/components/shelly/number.py b/homeassistant/components/shelly/number.py index 77d066a6106..5d35e71ce5d 100644 --- a/homeassistant/components/shelly/number.py +++ b/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.""" diff --git a/homeassistant/components/shelly/sensor.py b/homeassistant/components/shelly/sensor.py index c7d89f2d284..b439a19e318 100644 --- a/homeassistant/components/shelly/sensor.py +++ b/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: diff --git a/tests/components/shelly/test_sensor.py b/tests/components/shelly/test_sensor.py index 380f4f5999e..86c6356191b 100644 --- a/tests/components/shelly/test_sensor.py +++ b/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()