Update AirVisual to use DataUpdateCoordinator (#34796)

* Update AirVisual to use DataUpdateCoordinator

* Empty commit to re-trigger build

* Don't include history or trends in config flow

* Code review

homeassistant/components/airvisual/__init__.py

@@ -19,30 +19,23 @@
 )
 from homeassistant.core import callback
 from homeassistant.helpers import aiohttp_client, config_validation as cv
-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, UpdateFailed
 
 from .const import (
     CONF_CITY,
     CONF_COUNTRY,
     CONF_GEOGRAPHIES,
     CONF_INTEGRATION_TYPE,
-    DATA_CLIENT,
+    DATA_COORDINATOR,
     DOMAIN,
     INTEGRATION_TYPE_GEOGRAPHY,
     INTEGRATION_TYPE_NODE_PRO,
     LOGGER,
-    TOPIC_UPDATE,
 )
 
 PLATFORMS = ["air_quality", "sensor"]
 
-DATA_LISTENER = "listener"
-
 DEFAULT_ATTRIBUTION = "Data provided by AirVisual"
 DEFAULT_GEOGRAPHY_SCAN_INTERVAL = timedelta(minutes=10)
 DEFAULT_NODE_PRO_SCAN_INTERVAL = timedelta(minutes=1)
@@ -97,7 +90,7 @@ def async_get_geography_id(geography_dict):
 
 async def async_setup(hass, config):
     """Set up the AirVisual component."""
-    hass.data[DOMAIN] = {DATA_CLIENT: {}, DATA_LISTENER: {}}
+    hass.data[DOMAIN] = {DATA_COORDINATOR: {}}
 
     if DOMAIN not in config:
         return True
@@ -167,35 +160,71 @@ async def async_setup_entry(hass, config_entry):
 
     if CONF_API_KEY in config_entry.data:
         _standardize_geography_config_entry(hass, config_entry)
-        airvisual = AirVisualGeographyData(
+
+        client = Client(api_key=config_entry.data[CONF_API_KEY], session=websession)
+
+        async def async_update_data():
+            """Get new data from the API."""
+            if CONF_CITY in config_entry.data:
+                api_coro = client.api.city(
+                    config_entry.data[CONF_CITY],
+                    config_entry.data[CONF_STATE],
+                    config_entry.data[CONF_COUNTRY],
+                )
+            else:
+                api_coro = client.api.nearest_city(
+                    config_entry.data[CONF_LATITUDE], config_entry.data[CONF_LONGITUDE],
+                )
+
+            try:
+                return await api_coro
+            except AirVisualError as err:
+                raise UpdateFailed(f"Error while retrieving data: {err}")
+
+        coordinator = DataUpdateCoordinator(
             hass,
-            Client(api_key=config_entry.data[CONF_API_KEY], session=websession),
-            config_entry,
+            LOGGER,
+            name="geography data",
+            update_interval=DEFAULT_GEOGRAPHY_SCAN_INTERVAL,
+            update_method=async_update_data,
         )
 
         # Only geography-based entries have options:
         config_entry.add_update_listener(async_update_options)
     else:
         _standardize_node_pro_config_entry(hass, config_entry)
-        airvisual = AirVisualNodeProData(hass, Client(session=websession), config_entry)
 
-    await airvisual.async_update()
+        client = Client(session=websession)
 
-    hass.data[DOMAIN][DATA_CLIENT][config_entry.entry_id] = airvisual
+        async def async_update_data():
+            """Get new data from the API."""
+            try:
+                return await client.node.from_samba(
+                    config_entry.data[CONF_IP_ADDRESS],
+                    config_entry.data[CONF_PASSWORD],
+                    include_history=False,
+                    include_trends=False,
+                )
+            except NodeProError as err:
+                raise UpdateFailed(f"Error while retrieving data: {err}")
+
+        coordinator = DataUpdateCoordinator(
+            hass,
+            LOGGER,
+            name="Node/Pro data",
+            update_interval=DEFAULT_NODE_PRO_SCAN_INTERVAL,
+            update_method=async_update_data,
+        )
+
+    await coordinator.async_refresh()
+
+    hass.data[DOMAIN][DATA_COORDINATOR][config_entry.entry_id] = coordinator
 
     for component in PLATFORMS:
         hass.async_create_task(
             hass.config_entries.async_forward_entry_setup(config_entry, component)
         )
 
-    async def refresh(event_time):
-        """Refresh data from AirVisual."""
-        await airvisual.async_update()
-
-    hass.data[DOMAIN][DATA_LISTENER][config_entry.entry_id] = async_track_time_interval(
-        hass, refresh, airvisual.scan_interval
-    )
-
     return True
 
 
@@ -248,28 +277,31 @@ async def async_unload_entry(hass, config_entry):
         )
     )
     if unload_ok:
