Move AirVisual Pro coordinator to separate module (#164742)

Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

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(

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

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"

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)

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