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>
This commit is contained in:
Shay Levy 2024-01-15 20:53:56 +02:00 committed by Franck Nijhof
parent 9c6f87dd11
commit ed31adc6db
No known key found for this signature in database
GPG key ID: D62583BA8AB11CA3
5 changed files with 24 additions and 74 deletions

View file

@ -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,

View file

@ -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

View file

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

View file

@ -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:

View file

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