Ensure AirVisual Pro uses long-running Samba connection (#83869)

homeassistant/components/airvisual/manifest.json

@@ -3,7 +3,7 @@
   "name": "AirVisual Cloud",
   "config_flow": true,
   "documentation": "https://www.home-assistant.io/integrations/airvisual",
-  "requirements": ["pyairvisual==2022.11.1"],
+  "requirements": ["pyairvisual==2022.12.0"],
   "codeowners": ["@bachya"],
   "iot_class": "cloud_polling",
   "loggers": ["pyairvisual", "pysmb"],

homeassistant/components/airvisual_pro/__init__.py

@@ -1,15 +1,24 @@
 """The AirVisual Pro integration."""
 from __future__ import annotations
 
+import asyncio
+from contextlib import suppress
+from dataclasses import dataclass
 from datetime import timedelta
 from typing import Any
 
 from pyairvisual import NodeSamba
-from pyairvisual.node import NodeProError
+from pyairvisual.node import NodeConnectionError, NodeProError
 
 from homeassistant.config_entries import ConfigEntry
-from homeassistant.const import CONF_IP_ADDRESS, CONF_PASSWORD, Platform
-from homeassistant.core import HomeAssistant, callback
+from homeassistant.const import (
+    CONF_IP_ADDRESS,
+    CONF_PASSWORD,
+    EVENT_HOMEASSISTANT_STOP,
+    Platform,
+)
+from homeassistant.core import Event, HomeAssistant, callback
+from homeassistant.exceptions import ConfigEntryNotReady
 from homeassistant.helpers.entity import DeviceInfo, EntityDescription
 from homeassistant.helpers.update_coordinator import (
     CoordinatorEntity,
@@ -24,19 +33,41 @@ PLATFORMS = [Platform.SENSOR]
 UPDATE_INTERVAL = timedelta(minutes=1)
 
 
+@dataclass
+class AirVisualProData:
+    """Define a data class."""
+
+    coordinator: DataUpdateCoordinator
+    node: NodeSamba
+
+
 async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Set up AirVisual Pro from a config entry."""
+    node = NodeSamba(entry.data[CONF_IP_ADDRESS], entry.data[CONF_PASSWORD])
+
+    try:
+        await node.async_connect()
+    except NodeProError as err:
+        raise ConfigEntryNotReady() from err
+
+    reload_task: asyncio.Task | None = None
 
     async def async_get_data() -> dict[str, Any]:
         """Get data from the device."""
         try:
-            async with NodeSamba(
-                entry.data[CONF_IP_ADDRESS], entry.data[CONF_PASSWORD]
-            ) as node:
-                return await node.async_get_latest_measurements()
+            data = await node.async_get_latest_measurements()
+        except NodeConnectionError as err:
+            nonlocal reload_task
+            if not reload_task:
+                reload_task = hass.async_create_task(
+                    hass.config_entries.async_reload(entry.entry_id)
+                )
+            raise UpdateFailed(f"Connection to Pro unit lost: {err}") from err
         except NodeProError as err:
             raise UpdateFailed(f"Error while retrieving data: {err}") from err
 
+        return data
+
     coordinator = DataUpdateCoordinator(
         hass,
         LOGGER,
@@ -46,7 +77,21 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     )
 
     await coordinator.async_config_entry_first_refresh()
-    hass.data.setdefault(DOMAIN, {})[entry.entry_id] = coordinator
+    hass.data.setdefault(DOMAIN, {})[entry.entry_id] = AirVisualProData(
+        coordinator=coordinator, node=node
+    )
+
+    async def async_shutdown(_: Event) -> None:
+        """Define an event handler to disconnect from the websocket."""
+        nonlocal reload_task
+        if reload_task:
+            with suppress(asyncio.CancelledError):
+                reload_task.cancel()
+        await node.async_disconnect()
+
+    entry.async_on_unload(
+        hass.bus.async_listen_once(EVENT_HOMEASSISTANT_STOP, async_shutdown)
+    )
 
     await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
 
@@ -56,7 +101,8 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
 async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Unload a config entry."""
     if unload_ok := await hass.config_entries.async_unload_platforms(entry, PLATFORMS):
-        hass.data[DOMAIN].pop(entry.entry_id)
+        data = hass.data[DOMAIN].pop(entry.entry_id)
+        await data.node.async_disconnect()
 
     return unload_ok
 

homeassistant/components/airvisual_pro/diagnostics.py

@@ -7,8 +7,8 @@ from homeassistant.components.diagnostics import async_redact_data
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_PASSWORD
 from homeassistant.core import HomeAssistant
-from homeassistant.helpers.update_coordinator import DataUpdateCoordinator
 
+from . import AirVisualProData
 from .const import DOMAIN
 
 CONF_MAC_ADDRESS = "mac_address"
@@ -25,12 +25,12 @@ async def async_get_config_entry_diagnostics(
     hass: HomeAssistant, entry: ConfigEntry
 ) -> dict[str, Any]:
     """Return diagnostics for a config entry."""
-    coordinator: DataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id]
+    data: AirVisualProData = hass.data[DOMAIN][entry.entry_id]
 
     return async_redact_data(
         {
             "entry": entry.as_dict(),
-            "data": coordinator.data,
+            "data": data.coordinator.data,
         },
         TO_REDACT,
     )

homeassistant/components/airvisual_pro/manifest.json

@@ -3,7 +3,7 @@
   "name": "AirVisual Pro",
   "config_flow": true,
   "documentation": "https://www.home-assistant.io/integrations/airvisual_pro",
-  "requirements": ["pyairvisual==2022.11.1"],
+  "requirements": ["pyairvisual==2022.12.0"],
   "codeowners": ["@bachya"],
   "iot_class": "local_polling",
   "loggers": ["pyairvisual", "pysmb"],

homeassistant/components/airvisual_pro/sensor.py

@@ -19,9 +19,8 @@ from homeassistant.const import (
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.entity import EntityCategory
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
-from homeassistant.helpers.update_coordinator import DataUpdateCoordinator
 
-from . import AirVisualProEntity
+from . import AirVisualProData, AirVisualProEntity
 from .const import DOMAIN
 
 SENSOR_KIND_AQI = "air_quality_index"
@@ -113,10 +112,10 @@ async def async_setup_entry(
     hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
 ) -> None:
     """Set up AirVisual sensors based on a config entry."""
-    coordinator: DataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id]
+    data: AirVisualProData = hass.data[DOMAIN][entry.entry_id]
 
     async_add_entities(
-        AirVisualProSensor(coordinator, description)
+        AirVisualProSensor(data.coordinator, description)
         for description in SENSOR_DESCRIPTIONS
     )
 

requirements_all.txt

@@ -1464,7 +1464,7 @@ pyairnow==1.1.0
 
 # homeassistant.components.airvisual
 # homeassistant.components.airvisual_pro
-pyairvisual==2022.11.1
+pyairvisual==2022.12.0
 
 # homeassistant.components.almond
 pyalmond==0.0.2

requirements_test_all.txt

@@ -1052,7 +1052,7 @@ pyairnow==1.1.0
 
 # homeassistant.components.airvisual
 # homeassistant.components.airvisual_pro
-pyairvisual==2022.11.1
+pyairvisual==2022.12.0
 
 # homeassistant.components.almond
 pyalmond==0.0.2