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
This commit is contained in:
Maciej Bieniek 2023-06-05 00:14:08 +00:00 committed by GitHub
parent 7c02e1ca99
commit fe61672792
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
7 changed files with 106 additions and 39 deletions

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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