Add smarter API usage for RainMachine (#31115)

* Make RainMachine smarter with API usage

* Remove debug statements

* Fix deregistration

* Code review comments

* Code review

* Use an asyncio.Lock

* Remove unnecessary guard clause

* Ensure registation lock per API category
This commit is contained in:
Aaron Bach 2020-01-24 22:31:14 -07:00 committed by Paulus Schoutsen
parent cf165cc35f
commit 550aa6a0a5
3 changed files with 142 additions and 50 deletions

View file

@ -16,6 +16,7 @@ from homeassistant.const import (
CONF_SCAN_INTERVAL,
CONF_SSL,
)
from homeassistant.core import callback
from homeassistant.exceptions import ConfigEntryNotReady
from homeassistant.helpers import aiohttp_client, config_validation as cv
from homeassistant.helpers.dispatcher import async_dispatcher_send
@ -127,7 +128,6 @@ async def async_setup(hass, config):
async def async_setup_entry(hass, config_entry):
"""Set up RainMachine as config entry."""
_verify_domain_control = verify_domain_control(hass, DOMAIN)
websession = aiohttp_client.async_get_clientsession(hass)
@ -141,9 +141,11 @@ async def async_setup_entry(hass, config_entry):
ssl=config_entry.data[CONF_SSL],
)
rainmachine = RainMachine(
client, config_entry.data.get(CONF_ZONE_RUN_TIME, DEFAULT_ZONE_RUN),
hass,
client,
config_entry.data.get(CONF_ZONE_RUN_TIME, DEFAULT_ZONE_RUN),
config_entry.data[CONF_SCAN_INTERVAL],
)
await rainmachine.async_update()
except RainMachineError as err:
_LOGGER.error("An error occurred: %s", err)
raise ConfigEntryNotReady
@ -155,16 +157,6 @@ async def async_setup_entry(hass, config_entry):
hass.config_entries.async_forward_entry_setup(config_entry, component)
)
async def refresh(event_time):
"""Refresh RainMachine sensor data."""
_LOGGER.debug("Updating RainMachine sensor data")
await rainmachine.async_update()
async_dispatcher_send(hass, SENSOR_UPDATE_TOPIC)
hass.data[DOMAIN][DATA_LISTENER][config_entry.entry_id] = async_track_time_interval(
hass, refresh, timedelta(seconds=config_entry.data[CONF_SCAN_INTERVAL])
)
@_verify_domain_control
async def disable_program(call):
"""Disable a program."""
@ -271,30 +263,86 @@ async def async_unload_entry(hass, config_entry):
class RainMachine:
"""Define a generic RainMachine object."""
def __init__(self, client, default_zone_runtime):
def __init__(self, hass, client, default_zone_runtime, scan_interval):
"""Initialize."""
self._async_unsub_dispatcher_connect = None
self._scan_interval_seconds = scan_interval
self.client = client
self.data = {}
self.default_zone_runtime = default_zone_runtime
self.device_mac = self.client.mac
self.hass = hass
self._api_category_count = {
PROVISION_SETTINGS: 0,
RESTRICTIONS_CURRENT: 0,
RESTRICTIONS_UNIVERSAL: 0,
}
self._api_category_locks = {
PROVISION_SETTINGS: asyncio.Lock(),
RESTRICTIONS_CURRENT: asyncio.Lock(),
RESTRICTIONS_UNIVERSAL: asyncio.Lock(),
}
async def _async_fetch_from_api(self, api_category):
"""Execute the appropriate coroutine to fetch particular data from the API."""
if api_category == PROVISION_SETTINGS:
data = await self.client.provisioning.settings()
elif api_category == RESTRICTIONS_CURRENT:
data = await self.client.restrictions.current()
elif api_category == RESTRICTIONS_UNIVERSAL:
data = await self.client.restrictions.universal()
return data
@callback
def async_deregister_api_interest(self, api_category):
"""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_unsub_dispatcher_connect:
self._async_unsub_dispatcher_connect()
self._async_unsub_dispatcher_connect = None
return
self._api_category_count[api_category] += 1
async def async_register_api_interest(self, api_category):
"""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_unsub_dispatcher_connect:
self._async_unsub_dispatcher_connect = async_track_time_interval(
self.hass,
self.async_update,
timedelta(seconds=self._scan_interval_seconds),
)
self._api_category_count[api_category] += 1
# Lock API updates in case multiple entities are trying to call the same API
# endpoint at once:
async with self._api_category_locks[api_category]:
if api_category not in self.data:
self.data[api_category] = await self._async_fetch_from_api(api_category)
async def async_update(self):
"""Update sensor/binary sensor data."""
tasks = {
PROVISION_SETTINGS: self.client.provisioning.settings(),
RESTRICTIONS_CURRENT: self.client.restrictions.current(),
RESTRICTIONS_UNIVERSAL: self.client.restrictions.universal(),
}
tasks = {}
for category, count in self._api_category_count.items():
if count == 0:
continue
tasks[category] = self._async_fetch_from_api(category)
results = await asyncio.gather(*tasks.values(), return_exceptions=True)
for operation, result in zip(tasks, results):
for api_category, result in zip(tasks, results):
if isinstance(result, RainMachineError):
_LOGGER.error(
"There was an error while updating %s: %s", operation, result
"There was an error while updating %s: %s", api_category, result
)
continue
self.data[api_category] = result
self.data[operation] = result
async_dispatcher_send(self.hass, SENSOR_UPDATE_TOPIC)
class RainMachineEntity(Entity):

