From 5bb3c7ca55bca73b42d04e850b8a64a61621bb32 Mon Sep 17 00:00:00 2001 From: Hessel Date: Wed, 8 Nov 2023 13:13:11 +0100 Subject: [PATCH] Wallbox Add Authentication Decorator (#102520) --- .../components/wallbox/coordinator.py | 114 +++++++++--------- tests/components/wallbox/test_number.py | 4 +- tests/components/wallbox/test_switch.py | 6 +- 3 files changed, 64 insertions(+), 60 deletions(-) diff --git a/homeassistant/components/wallbox/coordinator.py b/homeassistant/components/wallbox/coordinator.py index b9248d8ce5b..b3c5a9b4910 100644 --- a/homeassistant/components/wallbox/coordinator.py +++ b/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,15 +102,9 @@ 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 + self._wallbox.authenticate() def _validate(self) -> None: """Authenticate using Wallbox API.""" @@ -101,47 +119,41 @@ 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 - ] - 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: 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 + data[CHARGER_STATUS_DESCRIPTION_KEY] = CHARGER_STATUS.get( + data[CHARGER_STATUS_ID_KEY], ChargerStatus.UNKNOWN + ) + return data 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 + + self._wallbox.setEnergyCost(self._station, energy_cost) 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 + + if pause: + self._wallbox.pauseChargingSession(self._station) + else: + self._wallbox.resumeChargingSession(self._station) async def async_pause_charger(self, pause: bool) -> None: """Set wallbox to pause or resume.""" diff --git a/tests/components/wallbox/test_number.py b/tests/components/wallbox/test_number.py index 41ebedc91da..738b9bf7bd6 100644 --- a/tests/components/wallbox/test_number.py +++ b/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, diff --git a/tests/components/wallbox/test_switch.py b/tests/components/wallbox/test_switch.py index 9418b4d8765..edd85c6ccc7 100644 --- a/tests/components/wallbox/test_switch.py +++ b/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,