diff --git a/homeassistant/components/airvisual_pro/__init__.py b/homeassistant/components/airvisual_pro/__init__.py index 3b3ac6df232..2c56086d399 100644 --- a/homeassistant/components/airvisual_pro/__init__.py +++ b/homeassistant/components/airvisual_pro/__init__.py @@ -4,18 +4,9 @@ from __future__ import annotations import asyncio from contextlib import suppress -from dataclasses import dataclass -from datetime import timedelta -from typing import Any -from pyairvisual.node import ( - InvalidAuthenticationError, - NodeConnectionError, - NodeProError, - NodeSamba, -) +from pyairvisual.node import NodeProError, NodeSamba -from homeassistant.config_entries import ConfigEntry from homeassistant.const import ( CONF_IP_ADDRESS, CONF_PASSWORD, @@ -23,25 +14,16 @@ from homeassistant.const import ( Platform, ) from homeassistant.core import Event, HomeAssistant -from homeassistant.exceptions import ConfigEntryAuthFailed, ConfigEntryNotReady -from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed +from homeassistant.exceptions import ConfigEntryNotReady -from .const import LOGGER +from .coordinator import ( + AirVisualProConfigEntry, + AirVisualProCoordinator, + AirVisualProData, +) PLATFORMS = [Platform.SENSOR] -UPDATE_INTERVAL = timedelta(minutes=1) - -type AirVisualProConfigEntry = ConfigEntry[AirVisualProData] - - -@dataclass -class AirVisualProData: - """Define a data class.""" - - coordinator: DataUpdateCoordinator - node: NodeSamba - async def async_setup_entry( hass: HomeAssistant, entry: AirVisualProConfigEntry @@ -54,48 +36,15 @@ async def async_setup_entry( 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: - data = await node.async_get_latest_measurements() - data["history"] = {} - if data["settings"].get("follow_mode") == "device": - history = await node.async_get_history(include_trends=False) - data["history"] = history.get("measurements", [])[-1] - except InvalidAuthenticationError as err: - raise ConfigEntryAuthFailed("Invalid Samba password") from err - 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, - config_entry=entry, - name="Node/Pro data", - update_interval=UPDATE_INTERVAL, - update_method=async_get_data, - ) - + coordinator = AirVisualProCoordinator(hass, entry, node) await coordinator.async_config_entry_first_refresh() entry.runtime_data = 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: + if coordinator.reload_task: with suppress(asyncio.CancelledError): - reload_task.cancel() + coordinator.reload_task.cancel() await node.async_disconnect() entry.async_on_unload( diff --git a/homeassistant/components/airvisual_pro/coordinator.py b/homeassistant/components/airvisual_pro/coordinator.py new file mode 100644 index 00000000000..946a247ace1 --- /dev/null +++ b/homeassistant/components/airvisual_pro/coordinator.py @@ -0,0 +1,79 @@ +"""DataUpdateCoordinator for the AirVisual Pro integration.""" + +from __future__ import annotations + +import asyncio +from dataclasses import dataclass +from datetime import timedelta +from typing import Any + +from pyairvisual.node import ( + InvalidAuthenticationError, + NodeConnectionError, + NodeProError, + NodeSamba, +) + +from homeassistant.config_entries import ConfigEntry +from homeassistant.core import HomeAssistant +from homeassistant.exceptions import ConfigEntryAuthFailed +from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed + +from .const import LOGGER + +UPDATE_INTERVAL = timedelta(minutes=1) + + +@dataclass +class AirVisualProData: + """Define a data class.""" + + coordinator: AirVisualProCoordinator + node: NodeSamba + + +type AirVisualProConfigEntry = ConfigEntry[AirVisualProData] + + +class AirVisualProCoordinator(DataUpdateCoordinator[dict[str, Any]]): + """Coordinator for AirVisual Pro data.""" + + config_entry: AirVisualProConfigEntry + + def __init__( + self, + hass: HomeAssistant, + config_entry: AirVisualProConfigEntry, + node: NodeSamba, + ) -> None: + """Initialize.""" + super().__init__( + hass, + LOGGER, + config_entry=config_entry, + name="Node/Pro data", + update_interval=UPDATE_INTERVAL, + ) + self._node = node + self.reload_task: asyncio.Task[bool] | None = None + + async def _async_update_data(self) -> dict[str, Any]: + """Get data from the device.""" + try: + data = await self._node.async_get_latest_measurements() + data["history"] = {} + if data["settings"].get("follow_mode") == "device": + history = await self._node.async_get_history(include_trends=False) + data["history"] = history.get("measurements", [])[-1] + except InvalidAuthenticationError as err: + raise ConfigEntryAuthFailed("Invalid Samba password") from err + except NodeConnectionError as err: + if self.reload_task is None: + self.reload_task = self.hass.async_create_task( + self.hass.config_entries.async_reload(self.config_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 diff --git a/homeassistant/components/airvisual_pro/diagnostics.py b/homeassistant/components/airvisual_pro/diagnostics.py index da871442547..dc69483c78f 100644 --- a/homeassistant/components/airvisual_pro/diagnostics.py +++ b/homeassistant/components/airvisual_pro/diagnostics.py @@ -8,7 +8,7 @@ from homeassistant.components.diagnostics import async_redact_data from homeassistant.const import CONF_PASSWORD from homeassistant.core import HomeAssistant -from . import AirVisualProConfigEntry +from .coordinator import AirVisualProConfigEntry CONF_MAC_ADDRESS = "mac_address" CONF_SERIAL_NUMBER = "serial_number" diff --git a/homeassistant/components/airvisual_pro/entity.py b/homeassistant/components/airvisual_pro/entity.py index bc28fa36e52..b44c5ed8bce 100644 --- a/homeassistant/components/airvisual_pro/entity.py +++ b/homeassistant/components/airvisual_pro/entity.py @@ -4,19 +4,17 @@ from __future__ import annotations from homeassistant.helpers.device_registry import DeviceInfo from homeassistant.helpers.entity import EntityDescription -from homeassistant.helpers.update_coordinator import ( - CoordinatorEntity, - DataUpdateCoordinator, -) +from homeassistant.helpers.update_coordinator import CoordinatorEntity from .const import DOMAIN +from .coordinator import AirVisualProCoordinator -class AirVisualProEntity(CoordinatorEntity): +class AirVisualProEntity(CoordinatorEntity[AirVisualProCoordinator]): """Define a generic AirVisual Pro entity.""" def __init__( - self, coordinator: DataUpdateCoordinator, description: EntityDescription + self, coordinator: AirVisualProCoordinator, description: EntityDescription ) -> None: """Initialize.""" super().__init__(coordinator) diff --git a/homeassistant/components/airvisual_pro/sensor.py b/homeassistant/components/airvisual_pro/sensor.py index 215370736fe..3fac272e655 100644 --- a/homeassistant/components/airvisual_pro/sensor.py +++ b/homeassistant/components/airvisual_pro/sensor.py @@ -22,7 +22,7 @@ from homeassistant.const import ( from homeassistant.core import HomeAssistant, callback from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback -from . import AirVisualProConfigEntry +from .coordinator import AirVisualProConfigEntry from .entity import AirVisualProEntity