View file

@ -28,40 +28,64 @@ TYPE_RAINSENSOR = "rainsensor"
TYPE_WEEKDAY = "weekday"
BINARY_SENSORS = {
TYPE_FLOW_SENSOR: ("Flow Sensor", "mdi:water-pump", True),
TYPE_FREEZE: ("Freeze Restrictions", "mdi:cancel", True),
TYPE_FREEZE_PROTECTION: ("Freeze Protection", "mdi:weather-snowy", True),
TYPE_HOT_DAYS: ("Extra Water on Hot Days", "mdi:thermometer-lines", True),
TYPE_HOURLY: ("Hourly Restrictions", "mdi:cancel", False),
TYPE_MONTH: ("Month Restrictions", "mdi:cancel", False),
TYPE_RAINDELAY: ("Rain Delay Restrictions", "mdi:cancel", False),
TYPE_RAINSENSOR: ("Rain Sensor Restrictions", "mdi:cancel", False),
TYPE_WEEKDAY: ("Weekday Restrictions", "mdi:cancel", False),
TYPE_FLOW_SENSOR: ("Flow Sensor", "mdi:water-pump", True, PROVISION_SETTINGS),
TYPE_FREEZE: ("Freeze Restrictions", "mdi:cancel", True, RESTRICTIONS_CURRENT),
TYPE_FREEZE_PROTECTION: (
"Freeze Protection",
"mdi:weather-snowy",
True,
RESTRICTIONS_UNIVERSAL,
),
TYPE_HOT_DAYS: (
"Extra Water on Hot Days",
"mdi:thermometer-lines",
True,
RESTRICTIONS_UNIVERSAL,
),
TYPE_HOURLY: ("Hourly Restrictions", "mdi:cancel", False, RESTRICTIONS_CURRENT),
TYPE_MONTH: ("Month Restrictions", "mdi:cancel", False, RESTRICTIONS_CURRENT),
TYPE_RAINDELAY: (
"Rain Delay Restrictions",
"mdi:cancel",
False,
RESTRICTIONS_CURRENT,
),
TYPE_RAINSENSOR: (
"Rain Sensor Restrictions",
"mdi:cancel",
False,
RESTRICTIONS_CURRENT,
),
TYPE_WEEKDAY: ("Weekday Restrictions", "mdi:cancel", False, RESTRICTIONS_CURRENT),
}
async def async_setup_entry(hass, entry, async_add_entities):
"""Set up RainMachine binary sensors based on a config entry."""
rainmachine = hass.data[RAINMACHINE_DOMAIN][DATA_CLIENT][entry.entry_id]
binary_sensors = []
for sensor_type, (name, icon, enabled_by_default) in BINARY_SENSORS.items():
binary_sensors.append(
async_add_entities(
[
RainMachineBinarySensor(
rainmachine, sensor_type, name, icon, enabled_by_default
rainmachine, sensor_type, name, icon, enabled_by_default, api_category
)
)
async_add_entities(binary_sensors, True)
for (
sensor_type,
(name, icon, enabled_by_default, api_category),
) in BINARY_SENSORS.items()
],
)
class RainMachineBinarySensor(RainMachineEntity, BinarySensorDevice):
"""A sensor implementation for raincloud device."""
def __init__(self, rainmachine, sensor_type, name, icon, enabled_by_default):
def __init__(
self, rainmachine, sensor_type, name, icon, enabled_by_default, api_category
):
"""Initialize the sensor."""
super().__init__(rainmachine)
self._api_category = api_category
self._enabled_by_default = enabled_by_default
self._icon = icon
self._name = name
@ -106,6 +130,8 @@ class RainMachineBinarySensor(RainMachineEntity, BinarySensorDevice):
self._dispatcher_handlers.append(
async_dispatcher_connect(self.hass, SENSOR_UPDATE_TOPIC, update)
)
await self.rainmachine.async_register_api_interest(self._api_category)
await self.async_update()
async def async_update(self):
"""Update the state."""
@ -133,3 +159,8 @@ class RainMachineBinarySensor(RainMachineEntity, BinarySensorDevice):
self._state = self.rainmachine.data[RESTRICTIONS_CURRENT]["rainSensor"]
elif self._sensor_type == TYPE_WEEKDAY:
self._state = self.rainmachine.data[RESTRICTIONS_CURRENT]["weekDay"]
async def async_will_remove_from_hass(self):
"""Disconnect dispatcher listeners and deregister API interest."""
super().async_will_remove_from_hass()
self.rainmachine.async_deregister_api_interest(self._api_category)

