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

homeassistant/components/rainmachine/__init__.py

@@ -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 = {
+        tasks = {}
-            PROVISION_SETTINGS: self.client.provisioning.settings(),
+        for category, count in self._api_category_count.items():
-            RESTRICTIONS_CURRENT: self.client.restrictions.current(),
+            if count == 0:
-            RESTRICTIONS_UNIVERSAL: self.client.restrictions.universal(),
+                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):

homeassistant/components/rainmachine/binary_sensor.py

@@ -28,40 +28,64 @@ TYPE_RAINSENSOR = "rainsensor"
 TYPE_WEEKDAY = "weekday"
 
 BINARY_SENSORS = {
-    TYPE_FLOW_SENSOR: ("Flow Sensor", "mdi:water-pump", True),
+    TYPE_FLOW_SENSOR: ("Flow Sensor", "mdi:water-pump", True, PROVISION_SETTINGS),
-    TYPE_FREEZE: ("Freeze Restrictions", "mdi:cancel", True),
+    TYPE_FREEZE: ("Freeze Restrictions", "mdi:cancel", True, RESTRICTIONS_CURRENT),
-    TYPE_FREEZE_PROTECTION: ("Freeze Protection", "mdi:weather-snowy", True),
+    TYPE_FREEZE_PROTECTION: (
-    TYPE_HOT_DAYS: ("Extra Water on Hot Days", "mdi:thermometer-lines", True),
+        "Freeze Protection",
-    TYPE_HOURLY: ("Hourly Restrictions", "mdi:cancel", False),
+        "mdi:weather-snowy",
-    TYPE_MONTH: ("Month Restrictions", "mdi:cancel", False),
+        True,
-    TYPE_RAINDELAY: ("Rain Delay Restrictions", "mdi:cancel", False),
+        RESTRICTIONS_UNIVERSAL,
-    TYPE_RAINSENSOR: ("Rain Sensor Restrictions", "mdi:cancel", False),
+    ),
-    TYPE_WEEKDAY: ("Weekday Restrictions", "mdi:cancel", False),
+    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]
-
+    async_add_entities(
-    binary_sensors = []
+        [
-    for sensor_type, (name, icon, enabled_by_default) in BINARY_SENSORS.items():
-        binary_sensors.append(
             RainMachineBinarySensor(
-                rainmachine, sensor_type, name, icon, enabled_by_default
+                rainmachine, sensor_type, name, icon, enabled_by_default, api_category
             )
-        )
+            for (
-
+                sensor_type,
-    async_add_entities(binary_sensors, True)
+                (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)

homeassistant/components/rainmachine/sensor.py

@@ -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]
-
+    async_add_entities(
-    sensors = []
+        [
-    for (
-        sensor_type,
-        (name, icon, unit, device_class, enabled_by_default),
-    ) in SENSORS.items():
-        sensors.append(
             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,
             )
-        )
+            for (
-
+                sensor_type,
-    async_add_entities(sensors, True)
+                (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)