From 9aaeefeb8e76d9c56b7bd32c300c3149b3ba0c22 Mon Sep 17 00:00:00 2001 From: Aaron Bach Date: Tue, 13 Dec 2022 23:31:13 -0700 Subject: [PATCH] Ensure AirVisual Pro uses long-running Samba connection (#83869) --- .../components/airvisual/manifest.json | 2 +- .../components/airvisual_pro/__init__.py | 64 ++++++++++++++++--- .../components/airvisual_pro/diagnostics.py | 6 +- .../components/airvisual_pro/manifest.json | 2 +- .../components/airvisual_pro/sensor.py | 7 +- requirements_all.txt | 2 +- requirements_test_all.txt | 2 +- 7 files changed, 65 insertions(+), 20 deletions(-) diff --git a/homeassistant/components/airvisual/manifest.json b/homeassistant/components/airvisual/manifest.json index b9823be2168..82104a08832 100644 --- a/homeassistant/components/airvisual/manifest.json +++ b/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"], diff --git a/homeassistant/components/airvisual_pro/__init__.py b/homeassistant/components/airvisual_pro/__init__.py index dace986db5b..8255019f14e 100644 --- a/homeassistant/components/airvisual_pro/__init__.py +++ b/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 diff --git a/homeassistant/components/airvisual_pro/diagnostics.py b/homeassistant/components/airvisual_pro/diagnostics.py index 16759c97580..d6e60207214 100644 --- a/homeassistant/components/airvisual_pro/diagnostics.py +++ b/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, ) diff --git a/homeassistant/components/airvisual_pro/manifest.json b/homeassistant/components/airvisual_pro/manifest.json index 6e47f1ebf9b..96e97c8f794 100644 --- a/homeassistant/components/airvisual_pro/manifest.json +++ b/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"], diff --git a/homeassistant/components/airvisual_pro/sensor.py b/homeassistant/components/airvisual_pro/sensor.py index 98f14c79c4c..7635a52e9b6 100644 --- a/homeassistant/components/airvisual_pro/sensor.py +++ b/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 ) diff --git a/requirements_all.txt b/requirements_all.txt index 313fca54579..e92d5e0a87b 100644 --- a/requirements_all.txt +++ b/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 diff --git a/requirements_test_all.txt b/requirements_test_all.txt index e74ef288a32..acbfe5e742c 100644 --- a/requirements_test_all.txt +++ b/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