Wallbox Add Authentication Decorator

homeassistant/components/wallbox/coordinator.py

@@ -1,10 +1,11 @@
 """DataUpdateCoordinator for the wallbox integration."""
 from __future__ import annotations
 
+from collections.abc import Callable
 from datetime import timedelta
 from http import HTTPStatus
 import logging
-from typing import Any
+from typing import Any, Concatenate, ParamSpec, TypeVar
 
 import requests
 from wallbox import Wallbox
@@ -62,6 +63,29 @@
     210: ChargerStatus.LOCKED_CAR_CONNECTED,
 }
 
+_WallboxCoordinatorT = TypeVar("_WallboxCoordinatorT", bound="WallboxCoordinator")
+_P = ParamSpec("_P")
+
+
+def _require_authentication(
+    func: Callable[Concatenate[_WallboxCoordinatorT, _P], Any]
+) -> Callable[Concatenate[_WallboxCoordinatorT, _P], Any]:
+    """Authenticate with decorator using Wallbox API."""
+
+    def require_authentication(
+        self: _WallboxCoordinatorT, *args: _P.args, **kwargs: _P.kwargs
+    ) -> Any:
+        """Authenticate using Wallbox API."""
+        try:
+            self.authenticate()
+            return func(self, *args, **kwargs)
+        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
+
+    return require_authentication
+
 
 class WallboxCoordinator(DataUpdateCoordinator[dict[str, Any]]):
     """Wallbox Coordinator class."""
@@ -78,15 +102,9 @@ def __init__(self, station: str, wallbox: Wallbox, hass: HomeAssistant) -> None:
             update_interval=timedelta(seconds=UPDATE_INTERVAL),
         )
 
-    def _authenticate(self) -> None:
+    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
+        self._wallbox.authenticate()
 
     def _validate(self) -> None:
         """Authenticate using Wallbox API."""
@@ -101,10 +119,10 @@ async def async_validate_input(self) -> None:
         """Get new sensor data for Wallbox component."""
         await self.hass.async_add_executor_job(self._validate)
 
+    @_require_authentication
     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
@@ -133,10 +151,10 @@ 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)
 
+    @_require_authentication
     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:
@@ -150,10 +168,10 @@ async def async_set_charging_current(self, charging_current: float) -> None:
         )
         await self.async_request_refresh()
 
+    @_require_authentication
     def _set_energy_cost(self, energy_cost: float) -> None:
         """Set energy cost for Wallbox."""
         try:
-            self._authenticate()
             self._wallbox.setEnergyCost(self._station, energy_cost)
         except requests.exceptions.HTTPError as wallbox_connection_error:
             if wallbox_connection_error.response.status_code == 403:
@@ -165,10 +183,10 @@ async def async_set_energy_cost(self, energy_cost: float) -> None:
         await self.hass.async_add_executor_job(self._set_energy_cost, energy_cost)
         await self.async_request_refresh()
 
+    @_require_authentication
     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:
@@ -183,10 +201,10 @@ async def async_set_lock_unlock(self, lock: bool) -> None:
         await self.hass.async_add_executor_job(self._set_lock_unlock, lock)
         await self.async_request_refresh()
 
+    @_require_authentication
     def _pause_charger(self, pause: bool) -> None:
         """Set wallbox to pause or resume."""
         try:
-            self._authenticate()
             if pause:
                 self._wallbox.pauseChargingSession(self._station)
             else: