Move NAM Data Update Coordinator to the coordinator module (#116686)

Move Data Update Coordinator to the coordinator module

Co-authored-by: Maciej Bieniek <478555+bieniu@users.noreply.github.com>

homeassistant/components/nam/__init__.py

@@ -2,17 +2,13 @@
 
 from __future__ import annotations
 
-import asyncio
 import logging
-from typing import cast
 
 from aiohttp.client_exceptions import ClientConnectorError, ClientError
 from nettigo_air_monitor import (
     ApiError,
     AuthFailedError,
     ConnectionOptions,
-    InvalidSensorDataError,
-    NAMSensors,
     NettigoAirMonitor,
 )
 
@@ -21,25 +17,20 @@ from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_HOST, CONF_PASSWORD, CONF_USERNAME, Platform
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import ConfigEntryAuthFailed, ConfigEntryNotReady
-from homeassistant.helpers import device_registry as dr, entity_registry as er
+from homeassistant.helpers import entity_registry as er
 from homeassistant.helpers.aiohttp_client import async_get_clientsession
-from homeassistant.helpers.device_registry import DeviceInfo
-from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
 
-from .const import (
-    ATTR_SDS011,
-    ATTR_SPS30,
-    DEFAULT_UPDATE_INTERVAL,
-    DOMAIN,
-    MANUFACTURER,
-)
+from .const import ATTR_SDS011, ATTR_SPS30, DOMAIN
+from .coordinator import NAMDataUpdateCoordinator
 
 _LOGGER = logging.getLogger(__name__)
 
 PLATFORMS = [Platform.BUTTON, Platform.SENSOR]
 
+NAMConfigEntry = ConfigEntry[NAMDataUpdateCoordinator]
 
-async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+
+async def async_setup_entry(hass: HomeAssistant, entry: NAMConfigEntry) -> bool:
     """Set up Nettigo as config entry."""
     host: str = entry.data[CONF_HOST]
     username: str | None = entry.data.get(CONF_USERNAME)
@@ -63,8 +54,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     coordinator = NAMDataUpdateCoordinator(hass, nam, entry.unique_id)
     await coordinator.async_config_entry_first_refresh()
 
-    hass.data.setdefault(DOMAIN, {})
-    hass.data[DOMAIN][entry.entry_id] = coordinator
+    entry.runtime_data = coordinator
 
     await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
 
@@ -81,57 +71,6 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     return True
 
 
-async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+async def async_unload_entry(hass: HomeAssistant, entry: NAMConfigEntry) -> bool:
     """Unload a config entry."""
-    unload_ok = await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
-
-    if unload_ok:
-        hass.data[DOMAIN].pop(entry.entry_id)
-
-    return unload_ok
-
-
-class NAMDataUpdateCoordinator(DataUpdateCoordinator[NAMSensors]):  # pylint: disable=hass-enforce-coordinator-module
-    """Class to manage fetching Nettigo Air Monitor data."""
-
-    def __init__(
-        self,
-        hass: HomeAssistant,
-        nam: NettigoAirMonitor,
-        unique_id: str | None,
-    ) -> None:
-        """Initialize."""
-        self._unique_id = unique_id
-        self.nam = nam
-
-        super().__init__(
-            hass, _LOGGER, name=DOMAIN, update_interval=DEFAULT_UPDATE_INTERVAL
-        )
-
-    async def _async_update_data(self) -> NAMSensors:
-        """Update data via library."""
-        try:
-            async with asyncio.timeout(10):
-                data = await self.nam.async_update()
-        # We do not need to catch AuthFailed exception here because sensor data is
-        # always available without authorization.
-        except (ApiError, ClientConnectorError, InvalidSensorDataError) as error:
-            raise UpdateFailed(error) from error
-
-        return data
-
-    @property
-    def unique_id(self) -> str | None:
-        """Return a unique_id."""
-        return self._unique_id
-
-    @property
-    def device_info(self) -> DeviceInfo:
-        """Return the device info."""
-        return DeviceInfo(
-            connections={(dr.CONNECTION_NETWORK_MAC, cast(str, self._unique_id))},
-            name="Nettigo Air Monitor",
-            sw_version=self.nam.software_version,
-            manufacturer=MANUFACTURER,
-            configuration_url=f"http://{self.nam.host}/",
-        )
+    return await hass.config_entries.async_unload_platforms(entry, PLATFORMS)

homeassistant/components/nam/button.py

@@ -9,14 +9,12 @@ from homeassistant.components.button import (
     ButtonEntity,
     ButtonEntityDescription,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import EntityCategory
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.update_coordinator import CoordinatorEntity
 
-from . import NAMDataUpdateCoordinator
-from .const import DOMAIN
+from . import NAMConfigEntry, NAMDataUpdateCoordinator
 
 PARALLEL_UPDATES = 1
 
@@ -30,10 +28,10 @@ RESTART_BUTTON: ButtonEntityDescription = ButtonEntityDescription(
 
 
 async def async_setup_entry(
-    hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
+    hass: HomeAssistant, entry: NAMConfigEntry, async_add_entities: AddEntitiesCallback
 ) -> None:
     """Add a Nettigo Air Monitor entities from a config_entry."""
-    coordinator: NAMDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id]
+    coordinator = entry.runtime_data
 
     buttons: list[NAMButton] = []
     buttons.append(NAMButton(coordinator, RESTART_BUTTON))

homeassistant/components/nam/coordinator.py

@@ -0,0 +1,67 @@
+"""The Nettigo Air Monitor coordinator."""
+
+import asyncio
+import logging
+from typing import cast
+
+from aiohttp.client_exceptions import ClientConnectorError
+from nettigo_air_monitor import (
+    ApiError,
+    InvalidSensorDataError,
+    NAMSensors,
+    NettigoAirMonitor,
+)
+
+from homeassistant.core import HomeAssistant
+from homeassistant.helpers.device_registry import CONNECTION_NETWORK_MAC, DeviceInfo
+from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
+
+from .const import DEFAULT_UPDATE_INTERVAL, DOMAIN, MANUFACTURER
+
+_LOGGER = logging.getLogger(__name__)
+
+
+class NAMDataUpdateCoordinator(DataUpdateCoordinator[NAMSensors]):
+    """Class to manage fetching Nettigo Air Monitor data."""
+
+    def __init__(
+        self,
+        hass: HomeAssistant,
+        nam: NettigoAirMonitor,
+        unique_id: str | None,
+    ) -> None:
+        """Initialize."""
+        self._unique_id = unique_id
+        self.nam = nam
+
+        super().__init__(
+            hass, _LOGGER, name=DOMAIN, update_interval=DEFAULT_UPDATE_INTERVAL
+        )
+
+    async def _async_update_data(self) -> NAMSensors:
+        """Update data via library."""
+        try:
+            async with asyncio.timeout(10):
+                data = await self.nam.async_update()
+        # We do not need to catch AuthFailed exception here because sensor data is
+        # always available without authorization.
+        except (ApiError, ClientConnectorError, InvalidSensorDataError) as error:
+            raise UpdateFailed(error) from error
+
+        return data
+
+    @property
+    def unique_id(self) -> str | None:
+        """Return a unique_id."""
+        return self._unique_id
+
+    @property
+    def device_info(self) -> DeviceInfo:
+        """Return the device info."""
+        return DeviceInfo(
+            connections={(CONNECTION_NETWORK_MAC, cast(str, self._unique_id))},
+            name="Nettigo Air Monitor",
+            sw_version=self.nam.software_version,
+            manufacturer=MANUFACTURER,
+            configuration_url=f"http://{self.nam.host}/",
+        )

homeassistant/components/nam/diagnostics.py

@@ -6,21 +6,19 @@ from dataclasses import asdict
 from typing import Any
 
 from homeassistant.components.diagnostics import async_redact_data
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_PASSWORD, CONF_USERNAME
 from homeassistant.core import HomeAssistant
 
-from . import NAMDataUpdateCoordinator
-from .const import DOMAIN
+from . import NAMConfigEntry
 
 TO_REDACT = {CONF_PASSWORD, CONF_USERNAME}
 
 
 async def async_get_config_entry_diagnostics(
-    hass: HomeAssistant, config_entry: ConfigEntry
+    hass: HomeAssistant, config_entry: NAMConfigEntry
 ) -> dict[str, Any]:
     """Return diagnostics for a config entry."""
-    coordinator: NAMDataUpdateCoordinator = hass.data[DOMAIN][config_entry.entry_id]
+    coordinator = config_entry.runtime_data
 
     return {
         "info": async_redact_data(config_entry.data, TO_REDACT),

homeassistant/components/nam/sensor.py

@@ -16,7 +16,6 @@ from homeassistant.components.sensor import (
     SensorEntityDescription,
     SensorStateClass,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import (
     CONCENTRATION_MICROGRAMS_PER_CUBIC_METER,
     CONCENTRATION_PARTS_PER_MILLION,
@@ -33,7 +32,7 @@ from homeassistant.helpers.typing import StateType
 from homeassistant.helpers.update_coordinator import CoordinatorEntity
 from homeassistant.util.dt import utcnow
 
-from . import NAMDataUpdateCoordinator
+from . import NAMConfigEntry, NAMDataUpdateCoordinator
 from .const import (
     ATTR_BME280_HUMIDITY,
     ATTR_BME280_PRESSURE,
@@ -347,10 +346,10 @@ SENSORS: tuple[NAMSensorEntityDescription, ...] = (
 
 
 async def async_setup_entry(
-    hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
+    hass: HomeAssistant, entry: NAMConfigEntry, async_add_entities: AddEntitiesCallback
 ) -> None:
     """Add a Nettigo Air Monitor entities from a config_entry."""
-    coordinator: NAMDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id]
+    coordinator = entry.runtime_data
 
     # Due to the change of the attribute name of two sensors, it is necessary to migrate
     # the unique_ids to the new names.