Move DataUpdateCoordinator to coordinator module in tesla_wall_connector (#164937)

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>

homeassistant/components/tesla_wall_connector/__init__.py

@@ -2,35 +2,20 @@
 
 from __future__ import annotations
 
-from dataclasses import dataclass
-from datetime import timedelta
-import logging
-
 from tesla_wall_connector import WallConnector
-from tesla_wall_connector.exceptions import (
-    WallConnectorConnectionError,
-    WallConnectorConnectionTimeoutError,
-    WallConnectorError,
-)
+from tesla_wall_connector.exceptions import WallConnectorError
 
 from homeassistant.config_entries import ConfigEntry
-from homeassistant.const import CONF_HOST, CONF_SCAN_INTERVAL, Platform
+from homeassistant.const import CONF_HOST, Platform
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import ConfigEntryNotReady
 from homeassistant.helpers.aiohttp_client import async_get_clientsession
-from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
 
-from .const import (
-    DEFAULT_SCAN_INTERVAL,
-    DOMAIN,
-    WALLCONNECTOR_DATA_LIFETIME,
-    WALLCONNECTOR_DATA_VITALS,
-)
+from .const import DOMAIN
+from .coordinator import WallConnectorCoordinator, WallConnectorData, get_poll_interval
 
 PLATFORMS: list[Platform] = [Platform.BINARY_SENSOR, Platform.SENSOR]
 
-_LOGGER = logging.getLogger(__name__)
-
 
 async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Set up Tesla Wall Connector from a config entry."""
@@ -44,39 +29,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     except WallConnectorError as ex:
         raise ConfigEntryNotReady from ex
 
-    async def async_update_data():
-        """Fetch new data from the Wall Connector."""
-        try:
-            vitals = await wall_connector.async_get_vitals()
-            lifetime = await wall_connector.async_get_lifetime()
-        except WallConnectorConnectionTimeoutError as ex:
-            raise UpdateFailed(
-                f"Could not fetch data from Tesla WallConnector at {hostname}: Timeout"
-            ) from ex
-        except WallConnectorConnectionError as ex:
-            raise UpdateFailed(
-                f"Could not fetch data from Tesla WallConnector at {hostname}: Cannot"
-                " connect"
-            ) from ex
-        except WallConnectorError as ex:
-            raise UpdateFailed(
-                f"Could not fetch data from Tesla WallConnector at {hostname}: {ex}"
-            ) from ex
-
-        return {
-            WALLCONNECTOR_DATA_VITALS: vitals,
-            WALLCONNECTOR_DATA_LIFETIME: lifetime,
-        }
-
-    coordinator: DataUpdateCoordinator = DataUpdateCoordinator(
-        hass,
-        _LOGGER,
-        config_entry=entry,
-        name="tesla-wallconnector",
-        update_interval=get_poll_interval(entry),
-        update_method=async_update_data,
-    )
-
+    coordinator = WallConnectorCoordinator(hass, entry, hostname, wall_connector)
     await coordinator.async_config_entry_first_refresh()
 
     hass.data[DOMAIN][entry.entry_id] = WallConnectorData(
@@ -95,13 +48,6 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     return True
 
 
-def get_poll_interval(entry: ConfigEntry) -> timedelta:
-    """Get the poll interval from config."""
-    return timedelta(
-        seconds=entry.options.get(CONF_SCAN_INTERVAL, DEFAULT_SCAN_INTERVAL)
-    )
-
-
 async def update_listener(hass: HomeAssistant, entry: ConfigEntry) -> None:
     """Handle options update."""
     wall_connector_data: WallConnectorData = hass.data[DOMAIN][entry.entry_id]
@@ -114,15 +60,3 @@ async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
         hass.data[DOMAIN].pop(entry.entry_id)
 
     return unload_ok
-
-
-@dataclass
-class WallConnectorData:
-    """Data for the Tesla Wall Connector integration."""
-
-    wall_connector_client: WallConnector
-    update_coordinator: DataUpdateCoordinator
-    hostname: str
-    part_number: str
-    firmware_version: str
-    serial_number: str

homeassistant/components/tesla_wall_connector/binary_sensor.py

@@ -13,8 +13,8 @@ from homeassistant.const import EntityCategory
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
 
-from . import WallConnectorData
 from .const import DOMAIN, WALLCONNECTOR_DATA_VITALS
+from .coordinator import WallConnectorData
 from .entity import WallConnectorEntity, WallConnectorLambdaValueGetterMixin
 
 _LOGGER = logging.getLogger(__name__)

homeassistant/components/tesla_wall_connector/coordinator.py

@@ -0,0 +1,96 @@
+"""DataUpdateCoordinator for the Tesla Wall Connector integration."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from datetime import timedelta
+import logging
+
+from tesla_wall_connector import WallConnector
+from tesla_wall_connector.exceptions import (
+    WallConnectorConnectionError,
+    WallConnectorConnectionTimeoutError,
+    WallConnectorError,
+)
+
+from homeassistant.config_entries import ConfigEntry
+from homeassistant.const import CONF_SCAN_INTERVAL
+from homeassistant.core import HomeAssistant
+from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
+
+from .const import (
+    DEFAULT_SCAN_INTERVAL,
+    WALLCONNECTOR_DATA_LIFETIME,
+    WALLCONNECTOR_DATA_VITALS,
+)
+
+_LOGGER = logging.getLogger(__name__)
+
+
+@dataclass
+class WallConnectorData:
+    """Data for the Tesla Wall Connector integration."""
+
+    wall_connector_client: WallConnector
+    update_coordinator: WallConnectorCoordinator
+    hostname: str
+    part_number: str
+    firmware_version: str
+    serial_number: str
+
+
+def get_poll_interval(entry: ConfigEntry) -> timedelta:
+    """Get the poll interval from config."""
+    return timedelta(
+        seconds=entry.options.get(CONF_SCAN_INTERVAL, DEFAULT_SCAN_INTERVAL)
+    )
+
+
+class WallConnectorCoordinator(DataUpdateCoordinator[dict]):
+    """Class to manage fetching Tesla Wall Connector data."""
+
+    config_entry: ConfigEntry
+
+    def __init__(
+        self,
+        hass: HomeAssistant,
+        entry: ConfigEntry,
+        hostname: str,
+        wall_connector: WallConnector,
+    ) -> None:
+        """Initialize the coordinator."""
+        super().__init__(
+            hass,
+            _LOGGER,
+            config_entry=entry,
+            name="tesla-wallconnector",
+            update_interval=get_poll_interval(entry),
+        )
+        self._hostname = hostname
+        self._wall_connector = wall_connector
+
+    async def _async_update_data(self) -> dict:
+        """Fetch new data from the Wall Connector."""
+        try:
+            vitals = await self._wall_connector.async_get_vitals()
+            lifetime = await self._wall_connector.async_get_lifetime()
+        except WallConnectorConnectionTimeoutError as ex:
+            raise UpdateFailed(
+                f"Could not fetch data from Tesla WallConnector at {self._hostname}:"
+                " Timeout"
+            ) from ex
+        except WallConnectorConnectionError as ex:
+            raise UpdateFailed(
+                f"Could not fetch data from Tesla WallConnector at {self._hostname}:"
+                " Cannot connect"
+            ) from ex
+        except WallConnectorError as ex:
+            raise UpdateFailed(
+                f"Could not fetch data from Tesla WallConnector at {self._hostname}:"
+                f" {ex}"
+            ) from ex
+
+        return {
+            WALLCONNECTOR_DATA_VITALS: vitals,
+            WALLCONNECTOR_DATA_LIFETIME: lifetime,
+        }

homeassistant/components/tesla_wall_connector/entity.py

@@ -9,8 +9,8 @@ from typing import Any
 from homeassistant.helpers.device_registry import DeviceInfo
 from homeassistant.helpers.update_coordinator import CoordinatorEntity
 
-from . import WallConnectorData
 from .const import DOMAIN, WALLCONNECTOR_DEVICE_NAME
+from .coordinator import WallConnectorCoordinator, WallConnectorData
 
 
 @dataclass(frozen=True)
@@ -25,7 +25,7 @@ def _get_unique_id(serial_number: str, key: str) -> str:
     return f"{serial_number}-{key}"
 
 
-class WallConnectorEntity(CoordinatorEntity):
+class WallConnectorEntity(CoordinatorEntity[WallConnectorCoordinator]):
     """Base class for Wall Connector entities."""
 
     _attr_has_entity_name = True

homeassistant/components/tesla_wall_connector/sensor.py

@@ -22,8 +22,8 @@ from homeassistant.const import (
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
 
-from . import WallConnectorData
 from .const import DOMAIN, WALLCONNECTOR_DATA_LIFETIME, WALLCONNECTOR_DATA_VITALS
+from .coordinator import WallConnectorData
 from .entity import WallConnectorEntity, WallConnectorLambdaValueGetterMixin
 
 _LOGGER = logging.getLogger(__name__)