Move wallbox coordinator to its own file (#101577)

homeassistant/components/wallbox/__init__.py

@@ -1,195 +1,17 @@
 """The Wallbox integration."""
 from __future__ import annotations
 
-from datetime import timedelta
-from http import HTTPStatus
-import logging
-from typing import Any
-
-import requests
 from wallbox import Wallbox
 
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_PASSWORD, CONF_USERNAME, Platform
 from homeassistant.core import HomeAssistant
-from homeassistant.exceptions import ConfigEntryAuthFailed, HomeAssistantError
+from homeassistant.exceptions import ConfigEntryAuthFailed
-from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
 
-from .const import (
+from .const import CONF_STATION, DOMAIN, UPDATE_INTERVAL
-    CHARGER_CURRENCY_KEY,
+from .coordinator import InvalidAuth, WallboxCoordinator
-    CHARGER_DATA_KEY,
-    CHARGER_ENERGY_PRICE_KEY,
-    CHARGER_LOCKED_UNLOCKED_KEY,
-    CHARGER_MAX_CHARGING_CURRENT_KEY,
-    CHARGER_STATUS_DESCRIPTION_KEY,
-    CHARGER_STATUS_ID_KEY,
-    CODE_KEY,
-    CONF_STATION,
-    DOMAIN,
-    ChargerStatus,
-)
-
-_LOGGER = logging.getLogger(__name__)
 
 PLATFORMS = [Platform.SENSOR, Platform.NUMBER, Platform.LOCK, Platform.SWITCH]
-UPDATE_INTERVAL = 30
-
-# Translation of StatusId based on Wallbox portal code:
-# https://my.wallbox.com/src/utilities/charger/chargerStatuses.js
-CHARGER_STATUS: dict[int, ChargerStatus] = {
-    0: ChargerStatus.DISCONNECTED,
-    14: ChargerStatus.ERROR,
-    15: ChargerStatus.ERROR,
-    161: ChargerStatus.READY,
-    162: ChargerStatus.READY,
-    163: ChargerStatus.DISCONNECTED,
-    164: ChargerStatus.WAITING,
-    165: ChargerStatus.LOCKED,
-    166: ChargerStatus.UPDATING,
-    177: ChargerStatus.SCHEDULED,
-    178: ChargerStatus.PAUSED,
-    179: ChargerStatus.SCHEDULED,
-    180: ChargerStatus.WAITING_FOR_CAR,
-    181: ChargerStatus.WAITING_FOR_CAR,
-    182: ChargerStatus.PAUSED,
-    183: ChargerStatus.WAITING_IN_QUEUE_POWER_SHARING,
-    184: ChargerStatus.WAITING_IN_QUEUE_POWER_SHARING,
-    185: ChargerStatus.WAITING_IN_QUEUE_POWER_BOOST,
-    186: ChargerStatus.WAITING_IN_QUEUE_POWER_BOOST,
-    187: ChargerStatus.WAITING_MID_FAILED,
-    188: ChargerStatus.WAITING_MID_SAFETY,
-    189: ChargerStatus.WAITING_IN_QUEUE_ECO_SMART,
-    193: ChargerStatus.CHARGING,
-    194: ChargerStatus.CHARGING,
-    195: ChargerStatus.CHARGING,
-    196: ChargerStatus.DISCHARGING,
-    209: ChargerStatus.LOCKED,
-    210: ChargerStatus.LOCKED_CAR_CONNECTED,
-}
-
-
-class WallboxCoordinator(DataUpdateCoordinator[dict[str, Any]]):
-    """Wallbox Coordinator class."""
-
-    def __init__(self, station: str, wallbox: Wallbox, hass: HomeAssistant) -> None:
-        """Initialize."""
-        self._station = station
-        self._wallbox = wallbox
-
-        super().__init__(
-            hass,
-            _LOGGER,
-            name=DOMAIN,
-            update_interval=timedelta(seconds=UPDATE_INTERVAL),
-        )
-
-    def _authenticate(self) -> None:
-        """Authenticate using Wallbox API."""
-        try:
-            self._wallbox.authenticate()
-
-        except requests.exceptions.HTTPError as wallbox_connection_error:
-            if wallbox_connection_error.response.status_code == HTTPStatus.FORBIDDEN:
-                raise ConfigEntryAuthFailed from wallbox_connection_error
-            raise ConnectionError from wallbox_connection_error
-
-    def _validate(self) -> None:
-        """Authenticate using Wallbox API."""
-        try:
-            self._wallbox.authenticate()
-        except requests.exceptions.HTTPError as wallbox_connection_error:
-            if wallbox_connection_error.response.status_code == 403:
-                raise InvalidAuth from wallbox_connection_error
-            raise ConnectionError from wallbox_connection_error
-
-    async def async_validate_input(self) -> None:
-        """Get new sensor data for Wallbox component."""
-        await self.hass.async_add_executor_job(self._validate)
-
-    def _get_data(self) -> dict[str, Any]:
-        """Get new sensor data for Wallbox component."""
-        try:
-            self._authenticate()
-            data: dict[str, Any] = self._wallbox.getChargerStatus(self._station)
-            data[CHARGER_MAX_CHARGING_CURRENT_KEY] = data[CHARGER_DATA_KEY][
-                CHARGER_MAX_CHARGING_CURRENT_KEY
-            ]
-            data[CHARGER_LOCKED_UNLOCKED_KEY] = data[CHARGER_DATA_KEY][
-                CHARGER_LOCKED_UNLOCKED_KEY
-            ]
-            data[CHARGER_ENERGY_PRICE_KEY] = data[CHARGER_DATA_KEY][
-                CHARGER_ENERGY_PRICE_KEY
-            ]
-            data[
-                CHARGER_CURRENCY_KEY
-            ] = f"{data[CHARGER_DATA_KEY][CHARGER_CURRENCY_KEY][CODE_KEY]}/kWh"
-
-            data[CHARGER_STATUS_DESCRIPTION_KEY] = CHARGER_STATUS.get(
-                data[CHARGER_STATUS_ID_KEY], ChargerStatus.UNKNOWN
-            )
-            return data
-        except (
-            ConnectionError,
-            requests.exceptions.HTTPError,
-        ) as wallbox_connection_error:
-            raise UpdateFailed from wallbox_connection_error
-
-    async def _async_update_data(self) -> dict[str, Any]:
-        """Get new sensor data for Wallbox component."""
-        return await self.hass.async_add_executor_job(self._get_data)
-
-    def _set_charging_current(self, charging_current: float) -> None:
-        """Set maximum charging current for Wallbox."""
-        try:
-            self._authenticate()
-            self._wallbox.setMaxChargingCurrent(self._station, charging_current)
-        except requests.exceptions.HTTPError as wallbox_connection_error:
-            if wallbox_connection_error.response.status_code == 403:
-                raise InvalidAuth from wallbox_connection_error
-            raise ConnectionError from wallbox_connection_error
-
-    async def async_set_charging_current(self, charging_current: float) -> None:
-        """Set maximum charging current for Wallbox."""
-        await self.hass.async_add_executor_job(
-            self._set_charging_current, charging_current
-        )
-        await self.async_request_refresh()
-
-    def _set_lock_unlock(self, lock: bool) -> None:
-        """Set wallbox to locked or unlocked."""
-        try:
-            self._authenticate()
-            if lock:
-                self._wallbox.lockCharger(self._station)
-            else:
-                self._wallbox.unlockCharger(self._station)
-        except requests.exceptions.HTTPError as wallbox_connection_error:
-            if wallbox_connection_error.response.status_code == 403:
-                raise InvalidAuth from wallbox_connection_error
-            raise ConnectionError from wallbox_connection_error
-
-    async def async_set_lock_unlock(self, lock: bool) -> None:
-        """Set wallbox to locked or unlocked."""
-        await self.hass.async_add_executor_job(self._set_lock_unlock, lock)
-        await self.async_request_refresh()
-
-    def _pause_charger(self, pause: bool) -> None:
-        """Set wallbox to pause or resume."""
-        try:
-            self._authenticate()
-            if pause:
-                self._wallbox.pauseChargingSession(self._station)
-            else:
-                self._wallbox.resumeChargingSession(self._station)
-        except requests.exceptions.HTTPError as wallbox_connection_error:
-            if wallbox_connection_error.response.status_code == 403:
-                raise InvalidAuth from wallbox_connection_error
-            raise ConnectionError from wallbox_connection_error
-
-    async def async_pause_charger(self, pause: bool) -> None:
-        """Set wallbox to pause or resume."""
-        await self.hass.async_add_executor_job(self._pause_charger, pause)
-        await self.async_request_refresh()
 
 
 async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
@@ -227,7 +49,3 @@ async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
         hass.data[DOMAIN].pop(entry.entry_id)
 
     return unload_ok
-
-
-class InvalidAuth(HomeAssistantError):
-    """Error to indicate there is invalid auth."""

homeassistant/components/wallbox/config_flow.py

@@ -11,8 +11,8 @@ from homeassistant import config_entries, core
 from homeassistant.const import CONF_PASSWORD, CONF_USERNAME
 from homeassistant.data_entry_flow import FlowResult
 
-from . import InvalidAuth, WallboxCoordinator
 from .const import CONF_STATION, DOMAIN
+from .coordinator import InvalidAuth, WallboxCoordinator
 
 COMPONENT_DOMAIN = DOMAIN
 

homeassistant/components/wallbox/const.py

@@ -2,6 +2,7 @@
 from enum import StrEnum
 
 DOMAIN = "wallbox"
+UPDATE_INTERVAL = 30
 
 BIDIRECTIONAL_MODEL_PREFIXES = ["QSX"]
 
@@ -55,3 +56,37 @@ class ChargerStatus(StrEnum):
     WAITING_MID_SAFETY = "Waiting MID safety margin exceeded"
     WAITING_IN_QUEUE_ECO_SMART = "Waiting in queue by Eco-Smart"
     UNKNOWN = "Unknown"
+
+
+# Translation of StatusId based on Wallbox portal code:
+# https://my.wallbox.com/src/utilities/charger/chargerStatuses.js
+CHARGER_STATUS: dict[int, ChargerStatus] = {
+    0: ChargerStatus.DISCONNECTED,
+    14: ChargerStatus.ERROR,
+    15: ChargerStatus.ERROR,
+    161: ChargerStatus.READY,
+    162: ChargerStatus.READY,
+    163: ChargerStatus.DISCONNECTED,
+    164: ChargerStatus.WAITING,
+    165: ChargerStatus.LOCKED,
+    166: ChargerStatus.UPDATING,
+    177: ChargerStatus.SCHEDULED,
+    178: ChargerStatus.PAUSED,
+    179: ChargerStatus.SCHEDULED,
+    180: ChargerStatus.WAITING_FOR_CAR,
+    181: ChargerStatus.WAITING_FOR_CAR,
+    182: ChargerStatus.PAUSED,
+    183: ChargerStatus.WAITING_IN_QUEUE_POWER_SHARING,
+    184: ChargerStatus.WAITING_IN_QUEUE_POWER_SHARING,
+    185: ChargerStatus.WAITING_IN_QUEUE_POWER_BOOST,
+    186: ChargerStatus.WAITING_IN_QUEUE_POWER_BOOST,
+    187: ChargerStatus.WAITING_MID_FAILED,
+    188: ChargerStatus.WAITING_MID_SAFETY,
+    189: ChargerStatus.WAITING_IN_QUEUE_ECO_SMART,
+    193: ChargerStatus.CHARGING,
+    194: ChargerStatus.CHARGING,
+    195: ChargerStatus.CHARGING,
+    196: ChargerStatus.DISCHARGING,
+    209: ChargerStatus.LOCKED,
+    210: ChargerStatus.LOCKED_CAR_CONNECTED,
+}

homeassistant/components/wallbox/coordinator.py

@@ -0,0 +1,159 @@
+"""DataUpdateCoordinator for the wallbox integration."""
+from __future__ import annotations
+
+from datetime import timedelta
+from http import HTTPStatus
+import logging
+from typing import Any
+
+import requests
+from wallbox import Wallbox
+
+from homeassistant.core import HomeAssistant
+from homeassistant.exceptions import ConfigEntryAuthFailed, HomeAssistantError
+from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
+
+from .const import (
+    CHARGER_CURRENCY_KEY,
+    CHARGER_DATA_KEY,
+    CHARGER_ENERGY_PRICE_KEY,
+    CHARGER_LOCKED_UNLOCKED_KEY,
+    CHARGER_MAX_CHARGING_CURRENT_KEY,
+    CHARGER_STATUS,
+    CHARGER_STATUS_DESCRIPTION_KEY,
+    CHARGER_STATUS_ID_KEY,
+    CODE_KEY,
+    DOMAIN,
+    UPDATE_INTERVAL,
+    ChargerStatus,
+)
+
+_LOGGER = logging.getLogger(__name__)
+
+
+class WallboxCoordinator(DataUpdateCoordinator[dict[str, Any]]):
+    """Wallbox Coordinator class."""
+
+    def __init__(self, station: str, wallbox: Wallbox, hass: HomeAssistant) -> None:
+        """Initialize."""
+        self._station = station
+        self._wallbox = wallbox
+
+        super().__init__(
+            hass,
+            _LOGGER,
+            name=DOMAIN,
+            update_interval=timedelta(seconds=UPDATE_INTERVAL),
+        )
+
+    def _authenticate(self) -> None:
+        """Authenticate using Wallbox API."""
+        try:
+            self._wallbox.authenticate()
+
+        except requests.exceptions.HTTPError as wallbox_connection_error:
+            if wallbox_connection_error.response.status_code == HTTPStatus.FORBIDDEN:
+                raise ConfigEntryAuthFailed from wallbox_connection_error
+            raise ConnectionError from wallbox_connection_error
+
+    def _validate(self) -> None:
+        """Authenticate using Wallbox API."""
+        try:
+            self._wallbox.authenticate()
+        except requests.exceptions.HTTPError as wallbox_connection_error:
+            if wallbox_connection_error.response.status_code == 403:
+                raise InvalidAuth from wallbox_connection_error
+            raise ConnectionError from wallbox_connection_error
+
+    async def async_validate_input(self) -> None:
+        """Get new sensor data for Wallbox component."""
+        await self.hass.async_add_executor_job(self._validate)
+
+    def _get_data(self) -> dict[str, Any]:
+        """Get new sensor data for Wallbox component."""
+        try:
+            self._authenticate()
+            data: dict[str, Any] = self._wallbox.getChargerStatus(self._station)
+            data[CHARGER_MAX_CHARGING_CURRENT_KEY] = data[CHARGER_DATA_KEY][
+                CHARGER_MAX_CHARGING_CURRENT_KEY
+            ]
+            data[CHARGER_LOCKED_UNLOCKED_KEY] = data[CHARGER_DATA_KEY][
+                CHARGER_LOCKED_UNLOCKED_KEY
+            ]
+            data[CHARGER_ENERGY_PRICE_KEY] = data[CHARGER_DATA_KEY][
+                CHARGER_ENERGY_PRICE_KEY
+            ]
+            data[
+                CHARGER_CURRENCY_KEY
+            ] = f"{data[CHARGER_DATA_KEY][CHARGER_CURRENCY_KEY][CODE_KEY]}/kWh"
+
+            data[CHARGER_STATUS_DESCRIPTION_KEY] = CHARGER_STATUS.get(
+                data[CHARGER_STATUS_ID_KEY], ChargerStatus.UNKNOWN
+            )
+            return data
+        except (
+            ConnectionError,
+            requests.exceptions.HTTPError,
+        ) as wallbox_connection_error:
+            raise UpdateFailed from wallbox_connection_error
+
+    async def _async_update_data(self) -> dict[str, Any]:
+        """Get new sensor data for Wallbox component."""
+        return await self.hass.async_add_executor_job(self._get_data)
+
+    def _set_charging_current(self, charging_current: float) -> None:
+        """Set maximum charging current for Wallbox."""
+        try:
+            self._authenticate()
+            self._wallbox.setMaxChargingCurrent(self._station, charging_current)
+        except requests.exceptions.HTTPError as wallbox_connection_error:
+            if wallbox_connection_error.response.status_code == 403:
+                raise InvalidAuth from wallbox_connection_error
+            raise ConnectionError from wallbox_connection_error
+
+    async def async_set_charging_current(self, charging_current: float) -> None:
+        """Set maximum charging current for Wallbox."""
+        await self.hass.async_add_executor_job(
+            self._set_charging_current, charging_current
+        )
+        await self.async_request_refresh()
+
+    def _set_lock_unlock(self, lock: bool) -> None:
+        """Set wallbox to locked or unlocked."""
+        try:
+            self._authenticate()
+            if lock:
+                self._wallbox.lockCharger(self._station)
+            else:
+                self._wallbox.unlockCharger(self._station)
+        except requests.exceptions.HTTPError as wallbox_connection_error:
+            if wallbox_connection_error.response.status_code == 403:
+                raise InvalidAuth from wallbox_connection_error
+            raise ConnectionError from wallbox_connection_error
+
+    async def async_set_lock_unlock(self, lock: bool) -> None:
+        """Set wallbox to locked or unlocked."""
+        await self.hass.async_add_executor_job(self._set_lock_unlock, lock)
+        await self.async_request_refresh()
+
+    def _pause_charger(self, pause: bool) -> None:
+        """Set wallbox to pause or resume."""
+        try:
+            self._authenticate()
+            if pause:
+                self._wallbox.pauseChargingSession(self._station)
+            else:
+                self._wallbox.resumeChargingSession(self._station)
+        except requests.exceptions.HTTPError as wallbox_connection_error:
+            if wallbox_connection_error.response.status_code == 403:
+                raise InvalidAuth from wallbox_connection_error
+            raise ConnectionError from wallbox_connection_error
+
+    async def async_pause_charger(self, pause: bool) -> None:
+        """Set wallbox to pause or resume."""
+        await self.hass.async_add_executor_job(self._pause_charger, pause)
+        await self.async_request_refresh()
+
+
+class InvalidAuth(HomeAssistantError):
+    """Error to indicate there is invalid auth."""

homeassistant/components/wallbox/entity.py

@@ -4,7 +4,6 @@ from __future__ import annotations
 from homeassistant.helpers.device_registry import DeviceInfo
 from homeassistant.helpers.update_coordinator import CoordinatorEntity
 
-from . import WallboxCoordinator
 from .const import (
     CHARGER_CURRENT_VERSION_KEY,
     CHARGER_DATA_KEY,
@@ -14,6 +13,7 @@ from .const import (
     CHARGER_SOFTWARE_KEY,
     DOMAIN,
 )
+from .coordinator import WallboxCoordinator
 
 
 class WallboxEntity(CoordinatorEntity[WallboxCoordinator]):

homeassistant/components/wallbox/lock.py

@@ -9,13 +9,13 @@ from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import PlatformNotReady
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from . import InvalidAuth, WallboxCoordinator
 from .const import (
     CHARGER_DATA_KEY,
     CHARGER_LOCKED_UNLOCKED_KEY,
     CHARGER_SERIAL_NUMBER_KEY,
     DOMAIN,
 )
+from .coordinator import InvalidAuth, WallboxCoordinator
 from .entity import WallboxEntity
 
 LOCK_TYPES: dict[str, LockEntityDescription] = {

homeassistant/components/wallbox/number.py

@@ -13,7 +13,6 @@ from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import PlatformNotReady
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from . import InvalidAuth, WallboxCoordinator
 from .const import (
     BIDIRECTIONAL_MODEL_PREFIXES,
     CHARGER_DATA_KEY,
@@ -23,6 +22,7 @@ from .const import (
     CHARGER_SERIAL_NUMBER_KEY,
     DOMAIN,
 )
+from .coordinator import InvalidAuth, WallboxCoordinator
 from .entity import WallboxEntity
 
 

homeassistant/components/wallbox/sensor.py

@@ -23,7 +23,6 @@ from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.typing import StateType
 
-from . import WallboxCoordinator
 from .const import (
     CHARGER_ADDED_DISCHARGED_ENERGY_KEY,
     CHARGER_ADDED_ENERGY_KEY,
@@ -43,6 +42,7 @@ from .const import (
     CHARGER_STATUS_DESCRIPTION_KEY,
     DOMAIN,
 )
+from .coordinator import WallboxCoordinator
 from .entity import WallboxEntity
 
 CHARGER_STATION = "station"

homeassistant/components/wallbox/switch.py

@@ -8,7 +8,6 @@ from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from . import WallboxCoordinator
 from .const import (
     CHARGER_DATA_KEY,
     CHARGER_PAUSE_RESUME_KEY,
@@ -17,6 +16,7 @@ from .const import (
     DOMAIN,
     ChargerStatus,
 )
+from .coordinator import WallboxCoordinator
 from .entity import WallboxEntity
 
 SWITCH_TYPES: dict[str, SwitchEntityDescription] = {

tests/components/wallbox/test_init.py

@@ -3,7 +3,10 @@ import json
 
 import requests_mock
 
-from homeassistant.components.wallbox import CHARGER_MAX_CHARGING_CURRENT_KEY, DOMAIN
+from homeassistant.components.wallbox.const import (
+    CHARGER_MAX_CHARGING_CURRENT_KEY,
+    DOMAIN,
+)
 from homeassistant.config_entries import ConfigEntryState
 from homeassistant.core import HomeAssistant
 

tests/components/wallbox/test_lock.py

@@ -5,7 +5,7 @@ import pytest
 import requests_mock
 
 from homeassistant.components.lock import SERVICE_LOCK, SERVICE_UNLOCK
-from homeassistant.components.wallbox import CHARGER_LOCKED_UNLOCKED_KEY
+from homeassistant.components.wallbox.const import CHARGER_LOCKED_UNLOCKED_KEY
 from homeassistant.const import ATTR_ENTITY_ID
 from homeassistant.core import HomeAssistant
 

tests/components/wallbox/test_number.py

@@ -5,7 +5,7 @@ import pytest
 import requests_mock
 
 from homeassistant.components.input_number import ATTR_VALUE, SERVICE_SET_VALUE
-from homeassistant.components.wallbox import CHARGER_MAX_CHARGING_CURRENT_KEY
+from homeassistant.components.wallbox.const import CHARGER_MAX_CHARGING_CURRENT_KEY
 from homeassistant.const import ATTR_ENTITY_ID
 from homeassistant.core import HomeAssistant
 

tests/components/wallbox/test_switch.py

@@ -5,8 +5,8 @@ import pytest
 import requests_mock
 
 from homeassistant.components.switch import SERVICE_TURN_OFF, SERVICE_TURN_ON
-from homeassistant.components.wallbox import InvalidAuth
 from homeassistant.components.wallbox.const import CHARGER_STATUS_ID_KEY
+from homeassistant.components.wallbox.coordinator import InvalidAuth
 from homeassistant.const import ATTR_ENTITY_ID
 from homeassistant.core import HomeAssistant