Refactor the SMA integration to use a dedicated DataUpdateCoordinator (home-assistant#154863)

Author: erwindouna

homeassistant/components/sma/__init__.py

@@ -1,153 +1,69 @@
-"""The sma integration."""
+"""The SMA integration."""
 
 from __future__ import annotations
 
-from datetime import timedelta
 import logging
-from typing import TYPE_CHECKING
 
-import pysma
+from pysma import SMA
 
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import (
-    ATTR_CONNECTIONS,
     CONF_HOST,
-    CONF_MAC,
     CONF_PASSWORD,
-    CONF_SCAN_INTERVAL,
     CONF_SSL,
     CONF_VERIFY_SSL,
     EVENT_HOMEASSISTANT_STOP,
+    Platform,
 )
-from homeassistant.core import HomeAssistant
-from homeassistant.exceptions import ConfigEntryAuthFailed, ConfigEntryNotReady
-from homeassistant.helpers import device_registry as dr
+from homeassistant.core import Event, HomeAssistant
 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 (
-    CONF_GROUP,
-    DEFAULT_SCAN_INTERVAL,
-    DOMAIN,
-    PLATFORMS,
-    PYSMA_COORDINATOR,
-    PYSMA_DEVICE_INFO,
-    PYSMA_OBJECT,
-    PYSMA_REMOVE_LISTENER,
-    PYSMA_SENSORS,
-)
+
+from .const import CONF_GROUP
+from .coordinator import SMADataUpdateCoordinator
+
+PLATFORMS = [Platform.SENSOR]
 
 _LOGGER = logging.getLogger(__name__)
 
 
-async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+type SMAConfigEntry = ConfigEntry[SMADataUpdateCoordinator]
+
+
+async def async_setup_entry(hass: HomeAssistant, entry: SMAConfigEntry) -> bool:
     """Set up sma from a config entry."""
-    # Init the SMA interface
+
     protocol = "https" if entry.data[CONF_SSL] else "http"
     url = f"{protocol}://{entry.data[CONF_HOST]}"
-    verify_ssl = entry.data[CONF_VERIFY_SSL]
-    group = entry.data[CONF_GROUP]
-    password = entry.data[CONF_PASSWORD]
-
-    session = async_get_clientsession(hass, verify_ssl=verify_ssl)
-    sma = pysma.SMA(session, url, password, group)
-
-    try:
-        # Get updated device info
-        sma_device_info = await sma.device_info()
-        # Get all device sensors
-        sensor_def = await sma.get_sensors()
-    except (
-        pysma.exceptions.SmaReadException,
-        pysma.exceptions.SmaConnectionException,
-    ) as exc:
-        raise ConfigEntryNotReady from exc
-    except pysma.exceptions.SmaAuthenticationException as exc:
-        raise ConfigEntryAuthFailed from exc
-
-    if TYPE_CHECKING:
-        assert entry.unique_id
-
-    # Create DeviceInfo object from sma_device_info
-    device_info = DeviceInfo(
-        configuration_url=url,
-        identifiers={(DOMAIN, entry.unique_id)},
-        manufacturer=sma_device_info["manufacturer"],
-        model=sma_device_info["type"],
-        name=sma_device_info["name"],
-        sw_version=sma_device_info["sw_version"],
-        serial_number=sma_device_info["serial"],
-    )
 
-    # Add the MAC address to connections, if it comes via DHCP
-    if CONF_MAC in entry.data:
-        device_info[ATTR_CONNECTIONS] = {
-            (dr.CONNECTION_NETWORK_MAC, entry.data[CONF_MAC])
-        }
-
-    # Define the coordinator
-    async def async_update_data():
-        """Update the used SMA sensors."""
-        try:
-            await sma.read(sensor_def)
-        except (
-            pysma.exceptions.SmaReadException,
-            pysma.exceptions.SmaConnectionException,
-        ) as exc:
-            raise UpdateFailed(exc) from exc
-
-    interval = timedelta(
-        seconds=entry.options.get(CONF_SCAN_INTERVAL, DEFAULT_SCAN_INTERVAL)
+    sma = SMA(
+        session=async_get_clientsession(
+            hass=hass, verify_ssl=entry.data[CONF_VERIFY_SSL]
+        ),
+        url=url,
+        password=entry.data[CONF_PASSWORD],
+        group=entry.data[CONF_GROUP],
     )
 
-    coordinator = DataUpdateCoordinator(
-        hass,
-        _LOGGER,
-        config_entry=entry,
-        name="sma",
-        update_method=async_update_data,
-        update_interval=interval,
-    )
+    coordinator = SMADataUpdateCoordinator(hass, entry, sma)
+    await coordinator.async_config_entry_first_refresh()
 
-    try:
-        await coordinator.async_config_entry_first_refresh()
-    except ConfigEntryNotReady:
-        await sma.close_session()
-        raise
+    entry.runtime_data = coordinator
+    await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
 
-    # Ensure we logout on shutdown
-    async def async_close_session(event):
-        """Close the session."""
-        await sma.close_session()
+    # Ensure the SMA session closes when Home Assistant stops
+    async def _async_handle_shutdown(event: Event) -> None:
+        await coordinator.async_close_sma_session()
 
-    remove_stop_listener = hass.bus.async_listen_once(
-        EVENT_HOMEASSISTANT_STOP, async_close_session
+    entry.async_on_unload(
+        hass.bus.async_listen_once(EVENT_HOMEASSISTANT_STOP, _async_handle_shutdown)
     )
 
-    hass.data.setdefault(DOMAIN, {})
-    hass.data[DOMAIN][entry.entry_id] = {
-        PYSMA_OBJECT: sma,
-        PYSMA_COORDINATOR: coordinator,
-        PYSMA_SENSORS: sensor_def,
-        PYSMA_REMOVE_LISTENER: remove_stop_listener,
-        PYSMA_DEVICE_INFO: device_info,
-    }
-
-    await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
-
     return True
 
 
-async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+async def async_unload_entry(hass: HomeAssistant, entry: SMAConfigEntry) -> bool:
     """Unload a config entry."""
-    unload_ok = await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
-    if unload_ok:
-        data = hass.data[DOMAIN].pop(entry.entry_id)
-        await data[PYSMA_OBJECT].close_session()
-        data[PYSMA_REMOVE_LISTENER]()
-
-    return unload_ok
+    return await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
 
 
 async def async_migrate_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
@@ -156,7 +72,6 @@ async def async_migrate_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     _LOGGER.debug("Migrating from version %s", entry.version)
 
     if entry.version == 1:
-        # 1 -> 2: Unique ID from integer to string
         if entry.minor_version == 1:
             minor_version = 2
             hass.config_entries.async_update_entry(

homeassistant/components/sma/const.py

@@ -1,7 +1,5 @@
 """Constants for the sma integration."""
 
-from homeassistant.const import Platform
-
 DOMAIN = "sma"
 
 PYSMA_COORDINATOR = "coordinator"
@@ -10,7 +8,6 @@
 PYSMA_SENSORS = "pysma_sensors"
 PYSMA_DEVICE_INFO = "device_info"
 
-PLATFORMS = [Platform.SENSOR]
 
 CONF_GROUP = "group"
 

homeassistant/components/sma/coordinator.py

@@ -0,0 +1,113 @@
+"""Coordinator for the SMA integration."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from datetime import timedelta
+import logging
+
+from pysma import SMA
+from pysma.exceptions import (
+    SmaAuthenticationException,
+    SmaConnectionException,
+    SmaReadException,
+)
+from pysma.sensor import Sensor
+
+from homeassistant.config_entries import ConfigEntry
+from homeassistant.const import CONF_SCAN_INTERVAL
+from homeassistant.core import HomeAssistant
+from homeassistant.exceptions import ConfigEntryAuthFailed, ConfigEntryNotReady
+from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
+
+from .const import DEFAULT_SCAN_INTERVAL, DOMAIN
+
+_LOGGER = logging.getLogger(__name__)
+
+
+@dataclass(slots=True)
+class SMACoordinatorData:
+    """Data class for SMA sensors."""
+
+    sma_device_info: dict[str, str]
+    sensors: list[Sensor]
+
+
+class SMADataUpdateCoordinator(DataUpdateCoordinator[SMACoordinatorData]):
+    """Data Update Coordinator for SMA."""
+
+    config_entry: ConfigEntry
+
+    def __init__(
+        self,
+        hass: HomeAssistant,
+        config_entry: ConfigEntry,
+        sma: SMA,
+    ) -> None:
+        """Initialize the SMA Data Update Coordinator."""
+        super().__init__(
+            hass,
+            _LOGGER,
+            config_entry=config_entry,
+            name=DOMAIN,
+            update_interval=timedelta(
+                seconds=config_entry.options.get(
+                    CONF_SCAN_INTERVAL, DEFAULT_SCAN_INTERVAL
+                )
+            ),
+        )
+        self.sma = sma
+        self._sma_device_info: dict[str, str] = {}
+        self._sensors: list[Sensor] = []
+
+    async def _async_setup(self) -> None:
+        """Setup the SMA Data Update Coordinator."""
+        try:
+            self._sma_device_info = await self.sma.device_info()
+            self._sensors = await self.sma.get_sensors()
+        except (
+            SmaReadException,
+            SmaConnectionException,
+        ) as err:
+            await self.async_close_sma_session()
+            raise ConfigEntryNotReady(
+                translation_domain=DOMAIN,
+                translation_key="cannot_connect",
+                translation_placeholders={"error": repr(err)},
+            ) from err
+        except SmaAuthenticationException as err:
+            raise ConfigEntryAuthFailed(
+                translation_domain=DOMAIN,
+                translation_key="invalid_auth",
+                translation_placeholders={"error": repr(err)},
+            ) from err
+
+    async def _async_update_data(self) -> SMACoordinatorData:
+        """Update the used SMA sensors."""
+        try:
+            await self.sma.read(self._sensors)
+        except (
+            SmaReadException,
+            SmaConnectionException,
+        ) as err:
+            raise UpdateFailed(
+                translation_domain=DOMAIN,
+                translation_key="cannot_connect",
+                translation_placeholders={"error": repr(err)},
+            ) from err
+        except SmaAuthenticationException as err:
+            raise ConfigEntryAuthFailed(
+                translation_domain=DOMAIN,
+                translation_key="invalid_auth",
+                translation_placeholders={"error": repr(err)},
+            ) from err
+
+        return SMACoordinatorData(
+            sma_device_info=self._sma_device_info,
+            sensors=self._sensors,
+        )
+
+    async def async_close_sma_session(self) -> None:
+        """Close the SMA session."""
+        await self.sma.close_session()
+        _LOGGER.debug("SMA session closed")