From 0067b6a84d2cce38fa7f74ef0de13044da3f5a09 Mon Sep 17 00:00:00 2001 From: Aaron Bach Date: Sun, 5 Jul 2020 16:09:40 -0600 Subject: [PATCH] Transition Guardian to use a DataUpdateCoordinator (#37380) * Migrate Guardian to use the DataUpdateCoordinator * Finish work * Cleanup * Don't use UpdateFailed error * Code cleanup * Code cleanup * Remove unnecessary change * Code review * Code review * Use a subclass of DataUpdateCoordinator * Make sure to pop client upon unload * Adjust coverage --- .coveragerc | 1 + homeassistant/components/guardian/__init__.py | 265 ++++++------------ .../components/guardian/binary_sensor.py | 72 ++++- homeassistant/components/guardian/const.py | 21 +- homeassistant/components/guardian/sensor.py | 73 +++-- homeassistant/components/guardian/switch.py | 90 ++++-- homeassistant/components/guardian/util.py | 49 ++++ 7 files changed, 320 insertions(+), 251 deletions(-) create mode 100644 homeassistant/components/guardian/util.py diff --git a/.coveragerc b/.coveragerc index 984a9c173cb..67bea8dae16 100644 --- a/.coveragerc +++ b/.coveragerc @@ -313,6 +313,7 @@ omit = homeassistant/components/guardian/binary_sensor.py homeassistant/components/guardian/sensor.py homeassistant/components/guardian/switch.py + homeassistant/components/guardian/util.py homeassistant/components/habitica/* homeassistant/components/hangouts/* homeassistant/components/hangouts/__init__.py diff --git a/homeassistant/components/guardian/__init__.py b/homeassistant/components/guardian/__init__.py index 816c9ccc0f1..03796415d65 100644 --- a/homeassistant/components/guardian/__init__.py +++ b/homeassistant/components/guardian/__init__.py @@ -1,67 +1,72 @@ """The Elexa Guardian integration.""" import asyncio from datetime import timedelta +from typing import Dict from aioguardian import Client -from aioguardian.errors import GuardianError from homeassistant.config_entries import ConfigEntry -from homeassistant.const import ATTR_ATTRIBUTION, CONF_IP_ADDRESS +from homeassistant.const import ATTR_ATTRIBUTION, CONF_IP_ADDRESS, CONF_PORT from homeassistant.core import HomeAssistant, callback -from homeassistant.helpers.dispatcher import ( - async_dispatcher_connect, - async_dispatcher_send, -) from homeassistant.helpers.entity import Entity -from homeassistant.helpers.event import async_track_time_interval +from homeassistant.helpers.update_coordinator import DataUpdateCoordinator from .const import ( + API_SYSTEM_DIAGNOSTICS, + API_SYSTEM_ONBOARD_SENSOR_STATUS, + API_VALVE_STATUS, + API_WIFI_STATUS, CONF_UID, DATA_CLIENT, - DATA_DIAGNOSTICS, - DATA_PAIR_DUMP, - DATA_PING, - DATA_SENSOR_STATUS, - DATA_VALVE_STATUS, - DATA_WIFI_STATUS, + DATA_COORDINATOR, DOMAIN, - LOGGER, - SENSOR_KIND_AP_INFO, - SENSOR_KIND_LEAK_DETECTED, - SENSOR_KIND_TEMPERATURE, - SWITCH_KIND_VALVE, - TOPIC_UPDATE, ) +from .util import GuardianDataUpdateCoordinator -DATA_ENTITY_TYPE_MAP = { - SENSOR_KIND_AP_INFO: DATA_WIFI_STATUS, - SENSOR_KIND_LEAK_DETECTED: DATA_SENSOR_STATUS, - SENSOR_KIND_TEMPERATURE: DATA_SENSOR_STATUS, - SWITCH_KIND_VALVE: DATA_VALVE_STATUS, -} - -DEFAULT_SCAN_INTERVAL = timedelta(seconds=30) +DEFAULT_UPDATE_INTERVAL = timedelta(seconds=30) PLATFORMS = ["binary_sensor", "sensor", "switch"] -@callback -def async_get_api_category(entity_kind: str): - """Get the API data category to which an entity belongs.""" - return DATA_ENTITY_TYPE_MAP.get(entity_kind) - - -async def async_setup(hass: HomeAssistant, config: dict): +async def async_setup(hass: HomeAssistant, config: dict) -> bool: """Set up the Elexa Guardian component.""" - hass.data[DOMAIN] = {DATA_CLIENT: {}} + hass.data[DOMAIN] = {DATA_CLIENT: {}, DATA_COORDINATOR: {}} return True -async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry): +async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: """Set up Elexa Guardian from a config entry.""" - guardian = Guardian(hass, entry) - await guardian.async_update() - hass.data[DOMAIN][DATA_CLIENT][entry.entry_id] = guardian + client = hass.data[DOMAIN][DATA_CLIENT][entry.entry_id] = Client( + entry.data[CONF_IP_ADDRESS], port=entry.data[CONF_PORT] + ) + hass.data[DOMAIN][DATA_COORDINATOR][entry.entry_id] = {} + + # The valve controller's UDP-based API can't handle concurrent requests very well, + # so we use a lock to ensure that only one API request is reaching it at a time: + api_lock = asyncio.Lock() + initial_fetch_tasks = [] + + for api, api_coro in [ + (API_SYSTEM_DIAGNOSTICS, client.system.diagnostics), + (API_SYSTEM_ONBOARD_SENSOR_STATUS, client.system.onboard_sensor_status), + (API_VALVE_STATUS, client.valve.status), + (API_WIFI_STATUS, client.wifi.status), + ]: + hass.data[DOMAIN][DATA_COORDINATOR][entry.entry_id][ + api + ] = GuardianDataUpdateCoordinator( + hass, + client=client, + api_name=api, + api_coro=api_coro, + api_lock=api_lock, + valve_controller_uid=entry.data[CONF_UID], + ) + initial_fetch_tasks.append( + hass.data[DOMAIN][DATA_COORDINATOR][entry.entry_id][api].async_refresh() + ) + + await asyncio.gather(*initial_fetch_tasks) for component in PLATFORMS: hass.async_create_task( @@ -71,7 +76,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry): return True -async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry): +async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: """Unload a config entry.""" unload_ok = all( await asyncio.gather( @@ -83,143 +88,52 @@ async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry): ) if unload_ok: hass.data[DOMAIN][DATA_CLIENT].pop(entry.entry_id) + hass.data[DOMAIN][DATA_COORDINATOR].pop(entry.entry_id) return unload_ok -class Guardian: - """Define a class to communicate with the Guardian device.""" - - def __init__(self, hass: HomeAssistant, entry: ConfigEntry): - """Initialize.""" - self._async_cancel_time_interval_listener = None - self._hass = hass - self.client = Client(entry.data[CONF_IP_ADDRESS]) - self.data = {} - self.uid = entry.data[CONF_UID] - - self._api_coros = { - DATA_DIAGNOSTICS: self.client.system.diagnostics, - DATA_PAIR_DUMP: self.client.sensor.pair_dump, - DATA_PING: self.client.system.ping, - DATA_SENSOR_STATUS: self.client.system.onboard_sensor_status, - DATA_VALVE_STATUS: self.client.valve.status, - DATA_WIFI_STATUS: self.client.wifi.status, - } - - self._api_category_count = { - DATA_SENSOR_STATUS: 0, - DATA_VALVE_STATUS: 0, - DATA_WIFI_STATUS: 0, - } - - self._api_lock = asyncio.Lock() - - async def _async_get_data_from_api(self, api_category: str): - """Update and save data for a particular API category.""" - if self._api_category_count.get(api_category) == 0: - return - - try: - result = await self._api_coros[api_category]() - except GuardianError as err: - LOGGER.error("Error while fetching %s data: %s", api_category, err) - self.data[api_category] = {} - else: - self.data[api_category] = result["data"] - - async def _async_update_listener_action(self, _): - """Define an async_track_time_interval action to update data.""" - await self.async_update() - - @callback - def async_deregister_api_interest(self, sensor_kind: str): - """Decrement the number of entities with data needs from an API category.""" - # If this deregistration should leave us with no registration at all, remove the - # time interval: - if sum(self._api_category_count.values()) == 0: - if self._async_cancel_time_interval_listener: - self._async_cancel_time_interval_listener() - self._async_cancel_time_interval_listener = None - return - - api_category = async_get_api_category(sensor_kind) - if api_category: - self._api_category_count[api_category] -= 1 - - async def async_register_api_interest(self, sensor_kind: str): - """Increment the number of entities with data needs from an API category.""" - # If this is the first registration we have, start a time interval: - if not self._async_cancel_time_interval_listener: - self._async_cancel_time_interval_listener = async_track_time_interval( - self._hass, self._async_update_listener_action, DEFAULT_SCAN_INTERVAL, - ) - - api_category = async_get_api_category(sensor_kind) - - if not api_category: - return - - self._api_category_count[api_category] += 1 - - # If a sensor registers interest in a particular API call and the data doesn't - # exist for it yet, make the API call and grab the data: - async with self._api_lock: - if api_category not in self.data: - async with self.client: - await self._async_get_data_from_api(api_category) - - async def async_update(self): - """Get updated data from the device.""" - async with self.client: - tasks = [ - self._async_get_data_from_api(api_category) - for api_category in self._api_coros - ] - - await asyncio.gather(*tasks) - - LOGGER.debug("Received new data: %s", self.data) - async_dispatcher_send(self._hass, TOPIC_UPDATE.format(self.uid)) - - class GuardianEntity(Entity): """Define a base Guardian entity.""" def __init__( - self, guardian: Guardian, kind: str, name: str, device_class: str, icon: str - ): + self, + entry: ConfigEntry, + client: Client, + coordinators: Dict[str, DataUpdateCoordinator], + kind: str, + name: str, + device_class: str, + icon: str, + ) -> None: """Initialize.""" self._attrs = {ATTR_ATTRIBUTION: "Data provided by Elexa"} self._available = True + self._client = client + self._coordinators = coordinators self._device_class = device_class - self._guardian = guardian self._icon = icon self._kind = kind self._name = name + self._valve_controller_uid = entry.data[CONF_UID] @property - def available(self): - """Return whether the entity is available.""" - return bool(self._guardian.data[DATA_PING]) - - @property - def device_class(self): + def device_class(self) -> str: """Return the device class.""" return self._device_class @property - def device_info(self): + def device_info(self) -> dict: """Return device registry information for this entity.""" return { - "identifiers": {(DOMAIN, self._guardian.uid)}, + "identifiers": {(DOMAIN, self._valve_controller_uid)}, "manufacturer": "Elexa", - "model": self._guardian.data[DATA_DIAGNOSTICS]["firmware"], - "name": f"Guardian {self._guardian.uid}", + "model": self._coordinators[API_SYSTEM_DIAGNOSTICS].data["firmware"], + "name": f"Guardian {self._valve_controller_uid}", } @property - def device_state_attributes(self): + def device_state_attributes(self) -> dict: """Return the state attributes.""" return self._attrs @@ -229,9 +143,9 @@ class GuardianEntity(Entity): return self._icon @property - def name(self): + def name(self) -> str: """Return the name of the entity.""" - return f"Guardian {self._guardian.uid}: {self._name}" + return f"Guardian {self._valve_controller_uid}: {self._name}" @property def should_poll(self) -> bool: @@ -241,32 +155,37 @@ class GuardianEntity(Entity): @property def unique_id(self): """Return the unique ID of the entity.""" - return f"{self._guardian.uid}_{self._kind}" + return f"{self._valve_controller_uid}_{self._kind}" - @callback - def _update_from_latest_data(self): - """Update the entity.""" + async def _async_internal_added_to_hass(self): + """Perform additional, internal tasks when the entity is about to be added. + + This should be extended by Guardian platforms. + """ raise NotImplementedError - async def async_added_to_hass(self): - """Register callbacks.""" + @callback + def _async_update_from_latest_data(self): + """Update the entity. + + This should be extended by Guardian platforms. + """ + raise NotImplementedError + + @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(): - """Update the state.""" - self._update_from_latest_data() + def async_update(): + """Update the entity's state.""" + self._async_update_from_latest_data() self.async_write_ha_state() - self.async_on_remove( - async_dispatcher_connect( - self.hass, TOPIC_UPDATE.format(self._guardian.uid), update - ) - ) + self.async_on_remove(self._coordinators[api].async_add_listener(async_update)) - await self._guardian.async_register_api_interest(self._kind) - - self._update_from_latest_data() - - async def async_will_remove_from_hass(self) -> None: - """Disconnect dispatcher listener when removed.""" - self._guardian.async_deregister_api_interest(self._kind) + async def async_added_to_hass(self) -> None: + """Perform tasks when the entity is added.""" + await self._async_internal_added_to_hass() + self.async_add_coordinator_update_listener(API_SYSTEM_DIAGNOSTICS) + self._async_update_from_latest_data() diff --git a/homeassistant/components/guardian/binary_sensor.py b/homeassistant/components/guardian/binary_sensor.py index f9d70d03d5d..495f325eb7f 100644 --- a/homeassistant/components/guardian/binary_sensor.py +++ b/homeassistant/components/guardian/binary_sensor.py @@ -1,31 +1,46 @@ """Binary sensors for the Elexa Guardian integration.""" +from typing import Callable, Dict + +from aioguardian import Client + from homeassistant.components.binary_sensor import BinarySensorEntity -from homeassistant.core import callback +from homeassistant.config_entries import ConfigEntry +from homeassistant.core import HomeAssistant, callback +from homeassistant.helpers.update_coordinator import DataUpdateCoordinator from . import GuardianEntity from .const import ( + API_SYSTEM_ONBOARD_SENSOR_STATUS, + API_WIFI_STATUS, DATA_CLIENT, - DATA_SENSOR_STATUS, - DATA_WIFI_STATUS, + DATA_COORDINATOR, DOMAIN, - SENSOR_KIND_AP_INFO, - SENSOR_KIND_LEAK_DETECTED, ) ATTR_CONNECTED_CLIENTS = "connected_clients" +SENSOR_KIND_AP_INFO = "ap_enabled" +SENSOR_KIND_LEAK_DETECTED = "leak_detected" SENSORS = [ (SENSOR_KIND_AP_INFO, "Onboard AP Enabled", "connectivity"), (SENSOR_KIND_LEAK_DETECTED, "Leak Detected", "moisture"), ] -async def async_setup_entry(hass, entry, async_add_entities): +async def async_setup_entry( + hass: HomeAssistant, entry: ConfigEntry, async_add_entities: Callable +) -> None: """Set up Guardian switches based on a config entry.""" - guardian = hass.data[DOMAIN][DATA_CLIENT][entry.entry_id] async_add_entities( [ - GuardianBinarySensor(guardian, kind, name, device_class) + GuardianBinarySensor( + entry, + hass.data[DOMAIN][DATA_CLIENT][entry.entry_id], + hass.data[DOMAIN][DATA_COORDINATOR][entry.entry_id], + kind, + name, + device_class, + ) for kind, name, device_class in SENSORS ], True, @@ -35,28 +50,55 @@ async def async_setup_entry(hass, entry, async_add_entities): class GuardianBinarySensor(GuardianEntity, BinarySensorEntity): """Define a generic Guardian sensor.""" - def __init__(self, guardian, kind, name, device_class): + def __init__( + self, + entry: ConfigEntry, + client: Client, + coordinators: Dict[str, DataUpdateCoordinator], + kind: str, + name: str, + device_class: str, + ) -> None: """Initialize.""" - super().__init__(guardian, kind, name, device_class, None) + super().__init__(entry, client, coordinators, kind, name, device_class, None) self._is_on = True @property - def is_on(self): + def available(self) -> bool: + """Return whether the entity is available.""" + if self._kind == SENSOR_KIND_AP_INFO: + return self._coordinators[API_WIFI_STATUS].last_update_success + if self._kind == SENSOR_KIND_LEAK_DETECTED: + return self._coordinators[ + API_SYSTEM_ONBOARD_SENSOR_STATUS + ].last_update_success + return False + + @property + def is_on(self) -> bool: """Return True if the binary sensor is on.""" return self._is_on + async def _async_internal_added_to_hass(self) -> None: + if self._kind == SENSOR_KIND_AP_INFO: + self.async_add_coordinator_update_listener(API_WIFI_STATUS) + elif self._kind == SENSOR_KIND_LEAK_DETECTED: + self.async_add_coordinator_update_listener(API_SYSTEM_ONBOARD_SENSOR_STATUS) + @callback - def _update_from_latest_data(self): + def _async_update_from_latest_data(self) -> None: """Update the entity.""" if self._kind == SENSOR_KIND_AP_INFO: - self._is_on = self._guardian.data[DATA_WIFI_STATUS]["ap_enabled"] + self._is_on = self._coordinators[API_WIFI_STATUS].data["ap_enabled"] self._attrs.update( { - ATTR_CONNECTED_CLIENTS: self._guardian.data[DATA_WIFI_STATUS][ + ATTR_CONNECTED_CLIENTS: self._coordinators[API_WIFI_STATUS].data[ "ap_clients" ] } ) elif self._kind == SENSOR_KIND_LEAK_DETECTED: - self._is_on = self._guardian.data[DATA_SENSOR_STATUS]["wet"] + self._is_on = self._coordinators[API_SYSTEM_ONBOARD_SENSOR_STATUS].data[ + "wet" + ] diff --git a/homeassistant/components/guardian/const.py b/homeassistant/components/guardian/const.py index f1d60fd07da..321a46a3ffc 100644 --- a/homeassistant/components/guardian/const.py +++ b/homeassistant/components/guardian/const.py @@ -5,21 +5,12 @@ DOMAIN = "guardian" LOGGER = logging.getLogger(__package__) +API_SYSTEM_DIAGNOSTICS = "system_diagnostics" +API_SYSTEM_ONBOARD_SENSOR_STATUS = "system_onboard_sensor_status" +API_VALVE_STATUS = "valve_status" +API_WIFI_STATUS = "wifi_status" + CONF_UID = "uid" DATA_CLIENT = "client" -DATA_DIAGNOSTICS = "diagnostics" -DATA_PAIR_DUMP = "pair_sensor" -DATA_PING = "ping" -DATA_SENSOR_STATUS = "sensor_status" -DATA_VALVE_STATUS = "valve_status" -DATA_WIFI_STATUS = "wifi_status" - -SENSOR_KIND_AP_INFO = "ap_enabled" -SENSOR_KIND_LEAK_DETECTED = "leak_detected" -SENSOR_KIND_TEMPERATURE = "temperature" -SENSOR_KIND_UPTIME = "uptime" - -SWITCH_KIND_VALVE = "valve" - -TOPIC_UPDATE = "guardian_update_{0}" +DATA_COORDINATOR = "coordinator" diff --git a/homeassistant/components/guardian/sensor.py b/homeassistant/components/guardian/sensor.py index 4da200224cf..ed22fc2e73a 100644 --- a/homeassistant/components/guardian/sensor.py +++ b/homeassistant/components/guardian/sensor.py @@ -1,17 +1,24 @@ """Sensors for the Elexa Guardian integration.""" -from homeassistant.const import DEVICE_CLASS_TEMPERATURE, TEMP_FAHRENHEIT, TIME_MINUTES -from homeassistant.core import callback +from typing import Callable, Dict -from . import Guardian, GuardianEntity +from aioguardian import Client + +from homeassistant.config_entries import ConfigEntry +from homeassistant.const import DEVICE_CLASS_TEMPERATURE, TEMP_FAHRENHEIT, TIME_MINUTES +from homeassistant.core import HomeAssistant, callback +from homeassistant.helpers.update_coordinator import DataUpdateCoordinator + +from . import GuardianEntity from .const import ( + API_SYSTEM_DIAGNOSTICS, + API_SYSTEM_ONBOARD_SENSOR_STATUS, DATA_CLIENT, - DATA_DIAGNOSTICS, - DATA_SENSOR_STATUS, + DATA_COORDINATOR, DOMAIN, - SENSOR_KIND_TEMPERATURE, - SENSOR_KIND_UPTIME, ) +SENSOR_KIND_TEMPERATURE = "temperature" +SENSOR_KIND_UPTIME = "uptime" SENSORS = [ ( SENSOR_KIND_TEMPERATURE, @@ -24,12 +31,22 @@ SENSORS = [ ] -async def async_setup_entry(hass, entry, async_add_entities): +async def async_setup_entry( + hass: HomeAssistant, entry: ConfigEntry, async_add_entities: Callable +) -> None: """Set up Guardian switches based on a config entry.""" - guardian = hass.data[DOMAIN][DATA_CLIENT][entry.entry_id] async_add_entities( [ - GuardianSensor(guardian, kind, name, device_class, icon, unit) + GuardianSensor( + entry, + hass.data[DOMAIN][DATA_CLIENT][entry.entry_id], + hass.data[DOMAIN][DATA_COORDINATOR][entry.entry_id], + kind, + name, + device_class, + icon, + unit, + ) for kind, name, device_class, icon, unit in SENSORS ], True, @@ -41,33 +58,53 @@ class GuardianSensor(GuardianEntity): def __init__( self, - guardian: Guardian, + entry: ConfigEntry, + client: Client, + coordinators: Dict[str, DataUpdateCoordinator], kind: str, name: str, device_class: str, icon: str, unit: str, - ): + ) -> None: """Initialize.""" - super().__init__(guardian, kind, name, device_class, icon) + super().__init__(entry, client, coordinators, kind, name, device_class, icon) self._state = None self._unit = unit @property - def state(self): + def available(self) -> bool: + """Return whether the entity is available.""" + if self._kind == SENSOR_KIND_TEMPERATURE: + return self._coordinators[ + API_SYSTEM_ONBOARD_SENSOR_STATUS + ].last_update_success + if self._kind == SENSOR_KIND_UPTIME: + return self._coordinators[API_SYSTEM_DIAGNOSTICS].last_update_success + return False + + @property + def state(self) -> str: """Return the sensor state.""" return self._state @property - def unit_of_measurement(self): + def unit_of_measurement(self) -> str: """Return the unit of measurement of this entity, if any.""" return self._unit + async def _async_internal_added_to_hass(self) -> None: + """Register API interest (and related tasks) when the entity is added.""" + if self._kind == SENSOR_KIND_TEMPERATURE: + self.async_add_coordinator_update_listener(API_SYSTEM_ONBOARD_SENSOR_STATUS) + @callback - def _update_from_latest_data(self): + def _async_update_from_latest_data(self) -> None: """Update the entity.""" if self._kind == SENSOR_KIND_TEMPERATURE: - self._state = self._guardian.data[DATA_SENSOR_STATUS]["temperature"] + self._state = self._coordinators[API_SYSTEM_ONBOARD_SENSOR_STATUS].data[ + "temperature" + ] elif self._kind == SENSOR_KIND_UPTIME: - self._state = self._guardian.data[DATA_DIAGNOSTICS]["uptime"] + self._state = self._coordinators[API_SYSTEM_DIAGNOSTICS].data["uptime"] diff --git a/homeassistant/components/guardian/switch.py b/homeassistant/components/guardian/switch.py index 461eeaaeedb..af42c5b3900 100644 --- a/homeassistant/components/guardian/switch.py +++ b/homeassistant/components/guardian/switch.py @@ -1,4 +1,7 @@ """Switches for the Elexa Guardian integration.""" +from typing import Callable, Dict + +from aioguardian import Client from aioguardian.commands.system import ( DEFAULT_FIRMWARE_UPGRADE_FILENAME, DEFAULT_FIRMWARE_UPGRADE_PORT, @@ -8,12 +11,14 @@ from aioguardian.errors import GuardianError import voluptuous as vol from homeassistant.components.switch import SwitchEntity +from homeassistant.config_entries import ConfigEntry from homeassistant.const import CONF_FILENAME, CONF_PORT, CONF_URL -from homeassistant.core import callback +from homeassistant.core import HomeAssistant, callback from homeassistant.helpers import config_validation as cv, entity_platform +from homeassistant.helpers.update_coordinator import DataUpdateCoordinator -from . import Guardian, GuardianEntity -from .const import DATA_CLIENT, DATA_VALVE_STATUS, DOMAIN, LOGGER, SWITCH_KIND_VALVE +from . import GuardianEntity +from .const import API_VALVE_STATUS, DATA_CLIENT, DATA_COORDINATOR, DOMAIN, LOGGER ATTR_AVG_CURRENT = "average_current" ATTR_INST_CURRENT = "instantaneous_current" @@ -37,10 +42,10 @@ SERVICE_UPGRADE_FIRMWARE_SCHEMA = vol.Schema( ) -async def async_setup_entry(hass, entry, async_add_entities): +async def async_setup_entry( + hass: HomeAssistant, entry: ConfigEntry, async_add_entities: Callable +) -> None: """Set up Guardian switches based on a config entry.""" - guardian = hass.data[DOMAIN][DATA_CLIENT][entry.entry_id] - platform = entity_platform.current_platform.get() for service_name, schema, method in [ @@ -56,27 +61,52 @@ async def async_setup_entry(hass, entry, async_add_entities): ]: platform.async_register_entity_service(service_name, schema, method) - async_add_entities([GuardianSwitch(guardian)], True) + async_add_entities( + [ + GuardianSwitch( + entry, + hass.data[DOMAIN][DATA_CLIENT][entry.entry_id], + hass.data[DOMAIN][DATA_COORDINATOR][entry.entry_id], + ) + ], + True, + ) class GuardianSwitch(GuardianEntity, SwitchEntity): """Define a switch to open/close the Guardian valve.""" - def __init__(self, guardian: Guardian): + def __init__( + self, + entry: ConfigEntry, + client: Client, + coordinators: Dict[str, DataUpdateCoordinator], + ): """Initialize.""" - super().__init__(guardian, SWITCH_KIND_VALVE, "Valve", None, "mdi:water") + super().__init__( + entry, client, coordinators, "valve", "Valve", None, "mdi:water" + ) self._is_on = True @property - def is_on(self): + def available(self) -> bool: + """Return whether the entity is available.""" + return self._coordinators[API_VALVE_STATUS].last_update_success + + @property + def is_on(self) -> bool: """Return True if the valve is open.""" return self._is_on + async def _async_internal_added_to_hass(self): + """Register API interest (and related tasks) when the entity is added.""" + self.async_add_coordinator_update_listener(API_VALVE_STATUS) + @callback - def _update_from_latest_data(self): + def _async_update_from_latest_data(self) -> None: """Update the entity.""" - self._is_on = self._guardian.data[DATA_VALVE_STATUS]["state"] in ( + self._is_on = self._coordinators[API_VALVE_STATUS].data["state"] in ( "start_opening", "opening", "finish_opening", @@ -85,16 +115,16 @@ class GuardianSwitch(GuardianEntity, SwitchEntity): self._attrs.update( { - ATTR_AVG_CURRENT: self._guardian.data[DATA_VALVE_STATUS][ + ATTR_AVG_CURRENT: self._coordinators[API_VALVE_STATUS].data[ "average_current" ], - ATTR_INST_CURRENT: self._guardian.data[DATA_VALVE_STATUS][ + ATTR_INST_CURRENT: self._coordinators[API_VALVE_STATUS].data[ "instantaneous_current" ], - ATTR_INST_CURRENT_DDT: self._guardian.data[DATA_VALVE_STATUS][ + ATTR_INST_CURRENT_DDT: self._coordinators[API_VALVE_STATUS].data[ "instantaneous_current_ddt" ], - ATTR_TRAVEL_COUNT: self._guardian.data[DATA_VALVE_STATUS][ + ATTR_TRAVEL_COUNT: self._coordinators[API_VALVE_STATUS].data[ "travel_count" ], } @@ -103,40 +133,40 @@ class GuardianSwitch(GuardianEntity, SwitchEntity): async def async_disable_ap(self): """Disable the device's onboard access point.""" try: - async with self._guardian.client: - await self._guardian.client.wifi.disable_ap() + async with self._client: + await self._client.wifi.disable_ap() except GuardianError as err: LOGGER.error("Error during service call: %s", err) async def async_enable_ap(self): """Enable the device's onboard access point.""" try: - async with self._guardian.client: - await self._guardian.client.wifi.enable_ap() + async with self._client: + await self._client.wifi.enable_ap() except GuardianError as err: LOGGER.error("Error during service call: %s", err) async def async_reboot(self): """Reboot the device.""" try: - async with self._guardian.client: - await self._guardian.client.system.reboot() + async with self._client: + await self._client.system.reboot() except GuardianError as err: LOGGER.error("Error during service call: %s", err) async def async_reset_valve_diagnostics(self): """Fully reset system motor diagnostics.""" try: - async with self._guardian.client: - await self._guardian.client.valve.reset() + async with self._client: + await self._client.valve.reset() except GuardianError as err: LOGGER.error("Error during service call: %s", err) async def async_upgrade_firmware(self, *, url, port, filename): """Upgrade the device firmware.""" try: - async with self._guardian.client: - await self._guardian.client.system.upgrade_firmware( + async with self._client: + await self._client.system.upgrade_firmware( url=url, port=port, filename=filename, ) except GuardianError as err: @@ -145,8 +175,8 @@ class GuardianSwitch(GuardianEntity, SwitchEntity): async def async_turn_off(self, **kwargs) -> None: """Turn the valve off (closed).""" try: - async with self._guardian.client: - await self._guardian.client.valve.close() + async with self._client: + await self._client.valve.close() except GuardianError as err: LOGGER.error("Error while closing the valve: %s", err) return @@ -157,8 +187,8 @@ class GuardianSwitch(GuardianEntity, SwitchEntity): async def async_turn_on(self, **kwargs) -> None: """Turn the valve on (open).""" try: - async with self._guardian.client: - await self._guardian.client.valve.open() + async with self._client: + await self._client.valve.open() except GuardianError as err: LOGGER.error("Error while opening the valve: %s", err) return diff --git a/homeassistant/components/guardian/util.py b/homeassistant/components/guardian/util.py new file mode 100644 index 00000000000..e5fe565bbf4 --- /dev/null +++ b/homeassistant/components/guardian/util.py @@ -0,0 +1,49 @@ +"""Define Guardian-specific utilities.""" +import asyncio +from datetime import timedelta +from typing import Awaitable, Callable + +from aioguardian import Client +from aioguardian.errors import GuardianError + +from homeassistant.core import HomeAssistant +from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed + +from .const import LOGGER + +DEFAULT_UPDATE_INTERVAL = timedelta(seconds=30) + + +class GuardianDataUpdateCoordinator(DataUpdateCoordinator): + """Define an extended DataUpdateCoordinator with some Guardian goodies.""" + + def __init__( + self, + hass: HomeAssistant, + *, + client: Client, + api_name: str, + api_coro: Callable[..., Awaitable], + api_lock: asyncio.Lock, + valve_controller_uid: str, + ): + """Initialize.""" + super().__init__( + hass, + LOGGER, + name=f"{valve_controller_uid}_{api_name}", + update_interval=DEFAULT_UPDATE_INTERVAL, + ) + + self._api_coro = api_coro + self._api_lock = api_lock + self._client = client + + async def _async_update_data(self) -> dict: + """Execute a "locked" API request against the valve controller.""" + async with self._api_lock, self._client: + try: + resp = await self._api_coro() + except GuardianError as err: + raise UpdateFailed(err) + return resp["data"]