Wallbox Add Authentication Decorator (#102520)

homeassistant/components/wallbox/coordinator.py

@@ -1,17 +1,18 @@
 """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
 
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import ConfigEntryAuthFailed, HomeAssistantError
-from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
+from homeassistant.helpers.update_coordinator import DataUpdateCoordinator
 
 from .const import (
     CHARGER_CURRENCY_KEY,
@@ -62,6 +63,29 @@ CHARGER_STATUS: dict[int, ChargerStatus] = {
     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,16 +102,10 @@ class WallboxCoordinator(DataUpdateCoordinator[dict[str, Any]]):
             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
-
     def _validate(self) -> None:
         """Authenticate using Wallbox API."""
         try:
@@ -101,10 +119,9 @@ class WallboxCoordinator(DataUpdateCoordinator[dict[str, Any]]):
         """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
@@ -123,25 +140,20 @@ class WallboxCoordinator(DataUpdateCoordinator[dict[str, Any]]):
             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)
 
+    @_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:
                 raise InvalidAuth from wallbox_connection_error
-            raise ConnectionError from wallbox_connection_error
+            raise wallbox_connection_error
 
     async def async_set_charging_current(self, charging_current: float) -> None:
         """Set maximum charging current for Wallbox."""
@@ -150,25 +162,21 @@ class WallboxCoordinator(DataUpdateCoordinator[dict[str, Any]]):
         )
         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:
-                raise InvalidAuth from wallbox_connection_error
-            raise ConnectionError from wallbox_connection_error
 
     async def async_set_energy_cost(self, energy_cost: float) -> None:
         """Set energy cost for Wallbox."""
         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:
@@ -176,25 +184,21 @@ class WallboxCoordinator(DataUpdateCoordinator[dict[str, Any]]):
         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
+            raise 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()
 
+    @_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:
             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."""

tests/components/wallbox/test_number.py

@@ -9,9 +9,9 @@ from homeassistant.components.wallbox.const import (
     CHARGER_ENERGY_PRICE_KEY,
     CHARGER_MAX_CHARGING_CURRENT_KEY,
 )
-from homeassistant.components.wallbox.coordinator import InvalidAuth
 from homeassistant.const import ATTR_ENTITY_ID
 from homeassistant.core import HomeAssistant
+from homeassistant.exceptions import ConfigEntryAuthFailed
 
 from . import (
     authorisation_response,
@@ -186,7 +186,7 @@ async def test_wallbox_number_class_energy_price_auth_error(
             status_code=403,
         )
 
-        with pytest.raises(InvalidAuth):
+        with pytest.raises(ConfigEntryAuthFailed):
             await hass.services.async_call(
                 "number",
                 SERVICE_SET_VALUE,

tests/components/wallbox/test_switch.py

@@ -6,9 +6,9 @@ import requests_mock
 
 from homeassistant.components.switch import SERVICE_TURN_OFF, SERVICE_TURN_ON
 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
+from homeassistant.exceptions import ConfigEntryAuthFailed
 
 from . import authorisation_response, setup_integration
 from .const import MOCK_SWITCH_ENTITY_ID
@@ -120,7 +120,7 @@ async def test_wallbox_switch_class_authentication_error(
             status_code=403,
         )
 
-        with pytest.raises(InvalidAuth):
+        with pytest.raises(ConfigEntryAuthFailed):
             await hass.services.async_call(
                 "switch",
                 SERVICE_TURN_ON,
@@ -129,7 +129,7 @@ async def test_wallbox_switch_class_authentication_error(
                 },
                 blocking=True,
             )
-        with pytest.raises(InvalidAuth):
+        with pytest.raises(ConfigEntryAuthFailed):
             await hass.services.async_call(
                 "switch",
                 SERVICE_TURN_OFF,