-        hass.data[DOMAIN][DATA_CLIENT].pop(config_entry.entry_id)
-        remove_listener = hass.data[DOMAIN][DATA_LISTENER].pop(config_entry.entry_id)
-        remove_listener()
+        hass.data[DOMAIN][DATA_COORDINATOR].pop(config_entry.entry_id)
 
     return unload_ok
 
 
 async def async_update_options(hass, config_entry):
     """Handle an options update."""
-    airvisual = hass.data[DOMAIN][DATA_CLIENT][config_entry.entry_id]
-    airvisual.async_update_options(config_entry.options)
+    coordinator = hass.data[DOMAIN][DATA_COORDINATOR][config_entry.entry_id]
+    await coordinator.async_request_refresh()
 
 
 class AirVisualEntity(Entity):
     """Define a generic AirVisual entity."""
 
-    def __init__(self, airvisual):
+    def __init__(self, coordinator):
         """Initialize."""
-        self._airvisual = airvisual
         self._attrs = {ATTR_ATTRIBUTION: DEFAULT_ATTRIBUTION}
         self._icon = None
         self._unit = None
+        self.coordinator = coordinator
+
+    @property
+    def available(self):
+        """Return if entity is available."""
+        return self.coordinator.last_update_success
 
     @property
     def device_state_attributes(self):
@@ -295,86 +327,11 @@ def update():
             self.update_from_latest_data()
             self.async_write_ha_state()
 
-        self.async_on_remove(
-            async_dispatcher_connect(self.hass, self._airvisual.topic_update, update)
-        )
+        self.async_on_remove(self.coordinator.async_add_listener(update))
 
         self.update_from_latest_data()
 
     @callback
     def update_from_latest_data(self):
         """Update the entity from the latest data."""
         raise NotImplementedError
