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

@@ -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

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_COORDINATOR,
-    DATA_PAIR_DUMP,
-    DATA_PING,
-    DATA_SENSOR_STATUS,
-    DATA_VALVE_STATUS,
-    DATA_WIFI_STATUS,
     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 = {
+DEFAULT_UPDATE_INTERVAL = timedelta(seconds=30)
-    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)
 
 PLATFORMS = ["binary_sensor", "sensor", "switch"]
 
 
-@callback
+async def async_setup(hass: HomeAssistant, config: dict) -> bool:
-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):
     """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)
+    client = hass.data[DOMAIN][DATA_CLIENT][entry.entry_id] = Client(
-    await guardian.async_update()
+        entry.data[CONF_IP_ADDRESS], port=entry.data[CONF_PORT]
-    hass.data[DOMAIN][DATA_CLIENT][entry.entry_id] = guardian
+    )
+    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):
+    def device_class(self) -> str:
-        """Return whether the entity is available."""
-        return bool(self._guardian.data[DATA_PING])
-
-    @property
-    def device_class(self):
         """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"],
+            "model": self._coordinators[API_SYSTEM_DIAGNOSTICS].data["firmware"],
-            "name": f"Guardian {self._guardian.uid}",
+            "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
+    async def _async_internal_added_to_hass(self):
-    def _update_from_latest_data(self):
+        """Perform additional, internal tasks when the entity is about to be added.
-        """Update the entity."""
+
+        This should be extended by Guardian platforms.
+        """
         raise NotImplementedError
 
-    async def async_added_to_hass(self):
+    @callback
-        """Register callbacks."""
+    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():
+        def async_update():
-            """Update the state."""
+            """Update the entity's state."""
-            self._update_from_latest_data()
+            self._async_update_from_latest_data()
             self.async_write_ha_state()
 
-        self.async_on_remove(
+        self.async_on_remove(self._coordinators[api].async_add_listener(async_update))
-            async_dispatcher_connect(
-                self.hass, TOPIC_UPDATE.format(self._guardian.uid), update
-            )
-        )
 
-        await self._guardian.async_register_api_interest(self._kind)
+    async def async_added_to_hass(self) -> None:
-
+        """Perform tasks when the entity is added."""
-        self._update_from_latest_data()
+        await self._async_internal_added_to_hass()
-
+        self.async_add_coordinator_update_listener(API_SYSTEM_DIAGNOSTICS)
-    async def async_will_remove_from_hass(self) -> None:
+        self._async_update_from_latest_data()
-        """Disconnect dispatcher listener when removed."""
-        self._guardian.async_deregister_api_interest(self._kind)

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_COORDINATOR,
-    DATA_WIFI_STATUS,
     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"
+            ]

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_COORDINATOR = "coordinator"
-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}"

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 typing import Callable, Dict
-from homeassistant.core import callback
 
-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_COORDINATOR,
-    DATA_SENSOR_STATUS,
     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"]

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 . import GuardianEntity
-from .const import DATA_CLIENT, DATA_VALVE_STATUS, DOMAIN, LOGGER, SWITCH_KIND_VALVE
+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:
+            async with self._client:
-                await self._guardian.client.wifi.disable_ap()
+                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:
+            async with self._client:
-                await self._guardian.client.wifi.enable_ap()
+                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:
+            async with self._client:
-                await self._guardian.client.system.reboot()
+                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:
+            async with self._client:
-                await self._guardian.client.valve.reset()
+                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:
+            async with self._client:
-                await self._guardian.client.system.upgrade_firmware(
+                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:
+            async with self._client:
-                await self._guardian.client.valve.close()
+                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:
+            async with self._client:
-                await self._guardian.client.valve.open()
+                await self._client.valve.open()
         except GuardianError as err:
             LOGGER.error("Error while opening the valve: %s", err)
             return

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"]