From b3ef6f4d047358048ecf12a4b4e07379cc90bcba Mon Sep 17 00:00:00 2001 From: Aaron Bach Date: Mon, 18 Jul 2022 09:18:07 -0600 Subject: [PATCH] Simplify Guardian entity inheritance hierarchy (#75274) --- homeassistant/components/guardian/__init__.py | 107 ++++++--------- .../components/guardian/binary_sensor.py | 127 +++++++++--------- homeassistant/components/guardian/button.py | 49 ++++--- homeassistant/components/guardian/sensor.py | 112 +++++++-------- homeassistant/components/guardian/switch.py | 78 ++++++----- 5 files changed, 225 insertions(+), 248 deletions(-) diff --git a/homeassistant/components/guardian/__init__.py b/homeassistant/components/guardian/__init__.py index f13ca1a7ff5..d15b7d57aea 100644 --- a/homeassistant/components/guardian/__init__.py +++ b/homeassistant/components/guardian/__init__.py @@ -3,6 +3,7 @@ from __future__ import annotations import asyncio from collections.abc import Awaitable, Callable +from dataclasses import dataclass from typing import cast from aioguardian import Client @@ -24,10 +25,7 @@ from homeassistant.exceptions import HomeAssistantError from homeassistant.helpers import config_validation as cv, device_registry as dr from homeassistant.helpers.dispatcher import async_dispatcher_send from homeassistant.helpers.entity import DeviceInfo, EntityDescription -from homeassistant.helpers.update_coordinator import ( - CoordinatorEntity, - DataUpdateCoordinator, -) +from homeassistant.helpers.update_coordinator import CoordinatorEntity from .const import ( API_SENSOR_PAIR_DUMP, @@ -132,7 +130,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: # so we use a lock to ensure that only one API request is reaching it at a time: api_lock = asyncio.Lock() - # Set up DataUpdateCoordinators for the valve controller: + # Set up GuardianDataUpdateCoordinators for the valve controller: coordinators: dict[str, GuardianDataUpdateCoordinator] = {} init_valve_controller_tasks = [] for api, api_coro in ( @@ -418,17 +416,18 @@ class PairedSensorManager: dev_reg.async_remove_device(device.id) -class GuardianEntity(CoordinatorEntity): +class GuardianEntity(CoordinatorEntity[GuardianDataUpdateCoordinator]): """Define a base Guardian entity.""" _attr_has_entity_name = True - def __init__( # pylint: disable=super-init-not-called - self, entry: ConfigEntry, description: EntityDescription + def __init__( + self, coordinator: GuardianDataUpdateCoordinator, description: EntityDescription ) -> None: """Initialize.""" + super().__init__(coordinator) + self._attr_extra_state_attributes = {} - self._entry = entry self.entity_description = description @callback @@ -438,6 +437,17 @@ class GuardianEntity(CoordinatorEntity): This should be extended by Guardian platforms. """ + @callback + def _handle_coordinator_update(self) -> None: + """Respond to a DataUpdateCoordinator update.""" + self._async_update_from_latest_data() + self.async_write_ha_state() + + async def async_added_to_hass(self) -> None: + """Handle entity which will be added.""" + await super().async_added_to_hass() + self._async_update_from_latest_data() + class PairedSensorEntity(GuardianEntity): """Define a Guardian paired sensor entity.""" @@ -445,11 +455,11 @@ class PairedSensorEntity(GuardianEntity): def __init__( self, entry: ConfigEntry, - coordinator: DataUpdateCoordinator, + coordinator: GuardianDataUpdateCoordinator, description: EntityDescription, ) -> None: """Initialize.""" - super().__init__(entry, description) + super().__init__(coordinator, description) paired_sensor_uid = coordinator.data["uid"] self._attr_device_info = DeviceInfo( @@ -460,11 +470,20 @@ class PairedSensorEntity(GuardianEntity): via_device=(DOMAIN, entry.data[CONF_UID]), ) self._attr_unique_id = f"{paired_sensor_uid}_{description.key}" - self.coordinator = coordinator - async def async_added_to_hass(self) -> None: - """Perform tasks when the entity is added.""" - self._async_update_from_latest_data() + +@dataclass +class ValveControllerEntityDescriptionMixin: + """Define an entity description mixin for valve controller entities.""" + + api_category: str + + +@dataclass +class ValveControllerEntityDescription( + EntityDescription, ValveControllerEntityDescriptionMixin +): + """Describe a Guardian valve controller entity.""" class ValveControllerEntity(GuardianEntity): @@ -473,64 +492,18 @@ class ValveControllerEntity(GuardianEntity): def __init__( self, entry: ConfigEntry, - coordinators: dict[str, DataUpdateCoordinator], - description: EntityDescription, + coordinators: dict[str, GuardianDataUpdateCoordinator], + description: ValveControllerEntityDescription, ) -> None: """Initialize.""" - super().__init__(entry, description) + super().__init__(coordinators[description.api_category], description) + + self._diagnostics_coordinator = coordinators[API_SYSTEM_DIAGNOSTICS] self._attr_device_info = DeviceInfo( identifiers={(DOMAIN, entry.data[CONF_UID])}, manufacturer="Elexa", - model=coordinators[API_SYSTEM_DIAGNOSTICS].data["firmware"], + model=self._diagnostics_coordinator.data["firmware"], name=f"Guardian valve controller {entry.data[CONF_UID]}", ) self._attr_unique_id = f"{entry.data[CONF_UID]}_{description.key}" - self.coordinators = coordinators - - @property - def available(self) -> bool: - """Return if entity is available.""" - return any( - coordinator.last_update_success - for coordinator in self.coordinators.values() - ) - - async def _async_continue_entity_setup(self) -> None: - """Perform additional, internal tasks when the entity is about to be added. - - This should be extended by Guardian platforms. - """ - - @callback - def async_add_coordinator_update_listener(self, api: str) -> None: - """Add a listener to a DataUpdateCoordinator based on the API referenced.""" - - @callback - def update() -> None: - """Update the entity's state.""" - self._async_update_from_latest_data() - self.async_write_ha_state() - - self.async_on_remove(self.coordinators[api].async_add_listener(update)) - - async def async_added_to_hass(self) -> None: - """Perform tasks when the entity is added.""" - await self._async_continue_entity_setup() - self.async_add_coordinator_update_listener(API_SYSTEM_DIAGNOSTICS) - self._async_update_from_latest_data() - - async def async_update(self) -> None: - """Update the entity. - - Only used by the generic entity update service. - """ - # Ignore manual update requests if the entity is disabled - if not self.enabled: - return - - refresh_tasks = [ - coordinator.async_request_refresh() - for coordinator in self.coordinators.values() - ] - await asyncio.gather(*refresh_tasks) diff --git a/homeassistant/components/guardian/binary_sensor.py b/homeassistant/components/guardian/binary_sensor.py index dc9febcfa91..9b824ab589f 100644 --- a/homeassistant/components/guardian/binary_sensor.py +++ b/homeassistant/components/guardian/binary_sensor.py @@ -1,6 +1,8 @@ """Binary sensors for the Elexa Guardian integration.""" from __future__ import annotations +from dataclasses import dataclass + from homeassistant.components.binary_sensor import ( BinarySensorDeviceClass, BinarySensorEntity, @@ -11,9 +13,12 @@ from homeassistant.core import HomeAssistant, callback from homeassistant.helpers.dispatcher import async_dispatcher_connect from homeassistant.helpers.entity import EntityCategory from homeassistant.helpers.entity_platform import AddEntitiesCallback -from homeassistant.helpers.update_coordinator import DataUpdateCoordinator -from . import PairedSensorEntity, ValveControllerEntity +from . import ( + PairedSensorEntity, + ValveControllerEntity, + ValveControllerEntityDescription, +) from .const import ( API_SYSTEM_ONBOARD_SENSOR_STATUS, API_WIFI_STATUS, @@ -23,6 +28,7 @@ from .const import ( DOMAIN, SIGNAL_PAIRED_SENSOR_COORDINATOR_ADDED, ) +from .util import GuardianDataUpdateCoordinator ATTR_CONNECTED_CLIENTS = "connected_clients" @@ -30,31 +36,42 @@ SENSOR_KIND_AP_INFO = "ap_enabled" SENSOR_KIND_LEAK_DETECTED = "leak_detected" SENSOR_KIND_MOVED = "moved" -SENSOR_DESCRIPTION_AP_ENABLED = BinarySensorEntityDescription( - key=SENSOR_KIND_AP_INFO, - name="Onboard AP enabled", - device_class=BinarySensorDeviceClass.CONNECTIVITY, - entity_category=EntityCategory.DIAGNOSTIC, -) -SENSOR_DESCRIPTION_LEAK_DETECTED = BinarySensorEntityDescription( - key=SENSOR_KIND_LEAK_DETECTED, - name="Leak detected", - device_class=BinarySensorDeviceClass.MOISTURE, -) -SENSOR_DESCRIPTION_MOVED = BinarySensorEntityDescription( - key=SENSOR_KIND_MOVED, - name="Recently moved", - device_class=BinarySensorDeviceClass.MOVING, - entity_category=EntityCategory.DIAGNOSTIC, -) + +@dataclass +class ValveControllerBinarySensorDescription( + BinarySensorEntityDescription, ValveControllerEntityDescription +): + """Describe a Guardian valve controller binary sensor.""" + PAIRED_SENSOR_DESCRIPTIONS = ( - SENSOR_DESCRIPTION_LEAK_DETECTED, - SENSOR_DESCRIPTION_MOVED, + BinarySensorEntityDescription( + key=SENSOR_KIND_LEAK_DETECTED, + name="Leak detected", + device_class=BinarySensorDeviceClass.MOISTURE, + ), + BinarySensorEntityDescription( + key=SENSOR_KIND_MOVED, + name="Recently moved", + device_class=BinarySensorDeviceClass.MOVING, + entity_category=EntityCategory.DIAGNOSTIC, + ), ) + VALVE_CONTROLLER_DESCRIPTIONS = ( - SENSOR_DESCRIPTION_AP_ENABLED, - SENSOR_DESCRIPTION_LEAK_DETECTED, + ValveControllerBinarySensorDescription( + key=SENSOR_KIND_LEAK_DETECTED, + name="Leak detected", + device_class=BinarySensorDeviceClass.MOISTURE, + api_category=API_SYSTEM_ONBOARD_SENSOR_STATUS, + ), + ValveControllerBinarySensorDescription( + key=SENSOR_KIND_AP_INFO, + name="Onboard AP enabled", + device_class=BinarySensorDeviceClass.CONNECTIVITY, + entity_category=EntityCategory.DIAGNOSTIC, + api_category=API_WIFI_STATUS, + ), ) @@ -62,19 +79,18 @@ async def async_setup_entry( hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback ) -> None: """Set up Guardian switches based on a config entry.""" + entry_data = hass.data[DOMAIN][entry.entry_id] + paired_sensor_coordinators = entry_data[DATA_COORDINATOR_PAIRED_SENSOR] + valve_controller_coordinators = entry_data[DATA_COORDINATOR] @callback def add_new_paired_sensor(uid: str) -> None: """Add a new paired sensor.""" - coordinator = hass.data[DOMAIN][entry.entry_id][DATA_COORDINATOR_PAIRED_SENSOR][ - uid - ] - async_add_entities( - [ - PairedSensorBinarySensor(entry, coordinator, description) - for description in PAIRED_SENSOR_DESCRIPTIONS - ] + PairedSensorBinarySensor( + entry, paired_sensor_coordinators[uid], description + ) + for description in PAIRED_SENSOR_DESCRIPTIONS ) # Handle adding paired sensors after HASS startup: @@ -88,9 +104,7 @@ async def async_setup_entry( # Add all valve controller-specific binary sensors: sensors: list[PairedSensorBinarySensor | ValveControllerBinarySensor] = [ - ValveControllerBinarySensor( - entry, hass.data[DOMAIN][entry.entry_id][DATA_COORDINATOR], description - ) + ValveControllerBinarySensor(entry, valve_controller_coordinators, description) for description in VALVE_CONTROLLER_DESCRIPTIONS ] @@ -98,9 +112,7 @@ async def async_setup_entry( sensors.extend( [ PairedSensorBinarySensor(entry, coordinator, description) - for coordinator in hass.data[DOMAIN][entry.entry_id][ - DATA_COORDINATOR_PAIRED_SENSOR - ].values() + for coordinator in paired_sensor_coordinators.values() for description in PAIRED_SENSOR_DESCRIPTIONS ] ) @@ -111,10 +123,12 @@ async def async_setup_entry( class PairedSensorBinarySensor(PairedSensorEntity, BinarySensorEntity): """Define a binary sensor related to a Guardian valve controller.""" + entity_description: BinarySensorEntityDescription + def __init__( self, entry: ConfigEntry, - coordinator: DataUpdateCoordinator, + coordinator: GuardianDataUpdateCoordinator, description: BinarySensorEntityDescription, ) -> None: """Initialize.""" @@ -134,45 +148,26 @@ class PairedSensorBinarySensor(PairedSensorEntity, BinarySensorEntity): class ValveControllerBinarySensor(ValveControllerEntity, BinarySensorEntity): """Define a binary sensor related to a Guardian valve controller.""" + entity_description: ValveControllerBinarySensorDescription + def __init__( self, entry: ConfigEntry, - coordinators: dict[str, DataUpdateCoordinator], - description: BinarySensorEntityDescription, + coordinators: dict[str, GuardianDataUpdateCoordinator], + description: ValveControllerBinarySensorDescription, ) -> None: """Initialize.""" super().__init__(entry, coordinators, description) self._attr_is_on = True - async def _async_continue_entity_setup(self) -> None: - """Add an API listener.""" - if self.entity_description.key == SENSOR_KIND_AP_INFO: - self.async_add_coordinator_update_listener(API_WIFI_STATUS) - elif self.entity_description.key == SENSOR_KIND_LEAK_DETECTED: - self.async_add_coordinator_update_listener(API_SYSTEM_ONBOARD_SENSOR_STATUS) - @callback def _async_update_from_latest_data(self) -> None: """Update the entity.""" if self.entity_description.key == SENSOR_KIND_AP_INFO: - self._attr_available = self.coordinators[ - API_WIFI_STATUS - ].last_update_success - self._attr_is_on = self.coordinators[API_WIFI_STATUS].data[ - "station_connected" - ] - self._attr_extra_state_attributes.update( - { - ATTR_CONNECTED_CLIENTS: self.coordinators[API_WIFI_STATUS].data.get( - "ap_clients" - ) - } - ) + self._attr_is_on = self.coordinator.data["station_connected"] + self._attr_extra_state_attributes[ + ATTR_CONNECTED_CLIENTS + ] = self.coordinator.data.get("ap_clients") elif self.entity_description.key == SENSOR_KIND_LEAK_DETECTED: - self._attr_available = self.coordinators[ - API_SYSTEM_ONBOARD_SENSOR_STATUS - ].last_update_success - self._attr_is_on = self.coordinators[API_SYSTEM_ONBOARD_SENSOR_STATUS].data[ - "wet" - ] + self._attr_is_on = self.coordinator.data["wet"] diff --git a/homeassistant/components/guardian/button.py b/homeassistant/components/guardian/button.py index e7cc757d367..e013cde85d6 100644 --- a/homeassistant/components/guardian/button.py +++ b/homeassistant/components/guardian/button.py @@ -17,24 +17,26 @@ from homeassistant.core import HomeAssistant from homeassistant.exceptions import HomeAssistantError from homeassistant.helpers.entity import EntityCategory from homeassistant.helpers.entity_platform import AddEntitiesCallback -from homeassistant.helpers.update_coordinator import DataUpdateCoordinator -from . import ValveControllerEntity -from .const import DATA_CLIENT, DATA_COORDINATOR, DOMAIN +from . import ValveControllerEntity, ValveControllerEntityDescription +from .const import API_SYSTEM_DIAGNOSTICS, DATA_CLIENT, DATA_COORDINATOR, DOMAIN +from .util import GuardianDataUpdateCoordinator @dataclass -class GuardianButtonDescriptionMixin: - """Define an entity description mixin for Guardian buttons.""" +class GuardianButtonEntityDescriptionMixin: + """Define an mixin for button entities.""" push_action: Callable[[Client], Awaitable] @dataclass -class GuardianButtonDescription( - ButtonEntityDescription, GuardianButtonDescriptionMixin +class ValveControllerButtonDescription( + ButtonEntityDescription, + ValveControllerEntityDescription, + GuardianButtonEntityDescriptionMixin, ): - """Describe a Guardian button description.""" + """Describe a Guardian valve controller button.""" BUTTON_KIND_REBOOT = "reboot" @@ -52,15 +54,21 @@ async def _async_valve_reset(client: Client) -> None: BUTTON_DESCRIPTIONS = ( - GuardianButtonDescription( + ValveControllerButtonDescription( key=BUTTON_KIND_REBOOT, name="Reboot", push_action=_async_reboot, + # Buttons don't actually need a coordinator; we give them one so they can + # properly inherit from GuardianEntity: + api_category=API_SYSTEM_DIAGNOSTICS, ), - GuardianButtonDescription( + ValveControllerButtonDescription( key=BUTTON_KIND_RESET_VALVE_DIAGNOSTICS, name="Reset valve diagnostics", push_action=_async_valve_reset, + # Buttons don't actually need a coordinator; we give them one so they can + # properly inherit from GuardianEntity: + api_category=API_SYSTEM_DIAGNOSTICS, ), ) @@ -69,16 +77,13 @@ async def async_setup_entry( hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback ) -> None: """Set up Guardian buttons based on a config entry.""" + entry_data = hass.data[DOMAIN][entry.entry_id] + client = entry_data[DATA_CLIENT] + valve_controller_coordinators = entry_data[DATA_COORDINATOR] + async_add_entities( - [ - GuardianButton( - entry, - hass.data[DOMAIN][entry.entry_id][DATA_CLIENT], - hass.data[DOMAIN][entry.entry_id][DATA_COORDINATOR], - description, - ) - for description in BUTTON_DESCRIPTIONS - ] + GuardianButton(entry, valve_controller_coordinators, description, client) + for description in BUTTON_DESCRIPTIONS ) @@ -88,14 +93,14 @@ class GuardianButton(ValveControllerEntity, ButtonEntity): _attr_device_class = ButtonDeviceClass.RESTART _attr_entity_category = EntityCategory.CONFIG - entity_description: GuardianButtonDescription + entity_description: ValveControllerButtonDescription def __init__( self, entry: ConfigEntry, + coordinators: dict[str, GuardianDataUpdateCoordinator], + description: ValveControllerButtonDescription, client: Client, - coordinators: dict[str, DataUpdateCoordinator], - description: GuardianButtonDescription, ) -> None: """Initialize.""" super().__init__(entry, coordinators, description) diff --git a/homeassistant/components/guardian/sensor.py b/homeassistant/components/guardian/sensor.py index 895452e0fda..c4fd1e110fa 100644 --- a/homeassistant/components/guardian/sensor.py +++ b/homeassistant/components/guardian/sensor.py @@ -1,6 +1,8 @@ """Sensors for the Elexa Guardian integration.""" from __future__ import annotations +from dataclasses import dataclass + from homeassistant.components.sensor import ( SensorDeviceClass, SensorEntity, @@ -14,7 +16,11 @@ from homeassistant.helpers.dispatcher import async_dispatcher_connect from homeassistant.helpers.entity import EntityCategory from homeassistant.helpers.entity_platform import AddEntitiesCallback -from . import PairedSensorEntity, ValveControllerEntity +from . import ( + PairedSensorEntity, + ValveControllerEntity, + ValveControllerEntityDescription, +) from .const import ( API_SYSTEM_DIAGNOSTICS, API_SYSTEM_ONBOARD_SENSOR_STATUS, @@ -29,35 +35,47 @@ SENSOR_KIND_BATTERY = "battery" SENSOR_KIND_TEMPERATURE = "temperature" SENSOR_KIND_UPTIME = "uptime" -SENSOR_DESCRIPTION_BATTERY = SensorEntityDescription( - key=SENSOR_KIND_BATTERY, - name="Battery", - device_class=SensorDeviceClass.BATTERY, - entity_category=EntityCategory.DIAGNOSTIC, - native_unit_of_measurement=PERCENTAGE, -) -SENSOR_DESCRIPTION_TEMPERATURE = SensorEntityDescription( - key=SENSOR_KIND_TEMPERATURE, - name="Temperature", - device_class=SensorDeviceClass.TEMPERATURE, - native_unit_of_measurement=TEMP_FAHRENHEIT, - state_class=SensorStateClass.MEASUREMENT, -) -SENSOR_DESCRIPTION_UPTIME = SensorEntityDescription( - key=SENSOR_KIND_UPTIME, - name="Uptime", - icon="mdi:timer", - entity_category=EntityCategory.DIAGNOSTIC, - native_unit_of_measurement=TIME_MINUTES, -) + +@dataclass +class ValveControllerSensorDescription( + SensorEntityDescription, ValveControllerEntityDescription +): + """Describe a Guardian valve controller sensor.""" + PAIRED_SENSOR_DESCRIPTIONS = ( - SENSOR_DESCRIPTION_BATTERY, - SENSOR_DESCRIPTION_TEMPERATURE, + SensorEntityDescription( + key=SENSOR_KIND_BATTERY, + name="Battery", + device_class=SensorDeviceClass.BATTERY, + entity_category=EntityCategory.DIAGNOSTIC, + native_unit_of_measurement=PERCENTAGE, + ), + SensorEntityDescription( + key=SENSOR_KIND_TEMPERATURE, + name="Temperature", + device_class=SensorDeviceClass.TEMPERATURE, + native_unit_of_measurement=TEMP_FAHRENHEIT, + state_class=SensorStateClass.MEASUREMENT, + ), ) VALVE_CONTROLLER_DESCRIPTIONS = ( - SENSOR_DESCRIPTION_TEMPERATURE, - SENSOR_DESCRIPTION_UPTIME, + ValveControllerSensorDescription( + key=SENSOR_KIND_TEMPERATURE, + name="Temperature", + device_class=SensorDeviceClass.TEMPERATURE, + native_unit_of_measurement=TEMP_FAHRENHEIT, + state_class=SensorStateClass.MEASUREMENT, + api_category=API_SYSTEM_ONBOARD_SENSOR_STATUS, + ), + ValveControllerSensorDescription( + key=SENSOR_KIND_UPTIME, + name="Uptime", + icon="mdi:timer", + entity_category=EntityCategory.DIAGNOSTIC, + native_unit_of_measurement=TIME_MINUTES, + api_category=API_SYSTEM_DIAGNOSTICS, + ), ) @@ -65,19 +83,16 @@ async def async_setup_entry( hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback ) -> None: """Set up Guardian switches based on a config entry.""" + entry_data = hass.data[DOMAIN][entry.entry_id] + paired_sensor_coordinators = entry_data[DATA_COORDINATOR_PAIRED_SENSOR] + valve_controller_coordinators = entry_data[DATA_COORDINATOR] @callback def add_new_paired_sensor(uid: str) -> None: """Add a new paired sensor.""" - coordinator = hass.data[DOMAIN][entry.entry_id][DATA_COORDINATOR_PAIRED_SENSOR][ - uid - ] - async_add_entities( - [ - PairedSensorSensor(entry, coordinator, description) - for description in PAIRED_SENSOR_DESCRIPTIONS - ] + PairedSensorSensor(entry, paired_sensor_coordinators[uid], description) + for description in PAIRED_SENSOR_DESCRIPTIONS ) # Handle adding paired sensors after HASS startup: @@ -91,9 +106,7 @@ async def async_setup_entry( # Add all valve controller-specific binary sensors: sensors: list[PairedSensorSensor | ValveControllerSensor] = [ - ValveControllerSensor( - entry, hass.data[DOMAIN][entry.entry_id][DATA_COORDINATOR], description - ) + ValveControllerSensor(entry, valve_controller_coordinators, description) for description in VALVE_CONTROLLER_DESCRIPTIONS ] @@ -101,9 +114,7 @@ async def async_setup_entry( sensors.extend( [ PairedSensorSensor(entry, coordinator, description) - for coordinator in hass.data[DOMAIN][entry.entry_id][ - DATA_COORDINATOR_PAIRED_SENSOR - ].values() + for coordinator in paired_sensor_coordinators.values() for description in PAIRED_SENSOR_DESCRIPTIONS ] ) @@ -114,6 +125,8 @@ async def async_setup_entry( class PairedSensorSensor(PairedSensorEntity, SensorEntity): """Define a binary sensor related to a Guardian valve controller.""" + entity_description: SensorEntityDescription + @callback def _async_update_from_latest_data(self) -> None: """Update the entity.""" @@ -126,25 +139,12 @@ class PairedSensorSensor(PairedSensorEntity, SensorEntity): class ValveControllerSensor(ValveControllerEntity, SensorEntity): """Define a generic Guardian sensor.""" - async def _async_continue_entity_setup(self) -> None: - """Register API interest (and related tasks) when the entity is added.""" - if self.entity_description.key == SENSOR_KIND_TEMPERATURE: - self.async_add_coordinator_update_listener(API_SYSTEM_ONBOARD_SENSOR_STATUS) + entity_description: ValveControllerSensorDescription @callback def _async_update_from_latest_data(self) -> None: """Update the entity.""" if self.entity_description.key == SENSOR_KIND_TEMPERATURE: - self._attr_available = self.coordinators[ - API_SYSTEM_ONBOARD_SENSOR_STATUS - ].last_update_success - self._attr_native_value = self.coordinators[ - API_SYSTEM_ONBOARD_SENSOR_STATUS - ].data["temperature"] + self._attr_native_value = self.coordinator.data["temperature"] elif self.entity_description.key == SENSOR_KIND_UPTIME: - self._attr_available = self.coordinators[ - API_SYSTEM_DIAGNOSTICS - ].last_update_success - self._attr_native_value = self.coordinators[API_SYSTEM_DIAGNOSTICS].data[ - "uptime" - ] + self._attr_native_value = self.coordinator.data["uptime"] diff --git a/homeassistant/components/guardian/switch.py b/homeassistant/components/guardian/switch.py index 485b0a3ffbc..c58f4548a87 100644 --- a/homeassistant/components/guardian/switch.py +++ b/homeassistant/components/guardian/switch.py @@ -1,6 +1,7 @@ """Switches for the Elexa Guardian integration.""" from __future__ import annotations +from dataclasses import dataclass from typing import Any from aioguardian import Client @@ -11,10 +12,10 @@ from homeassistant.config_entries import ConfigEntry from homeassistant.core import HomeAssistant, callback from homeassistant.exceptions import HomeAssistantError from homeassistant.helpers.entity_platform import AddEntitiesCallback -from homeassistant.helpers.update_coordinator import DataUpdateCoordinator -from . import ValveControllerEntity +from . import ValveControllerEntity, ValveControllerEntityDescription from .const import API_VALVE_STATUS, DATA_CLIENT, DATA_COORDINATOR, DOMAIN +from .util import GuardianDataUpdateCoordinator ATTR_AVG_CURRENT = "average_current" ATTR_INST_CURRENT = "instantaneous_current" @@ -23,10 +24,21 @@ ATTR_TRAVEL_COUNT = "travel_count" SWITCH_KIND_VALVE = "valve" -SWITCH_DESCRIPTION_VALVE = SwitchEntityDescription( - key=SWITCH_KIND_VALVE, - name="Valve controller", - icon="mdi:water", + +@dataclass +class ValveControllerSwitchDescription( + SwitchEntityDescription, ValveControllerEntityDescription +): + """Describe a Guardian valve controller switch.""" + + +VALVE_CONTROLLER_DESCRIPTIONS = ( + ValveControllerSwitchDescription( + key=SWITCH_KIND_VALVE, + name="Valve controller", + icon="mdi:water", + api_category=API_VALVE_STATUS, + ), ) @@ -34,61 +46,53 @@ async def async_setup_entry( hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback ) -> None: """Set up Guardian switches based on a config entry.""" + entry_data = hass.data[DOMAIN][entry.entry_id] + client = entry_data[DATA_CLIENT] + valve_controller_coordinators = entry_data[DATA_COORDINATOR] + async_add_entities( - [ - ValveControllerSwitch( - entry, - hass.data[DOMAIN][entry.entry_id][DATA_CLIENT], - hass.data[DOMAIN][entry.entry_id][DATA_COORDINATOR], - ) - ] + ValveControllerSwitch(entry, valve_controller_coordinators, description, client) + for description in VALVE_CONTROLLER_DESCRIPTIONS ) class ValveControllerSwitch(ValveControllerEntity, SwitchEntity): """Define a switch to open/close the Guardian valve.""" + entity_description: ValveControllerSwitchDescription + + ON_STATES = { + "start_opening", + "opening", + "finish_opening", + "opened", + } + def __init__( self, entry: ConfigEntry, + coordinators: dict[str, GuardianDataUpdateCoordinator], + description: ValveControllerSwitchDescription, client: Client, - coordinators: dict[str, DataUpdateCoordinator], ) -> None: """Initialize.""" - super().__init__(entry, coordinators, SWITCH_DESCRIPTION_VALVE) + super().__init__(entry, coordinators, description) self._attr_is_on = True self._client = client - async def _async_continue_entity_setup(self) -> None: - """Register API interest (and related tasks) when the entity is added.""" - self.async_add_coordinator_update_listener(API_VALVE_STATUS) - @callback def _async_update_from_latest_data(self) -> None: """Update the entity.""" - self._attr_available = self.coordinators[API_VALVE_STATUS].last_update_success - self._attr_is_on = self.coordinators[API_VALVE_STATUS].data["state"] in ( - "start_opening", - "opening", - "finish_opening", - "opened", - ) - + self._attr_is_on = self.coordinator.data["state"] in self.ON_STATES self._attr_extra_state_attributes.update( { - ATTR_AVG_CURRENT: self.coordinators[API_VALVE_STATUS].data[ - "average_current" - ], - ATTR_INST_CURRENT: self.coordinators[API_VALVE_STATUS].data[ - "instantaneous_current" - ], - ATTR_INST_CURRENT_DDT: self.coordinators[API_VALVE_STATUS].data[ + ATTR_AVG_CURRENT: self.coordinator.data["average_current"], + ATTR_INST_CURRENT: self.coordinator.data["instantaneous_current"], + ATTR_INST_CURRENT_DDT: self.coordinator.data[ "instantaneous_current_ddt" ], - ATTR_TRAVEL_COUNT: self.coordinators[API_VALVE_STATUS].data[ - "travel_count" - ], + ATTR_TRAVEL_COUNT: self.coordinator.data["travel_count"], } )