-
-
-class AirVisualGeographyData:
-    """Define a class to manage data from the AirVisual cloud API."""
-
-    def __init__(self, hass, client, config_entry):
-        """Initialize."""
-        self._client = client
-        self._hass = hass
-        self.data = {}
-        self.geography_data = config_entry.data
-        self.geography_id = config_entry.unique_id
-        self.integration_type = INTEGRATION_TYPE_GEOGRAPHY
-        self.options = config_entry.options
-        self.scan_interval = DEFAULT_GEOGRAPHY_SCAN_INTERVAL
-        self.topic_update = TOPIC_UPDATE.format(config_entry.unique_id)
-
-    async def async_update(self):
-        """Get new data for all locations from the AirVisual cloud API."""
-        if CONF_CITY in self.geography_data:
-            api_coro = self._client.api.city(
-                self.geography_data[CONF_CITY],
-                self.geography_data[CONF_STATE],
-                self.geography_data[CONF_COUNTRY],
-            )
-        else:
-            api_coro = self._client.api.nearest_city(
-                self.geography_data[CONF_LATITUDE], self.geography_data[CONF_LONGITUDE],
-            )
-
-        try:
-            self.data[self.geography_id] = await api_coro
-        except AirVisualError as err:
-            LOGGER.error("Error while retrieving data: %s", err)
-            self.data[self.geography_id] = {}
-
-        LOGGER.debug("Received new geography data")
-        async_dispatcher_send(self._hass, self.topic_update)
-
-    @callback
-    def async_update_options(self, options):
-        """Update the data manager's options."""
-        self.options = options
-        async_dispatcher_send(self._hass, self.topic_update)
-
-
-class AirVisualNodeProData:
-    """Define a class to manage data from an AirVisual Node/Pro."""
-
-    def __init__(self, hass, client, config_entry):
-        """Initialize."""
-        self._client = client
-        self._hass = hass
-        self._password = config_entry.data[CONF_PASSWORD]
-        self.data = {}
-        self.integration_type = INTEGRATION_TYPE_NODE_PRO
-        self.ip_address = config_entry.data[CONF_IP_ADDRESS]
-        self.scan_interval = DEFAULT_NODE_PRO_SCAN_INTERVAL
-        self.topic_update = TOPIC_UPDATE.format(config_entry.data[CONF_IP_ADDRESS])
-
-    async def async_update(self):
-        """Get new data from the Node/Pro."""
-        try:
-            self.data = await self._client.node.from_samba(
-                self.ip_address, self._password, include_history=False
-            )
-        except NodeProError as err:
-            LOGGER.error("Error while retrieving Node/Pro data: %s", err)
-            self.data = {}
-            return
-
-        LOGGER.debug("Received new Node/Pro data")
-        async_dispatcher_send(self._hass, self.topic_update)

homeassistant/components/airvisual/air_quality.py

@@ -4,7 +4,12 @@
 from homeassistant.core import callback
 
 from . import AirVisualEntity
-from .const import DATA_CLIENT, DOMAIN, INTEGRATION_TYPE_GEOGRAPHY
+from .const import (
+    CONF_INTEGRATION_TYPE,
+    DATA_COORDINATOR,
+    DOMAIN,
+    INTEGRATION_TYPE_GEOGRAPHY,
+)
 
 ATTR_HUMIDITY = "humidity"
 ATTR_SENSOR_LIFE = "{0}_sensor_life"
@@ -13,13 +18,13 @@
 
 async def async_setup_entry(hass, config_entry, async_add_entities):
     """Set up AirVisual air quality entities based on a config entry."""
-    airvisual = hass.data[DOMAIN][DATA_CLIENT][config_entry.entry_id]
+    coordinator = hass.data[DOMAIN][DATA_COORDINATOR][config_entry.entry_id]
 
     # Geography-based AirVisual integrations don't utilize this platform:
-    if airvisual.integration_type == INTEGRATION_TYPE_GEOGRAPHY:
+    if config_entry.data[CONF_INTEGRATION_TYPE] == INTEGRATION_TYPE_GEOGRAPHY:
         return
 
-    async_add_entities([AirVisualNodeProSensor(airvisual)], True)
+    async_add_entities([AirVisualNodeProSensor(coordinator)], True)
 
 
 class AirVisualNodeProSensor(AirVisualEntity, AirQualityEntity):
@@ -35,69 +40,71 @@ def __init__(self, airvisual):
     @property
     def air_quality_index(self):
         """Return the Air Quality Index (AQI)."""
-        if self._airvisual.data["current"]["settings"]["is_aqi_usa"]:
-            return self._airvisual.data["current"]["measurements"]["aqi_us"]
-        return self._airvisual.data["current"]["measurements"]["aqi_cn"]
+        if self.coordinator.data["current"]["settings"]["is_aqi_usa"]:
+            return self.coordinator.data["current"]["measurements"]["aqi_us"]
+        return self.coordinator.data["current"]["measurements"]["aqi_cn"]
 
     @property
     def available(self):
         """Return True if entity is available."""
-        return bool(self._airvisual.data)
+        return bool(self.coordinator.data)
 
     @property
     def carbon_dioxide(self):
         """Return the CO2 (carbon dioxide) level."""
