diff --git a/homeassistant/components/vesync/sensor.py b/homeassistant/components/vesync/sensor.py index 24ba6f2f0a0..6e0c09b2f60 100644 --- a/homeassistant/components/vesync/sensor.py +++ b/homeassistant/components/vesync/sensor.py @@ -1,8 +1,14 @@ """Support for power & energy sensors for VeSync outlets.""" +from __future__ import annotations + from collections.abc import Callable from dataclasses import dataclass import logging +from pyvesync.vesyncfan import VeSyncAirBypass +from pyvesync.vesyncoutlet import VeSyncOutlet +from pyvesync.vesyncswitch import VeSyncSwitch + from homeassistant.components.sensor import ( SensorDeviceClass, SensorEntity, @@ -22,7 +28,7 @@ from homeassistant.helpers.entity import EntityCategory from homeassistant.helpers.entity_platform import AddEntitiesCallback from homeassistant.helpers.typing import StateType -from .common import VeSyncBaseEntity, VeSyncDevice +from .common import VeSyncBaseEntity from .const import DEV_TYPE_TO_HA, DOMAIN, SKU_TO_BASE_DEVICE, VS_DISCOVERY, VS_SENSORS _LOGGER = logging.getLogger(__name__) @@ -32,7 +38,7 @@ _LOGGER = logging.getLogger(__name__) class VeSyncSensorEntityDescriptionMixin: """Mixin for required keys.""" - value_fn: Callable[[VeSyncDevice], StateType] + value_fn: Callable[[VeSyncAirBypass | VeSyncOutlet | VeSyncSwitch], StateType] @dataclass @@ -41,8 +47,12 @@ class VeSyncSensorEntityDescription( ): """Describe VeSync sensor entity.""" - exists_fn: Callable[[VeSyncDevice], bool] = lambda _: True - update_fn: Callable[[VeSyncDevice], None] = lambda _: None + exists_fn: Callable[ + [VeSyncAirBypass | VeSyncOutlet | VeSyncSwitch], bool + ] = lambda _: True + update_fn: Callable[ + [VeSyncAirBypass | VeSyncOutlet | VeSyncSwitch], None + ] = lambda _: None def update_energy(device): @@ -107,7 +117,7 @@ SENSORS: tuple[VeSyncSensorEntityDescription, ...] = ( device_class=SensorDeviceClass.ENERGY, native_unit_of_measurement=ENERGY_KILO_WATT_HOUR, state_class=SensorStateClass.TOTAL_INCREASING, - value_fn=lambda device: device.details["energy"], + value_fn=lambda device: device.energy_today, update_fn=update_energy, exists_fn=lambda device: ha_dev_type(device) == "outlet", ), @@ -151,7 +161,7 @@ class VeSyncSensorEntity(VeSyncBaseEntity, SensorEntity): def __init__( self, - device: VeSyncDevice, + device: VeSyncAirBypass | VeSyncOutlet | VeSyncSwitch, description: VeSyncSensorEntityDescription, ) -> None: """Initialize the VeSync outlet device."""