Add support for AirVisual Node/Pro units

.coveragerc

@@ -21,6 +21,7 @@ omit =
     homeassistant/components/airly/sensor.py
     homeassistant/components/airly/const.py
     homeassistant/components/airvisual/__init__.py
+    homeassistant/components/airvisual/air_quality.py
     homeassistant/components/airvisual/sensor.py
     homeassistant/components/aladdin_connect/cover.py
     homeassistant/components/alarmdecoder/*

homeassistant/components/airvisual/__init__.py

@@ -1,38 +1,50 @@
 """The airvisual component."""
-import logging
+import asyncio
+from datetime import timedelta
 
 from pyairvisual import Client
-from pyairvisual.errors import AirVisualError, InvalidKeyError
+from pyairvisual.errors import AirVisualError, NodeProError
 import voluptuous as vol
 
 from homeassistant.config_entries import SOURCE_IMPORT
 from homeassistant.const import (
+    ATTR_ATTRIBUTION,
     CONF_API_KEY,
+    CONF_IP_ADDRESS,
     CONF_LATITUDE,
     CONF_LONGITUDE,
+    CONF_PASSWORD,
     CONF_SHOW_ON_MAP,
     CONF_STATE,
 )
 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
+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 .const import (
     CONF_CITY,
     CONF_COUNTRY,
     CONF_GEOGRAPHIES,
     DATA_CLIENT,
-    DEFAULT_SCAN_INTERVAL,
     DOMAIN,
+    INTEGRATION_TYPE_GEOGRAPHY,
+    INTEGRATION_TYPE_NODE_PRO,
+    LOGGER,
     TOPIC_UPDATE,
 )
 
-_LOGGER = logging.getLogger(__name__)
+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)
 DEFAULT_OPTIONS = {CONF_SHOW_ON_MAP: True}
 
 GEOGRAPHY_COORDINATES_SCHEMA = vol.Schema(
@@ -66,6 +78,9 @@
 @callback
 def async_get_geography_id(geography_dict):
     """Generate a unique ID from a geography dict."""
+    if not geography_dict:
+        return
+
     if CONF_CITY in geography_dict:
         return ", ".join(
             (
@@ -103,53 +118,66 @@ async def async_setup(hass, config):
     return True
 
 
-async def async_setup_entry(hass, config_entry):
-    """Set up AirVisual as config entry."""
+@callback
+def _standardize_geography_config_entry(hass, config_entry):
+    """Ensure that geography observables have appropriate properties."""
     entry_updates = {}
+
     if not config_entry.unique_id:
         # If the config entry doesn't already have a unique ID, set one:
         entry_updates["unique_id"] = config_entry.data[CONF_API_KEY]
     if not config_entry.options:
         # If the config entry doesn't already have any options set, set defaults:
-        entry_updates["options"] = DEFAULT_OPTIONS
+        entry_updates["options"] = {CONF_SHOW_ON_MAP: True}
 
-    if entry_updates:
-        hass.config_entries.async_update_entry(config_entry, **entry_updates)
+    if not entry_updates:
+        return
 
+    hass.config_entries.async_update_entry(config_entry, **entry_updates)
+
+
+async def async_setup_entry(hass, config_entry):
+    """Set up AirVisual as config entry."""
     websession = aiohttp_client.async_get_clientsession(hass)
 
-    hass.data[DOMAIN][DATA_CLIENT][config_entry.entry_id] = AirVisualData(
-        hass, Client(websession, api_key=config_entry.data[CONF_API_KEY]), config_entry
-    )
+    if CONF_API_KEY in config_entry.data:
+        _standardize_geography_config_entry(hass, config_entry)
+        airvisual = AirVisualGeographyData(
+            hass,
+            Client(websession, api_key=config_entry.data[CONF_API_KEY]),
+            config_entry,
+        )
 
-    try:
-        await hass.data[DOMAIN][DATA_CLIENT][config_entry.entry_id].async_update()
-    except InvalidKeyError:
-        _LOGGER.error("Invalid API key provided")
-        raise ConfigEntryNotReady
+        # Only geography-based entries have options:
+        config_entry.add_update_listener(async_update_options)
+    else:
+        airvisual = AirVisualNodeProData(hass, Client(websession), config_entry)
 
-    hass.async_create_task(
-        hass.config_entries.async_forward_entry_setup(config_entry, "sensor")
-    )
+    await airvisual.async_update()
+
+    hass.data[DOMAIN][DATA_CLIENT][config_entry.entry_id] = airvisual
+
+    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 hass.data[DOMAIN][DATA_CLIENT][config_entry.entry_id].async_update()
+        await airvisual.async_update()
 
     hass.data[DOMAIN][DATA_LISTENER][config_entry.entry_id] = async_track_time_interval(
-        hass, refresh, DEFAULT_SCAN_INTERVAL
+        hass, refresh, airvisual.scan_interval
     )
 
-    config_entry.add_update_listener(async_update_options)
-
     return True
 
 
 async def async_migrate_entry(hass, config_entry):
     """Migrate an old config entry."""
     version = config_entry.version
 
-    _LOGGER.debug("Migrating from version %s", version)
+    LOGGER.debug("Migrating from version %s", version)
 
     # 1 -> 2: One geography per config entry
     if version == 1:
@@ -178,21 +206,27 @@ async def async_migrate_entry(hass, config_entry):
                 )
             )
 
-    _LOGGER.info("Migration to version %s successful", version)
+    LOGGER.info("Migration to version %s successful", version)
 
     return True
 
 
 async def async_unload_entry(hass, config_entry):
     """Unload an AirVisual config entry."""
-    hass.data[DOMAIN][DATA_CLIENT].pop(config_entry.entry_id)
-
-    remove_listener = hass.data[DOMAIN][DATA_LISTENER].pop(config_entry.entry_id)
-    remove_listener()
+    unload_ok = all(
+        await asyncio.gather(
+            *[
+                hass.config_entries.async_forward_entry_unload(config_entry, component)
+                for component in PLATFORMS
+            ]
+        )
+    )
+    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()
 
-    await hass.config_entries.async_forward_entry_unload(config_entry, "sensor")
-
-    return True
+    return unload_ok
 
 
 async def async_update_options(hass, config_entry):
@@ -201,7 +235,53 @@ async def async_update_options(hass, config_entry):
     airvisual.async_update_options(config_entry.options)
 
 
-class AirVisualData:
+class AirVisualEntity(Entity):
+    """Define a generic AirVisual entity."""
+
+    def __init__(self, airvisual):
+        """Initialize."""
+        self._airvisual = airvisual
+        self._attrs = {ATTR_ATTRIBUTION: DEFAULT_ATTRIBUTION}
+        self._icon = None
+        self._unit = None
+
+    @property
+    def device_state_attributes(self):
+        """Return the device state attributes."""
+        return self._attrs
+
+    @property
+    def icon(self):
+        """Return the icon."""
+        return self._icon
+
+    @property
+    def unit_of_measurement(self):
+        """Return the unit the value is expressed in."""
+        return self._unit
+
+    async def async_added_to_hass(self):
+        """Register callbacks."""
+
+        @callback
+        def update():
+            """Update the state."""
+            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.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):
@@ -211,7 +291,10 @@ def __init__(self, hass, client, config_entry):
         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."""
@@ -229,14 +312,43 @@ async def async_update(self):
         try:
             self.data[self.geography_id] = await api_coro
         except AirVisualError as err:
-            _LOGGER.error("Error while retrieving data: %s", err)
+            LOGGER.error("Error while retrieving data: %s", err)
             self.data[self.geography_id] = {}
 
-        _LOGGER.debug("Received new data")
-        async_dispatcher_send(self._hass, TOPIC_UPDATE)
+        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, TOPIC_UPDATE)
+        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)