-        return self._airvisual.data["current"]["measurements"].get("co2_ppm")
+        return self.coordinator.data["current"]["measurements"].get("co2")
 
     @property
     def device_info(self):
         """Return device registry information for this entity."""
         return {
-            "identifiers": {(DOMAIN, self._airvisual.data["current"]["serial_number"])},
-            "name": self._airvisual.data["current"]["settings"]["node_name"],
+            "identifiers": {
+                (DOMAIN, self.coordinator.data["current"]["serial_number"])
+            },
+            "name": self.coordinator.data["current"]["settings"]["node_name"],
             "manufacturer": "AirVisual",
-            "model": f'{self._airvisual.data["current"]["status"]["model"]}',
+            "model": f'{self.coordinator.data["current"]["status"]["model"]}',
             "sw_version": (
-                f'Version {self._airvisual.data["current"]["status"]["system_version"]}'
-                f'{self._airvisual.data["current"]["status"]["app_version"]}'
+                f'Version {self.coordinator.data["current"]["status"]["system_version"]}'
+                f'{self.coordinator.data["current"]["status"]["app_version"]}'
             ),
         }
 
     @property
     def name(self):
         """Return the name."""
-        node_name = self._airvisual.data["current"]["settings"]["node_name"]
+        node_name = self.coordinator.data["current"]["settings"]["node_name"]
         return f"{node_name} Node/Pro: Air Quality"
 
     @property
     def particulate_matter_2_5(self):
         """Return the particulate matter 2.5 level."""
-        return self._airvisual.data["current"]["measurements"].get("pm2_5")
+        return self.coordinator.data["current"]["measurements"].get("pm2_5")
 
     @property
     def particulate_matter_10(self):
         """Return the particulate matter 10 level."""
-        return self._airvisual.data["current"]["measurements"].get("pm1_0")
+        return self.coordinator.data["current"]["measurements"].get("pm1_0")
 
     @property
     def particulate_matter_0_1(self):
         """Return the particulate matter 0.1 level."""
-        return self._airvisual.data["current"]["measurements"].get("pm0_1")
+        return self.coordinator.data["current"]["measurements"].get("pm0_1")
 
     @property
     def unique_id(self):
         """Return a unique, Home Assistant friendly identifier for this entity."""
-        return self._airvisual.data["current"]["serial_number"]
+        return self.coordinator.data["current"]["serial_number"]
 
     @callback
     def update_from_latest_data(self):
-        """Update from the Node/Pro's data."""
+        """Update the entity from the latest data."""
         self._attrs.update(
             {
-                ATTR_VOC: self._airvisual.data["current"]["measurements"].get("voc"),
+                ATTR_VOC: self.coordinator.data["current"]["measurements"].get("voc"),
                 **{
                     ATTR_SENSOR_LIFE.format(pollutant): lifespan
-                    for pollutant, lifespan in self._airvisual.data["current"][
+                    for pollutant, lifespan in self.coordinator.data["current"][
                         "status"
                     ]["sensor_life"].items()
                 },

homeassistant/components/airvisual/config_flow.py

@@ -146,7 +146,10 @@ async def async_step_node_pro(self, user_input=None):
 
         try:
             await client.node.from_samba(
-                user_input[CONF_IP_ADDRESS], user_input[CONF_PASSWORD]
+                user_input[CONF_IP_ADDRESS],
+                user_input[CONF_PASSWORD],
+                include_history=False,
+                include_trends=False,
             )
         except NodeProError as err:
             LOGGER.error("Error connecting to Node/Pro unit: %s", err)

homeassistant/components/airvisual/const.py

@@ -12,6 +12,4 @@
 CONF_GEOGRAPHIES = "geographies"
 CONF_INTEGRATION_TYPE = "integration_type"
 
-DATA_CLIENT = "client"
-
-TOPIC_UPDATE = f"airvisual_update_{0}"
+DATA_COORDINATOR = "coordinator"