View file

@ -28,6 +28,7 @@ SENSORS = {
"clicks/m^3",
None,
False,
PROVISION_SETTINGS,
),
TYPE_FLOW_SENSOR_CONSUMED_LITERS: (
"Flow Sensor Consumed Liters",
@ -35,6 +36,7 @@ SENSORS = {
"liter",
None,
False,
PROVISION_SETTINGS,
),
TYPE_FLOW_SENSOR_START_INDEX: (
"Flow Sensor Start Index",
@ -42,6 +44,7 @@ SENSORS = {
"index",
None,
False,
PROVISION_SETTINGS,
),
TYPE_FLOW_SENSOR_WATERING_CLICKS: (
"Flow Sensor Clicks",
@ -49,6 +52,7 @@ SENSORS = {
"clicks",
None,
False,
PROVISION_SETTINGS,
),
TYPE_FREEZE_TEMP: (
"Freeze Protect Temperature",
@ -56,6 +60,7 @@ SENSORS = {
"°C",
"temperature",
True,
RESTRICTIONS_UNIVERSAL,
),
}
@ -63,13 +68,8 @@ SENSORS = {
async def async_setup_entry(hass, entry, async_add_entities):
"""Set up RainMachine sensors based on a config entry."""
rainmachine = hass.data[RAINMACHINE_DOMAIN][DATA_CLIENT][entry.entry_id]
sensors = []
for (
sensor_type,
(name, icon, unit, device_class, enabled_by_default),
) in SENSORS.items():
sensors.append(
async_add_entities(
[
RainMachineSensor(
rainmachine,
sensor_type,
@ -78,10 +78,14 @@ async def async_setup_entry(hass, entry, async_add_entities):
unit,
device_class,
enabled_by_default,
api_category,
)
)
async_add_entities(sensors, True)
for (
sensor_type,
(name, icon, unit, device_class, enabled_by_default, api_category),
) in SENSORS.items()
],
)
class RainMachineSensor(RainMachineEntity):
@ -96,10 +100,12 @@ class RainMachineSensor(RainMachineEntity):
unit,
device_class,
enabled_by_default,
api_category,
):
"""Initialize."""
super().__init__(rainmachine)
self._api_category = api_category
self._device_class = device_class
self._enabled_by_default = enabled_by_default
self._icon = icon
@ -151,6 +157,8 @@ class RainMachineSensor(RainMachineEntity):
self._dispatcher_handlers.append(
async_dispatcher_connect(self.hass, SENSOR_UPDATE_TOPIC, update)
)
await self.rainmachine.async_register_api_interest(self._api_category)
await self.async_update()
async def async_update(self):
"""Update the sensor's state."""
@ -182,3 +190,8 @@ class RainMachineSensor(RainMachineEntity):
self._state = self.rainmachine.data[RESTRICTIONS_UNIVERSAL][
"freezeProtectTemp"
]
async def async_will_remove_from_hass(self):
"""Disconnect dispatcher listeners and deregister API interest."""
super().async_will_remove_from_hass()
self.rainmachine.async_deregister_api_interest(self._api_category)