From d99044572a1ea6a047975bf91b5c29002cb1963d Mon Sep 17 00:00:00 2001 From: IceBotYT <34712694+IceBotYT@users.noreply.github.com> Date: Fri, 14 Feb 2025 14:03:21 -0500 Subject: [PATCH] Improved auth failure handling in Nice G.O. (#136607) --- .../components/nice_go/coordinator.py | 4 +- homeassistant/components/nice_go/cover.py | 26 ++---- homeassistant/components/nice_go/light.py | 26 ++---- homeassistant/components/nice_go/switch.py | 26 ++---- homeassistant/components/nice_go/util.py | 66 ++++++++++++++ tests/components/nice_go/test_cover.py | 85 ++++++++++++++++++- tests/components/nice_go/test_light.py | 85 ++++++++++++++++++- tests/components/nice_go/test_switch.py | 85 ++++++++++++++++++- 8 files changed, 335 insertions(+), 68 deletions(-) create mode 100644 homeassistant/components/nice_go/util.py diff --git a/homeassistant/components/nice_go/coordinator.py b/homeassistant/components/nice_go/coordinator.py index e486263fbe5..ffdd9dbd518 100644 --- a/homeassistant/components/nice_go/coordinator.py +++ b/homeassistant/components/nice_go/coordinator.py @@ -153,7 +153,7 @@ class NiceGOUpdateCoordinator(DataUpdateCoordinator[dict[str, NiceGODevice]]): ) try: if datetime.now().timestamp() >= expiry_time: - await self._update_refresh_token() + await self.update_refresh_token() else: await self.api.authenticate_refresh( self.refresh_token, async_get_clientsession(self.hass) @@ -178,7 +178,7 @@ class NiceGOUpdateCoordinator(DataUpdateCoordinator[dict[str, NiceGODevice]]): else: self.async_set_updated_data(devices) - async def _update_refresh_token(self) -> None: + async def update_refresh_token(self) -> None: """Update the refresh token with Nice G.O. API.""" _LOGGER.debug("Updating the refresh token with Nice G.O. API") try: diff --git a/homeassistant/components/nice_go/cover.py b/homeassistant/components/nice_go/cover.py index 03124971410..b9b39711a01 100644 --- a/homeassistant/components/nice_go/cover.py +++ b/homeassistant/components/nice_go/cover.py @@ -2,21 +2,17 @@ from typing import Any -from aiohttp import ClientError -from nice_go import ApiError - from homeassistant.components.cover import ( CoverDeviceClass, CoverEntity, CoverEntityFeature, ) from homeassistant.core import HomeAssistant -from homeassistant.exceptions import HomeAssistantError from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback -from .const import DOMAIN from .coordinator import NiceGOConfigEntry from .entity import NiceGOEntity +from .util import retry DEVICE_CLASSES = { "WallStation": CoverDeviceClass.GARAGE, @@ -71,30 +67,18 @@ class NiceGOCoverEntity(NiceGOEntity, CoverEntity): """Return if cover is closing.""" return self.data.barrier_status == "closing" + @retry("close_cover_error") async def async_close_cover(self, **kwargs: Any) -> None: """Close the garage door.""" if self.is_closed: return - try: - await self.coordinator.api.close_barrier(self._device_id) - except (ApiError, ClientError) as err: - raise HomeAssistantError( - translation_domain=DOMAIN, - translation_key="close_cover_error", - translation_placeholders={"exception": str(err)}, - ) from err + await self.coordinator.api.close_barrier(self._device_id) + @retry("open_cover_error") async def async_open_cover(self, **kwargs: Any) -> None: """Open the garage door.""" if self.is_opened: return - try: - await self.coordinator.api.open_barrier(self._device_id) - except (ApiError, ClientError) as err: - raise HomeAssistantError( - translation_domain=DOMAIN, - translation_key="open_cover_error", - translation_placeholders={"exception": str(err)}, - ) from err + await self.coordinator.api.open_barrier(self._device_id) diff --git a/homeassistant/components/nice_go/light.py b/homeassistant/components/nice_go/light.py index 5b06c02f5db..bf283ed6eff 100644 --- a/homeassistant/components/nice_go/light.py +++ b/homeassistant/components/nice_go/light.py @@ -3,23 +3,19 @@ import logging from typing import TYPE_CHECKING, Any -from aiohttp import ClientError -from nice_go import ApiError - from homeassistant.components.light import ColorMode, LightEntity from homeassistant.const import Platform from homeassistant.core import HomeAssistant -from homeassistant.exceptions import HomeAssistantError from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback from .const import ( - DOMAIN, KNOWN_UNSUPPORTED_DEVICE_TYPES, SUPPORTED_DEVICE_TYPES, UNSUPPORTED_DEVICE_WARNING, ) from .coordinator import NiceGOConfigEntry from .entity import NiceGOEntity +from .util import retry _LOGGER = logging.getLogger(__name__) @@ -63,26 +59,14 @@ class NiceGOLightEntity(NiceGOEntity, LightEntity): assert self.data.light_status is not None return self.data.light_status + @retry("light_on_error") async def async_turn_on(self, **kwargs: Any) -> None: """Turn on the light.""" - try: - await self.coordinator.api.light_on(self._device_id) - except (ApiError, ClientError) as error: - raise HomeAssistantError( - translation_domain=DOMAIN, - translation_key="light_on_error", - translation_placeholders={"exception": str(error)}, - ) from error + await self.coordinator.api.light_on(self._device_id) + @retry("light_off_error") async def async_turn_off(self, **kwargs: Any) -> None: """Turn off the light.""" - try: - await self.coordinator.api.light_off(self._device_id) - except (ApiError, ClientError) as error: - raise HomeAssistantError( - translation_domain=DOMAIN, - translation_key="light_off_error", - translation_placeholders={"exception": str(error)}, - ) from error + await self.coordinator.api.light_off(self._device_id) diff --git a/homeassistant/components/nice_go/switch.py b/homeassistant/components/nice_go/switch.py index e81ea489d2f..f043a23eab5 100644 --- a/homeassistant/components/nice_go/switch.py +++ b/homeassistant/components/nice_go/switch.py @@ -5,23 +5,19 @@ from __future__ import annotations import logging from typing import TYPE_CHECKING, Any -from aiohttp import ClientError -from nice_go import ApiError - from homeassistant.components.switch import SwitchDeviceClass, SwitchEntity from homeassistant.const import Platform from homeassistant.core import HomeAssistant -from homeassistant.exceptions import HomeAssistantError from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback from .const import ( - DOMAIN, KNOWN_UNSUPPORTED_DEVICE_TYPES, SUPPORTED_DEVICE_TYPES, UNSUPPORTED_DEVICE_WARNING, ) from .coordinator import NiceGOConfigEntry from .entity import NiceGOEntity +from .util import retry _LOGGER = logging.getLogger(__name__) @@ -65,26 +61,14 @@ class NiceGOSwitchEntity(NiceGOEntity, SwitchEntity): assert self.data.vacation_mode is not None return self.data.vacation_mode + @retry("switch_on_error") async def async_turn_on(self, **kwargs: Any) -> None: """Turn the switch on.""" - try: - await self.coordinator.api.vacation_mode_on(self.data.id) - except (ApiError, ClientError) as error: - raise HomeAssistantError( - translation_domain=DOMAIN, - translation_key="switch_on_error", - translation_placeholders={"exception": str(error)}, - ) from error + await self.coordinator.api.vacation_mode_on(self.data.id) + @retry("switch_off_error") async def async_turn_off(self, **kwargs: Any) -> None: """Turn the switch off.""" - try: - await self.coordinator.api.vacation_mode_off(self.data.id) - except (ApiError, ClientError) as error: - raise HomeAssistantError( - translation_domain=DOMAIN, - translation_key="switch_off_error", - translation_placeholders={"exception": str(error)}, - ) from error + await self.coordinator.api.vacation_mode_off(self.data.id) diff --git a/homeassistant/components/nice_go/util.py b/homeassistant/components/nice_go/util.py new file mode 100644 index 00000000000..02dee6b0ac1 --- /dev/null +++ b/homeassistant/components/nice_go/util.py @@ -0,0 +1,66 @@ +"""Utilities for Nice G.O.""" + +from collections.abc import Callable, Coroutine +from functools import wraps +from typing import Any, Protocol, runtime_checkable + +from aiohttp import ClientError +from nice_go import ApiError, AuthFailedError + +from homeassistant.exceptions import ConfigEntryAuthFailed, HomeAssistantError +from homeassistant.helpers.update_coordinator import UpdateFailed + +from .const import DOMAIN + + +@runtime_checkable +class _ArgsProtocol(Protocol): + coordinator: Any + hass: Any + + +def retry[_R, **P]( + translation_key: str, +) -> Callable[ + [Callable[P, Coroutine[Any, Any, _R]]], Callable[P, Coroutine[Any, Any, _R]] +]: + """Retry decorator to handle API errors.""" + + def decorator( + func: Callable[P, Coroutine[Any, Any, _R]], + ) -> Callable[P, Coroutine[Any, Any, _R]]: + @wraps(func) + async def wrapper(*args: P.args, **kwargs: P.kwargs): + instance = args[0] + if not isinstance(instance, _ArgsProtocol): + raise TypeError("First argument must have correct attributes") + try: + return await func(*args, **kwargs) + except (ApiError, ClientError) as err: + raise HomeAssistantError( + translation_domain=DOMAIN, + translation_key=translation_key, + translation_placeholders={"exception": str(err)}, + ) from err + except AuthFailedError: + # Try refreshing token and retry + try: + await instance.coordinator.update_refresh_token() + return await func(*args, **kwargs) + except (ApiError, ClientError, UpdateFailed) as err: + raise HomeAssistantError( + translation_domain=DOMAIN, + translation_key=translation_key, + translation_placeholders={"exception": str(err)}, + ) from err + except (AuthFailedError, ConfigEntryAuthFailed) as err: + instance.coordinator.config_entry.async_start_reauth(instance.hass) + raise HomeAssistantError( + translation_domain=DOMAIN, + translation_key=translation_key, + translation_placeholders={"exception": str(err)}, + ) from err + + return wrapper + + return decorator diff --git a/tests/components/nice_go/test_cover.py b/tests/components/nice_go/test_cover.py index f90c2d438b0..542b1717d88 100644 --- a/tests/components/nice_go/test_cover.py +++ b/tests/components/nice_go/test_cover.py @@ -4,7 +4,7 @@ from unittest.mock import AsyncMock from aiohttp import ClientError from freezegun.api import FrozenDateTimeFactory -from nice_go import ApiError +from nice_go import ApiError, AuthFailedError import pytest from syrupy import SnapshotAssertion @@ -154,3 +154,86 @@ async def test_cover_exceptions( {ATTR_ENTITY_ID: entity_id}, blocking=True, ) + + +async def test_auth_failed_error( + hass: HomeAssistant, + mock_nice_go: AsyncMock, + mock_config_entry: MockConfigEntry, +) -> None: + """Test that if an auth failed error occurs, the integration attempts a token refresh and a retry before throwing an error.""" + + await setup_integration(hass, mock_config_entry, [Platform.COVER]) + + def _open_side_effect(*args, **kwargs): + if mock_nice_go.open_barrier.call_count <= 3: + raise AuthFailedError + if mock_nice_go.open_barrier.call_count == 5: + raise AuthFailedError + if mock_nice_go.open_barrier.call_count == 6: + raise ApiError + + def _close_side_effect(*args, **kwargs): + if mock_nice_go.close_barrier.call_count <= 3: + raise AuthFailedError + if mock_nice_go.close_barrier.call_count == 4: + raise ApiError + + mock_nice_go.open_barrier.side_effect = _open_side_effect + mock_nice_go.close_barrier.side_effect = _close_side_effect + + with pytest.raises(HomeAssistantError, match="Error opening the barrier"): + await hass.services.async_call( + COVER_DOMAIN, + SERVICE_OPEN_COVER, + {ATTR_ENTITY_ID: "cover.test_garage_1"}, + blocking=True, + ) + + assert mock_nice_go.authenticate.call_count == 1 + assert mock_nice_go.open_barrier.call_count == 2 + + with pytest.raises(HomeAssistantError, match="Error closing the barrier"): + await hass.services.async_call( + COVER_DOMAIN, + SERVICE_CLOSE_COVER, + {ATTR_ENTITY_ID: "cover.test_garage_2"}, + blocking=True, + ) + + assert mock_nice_go.authenticate.call_count == 2 + assert mock_nice_go.close_barrier.call_count == 2 + + # Try again, but this time the auth failed error should not be raised + + await hass.services.async_call( + COVER_DOMAIN, + SERVICE_OPEN_COVER, + {ATTR_ENTITY_ID: "cover.test_garage_1"}, + blocking=True, + ) + + assert mock_nice_go.authenticate.call_count == 3 + assert mock_nice_go.open_barrier.call_count == 4 + + # One more time but with an ApiError instead of AuthFailed + + with pytest.raises(HomeAssistantError, match="Error opening the barrier"): + await hass.services.async_call( + COVER_DOMAIN, + SERVICE_OPEN_COVER, + {ATTR_ENTITY_ID: "cover.test_garage_1"}, + blocking=True, + ) + + with pytest.raises(HomeAssistantError, match="Error closing the barrier"): + await hass.services.async_call( + COVER_DOMAIN, + SERVICE_CLOSE_COVER, + {ATTR_ENTITY_ID: "cover.test_garage_2"}, + blocking=True, + ) + + assert mock_nice_go.authenticate.call_count == 5 + assert mock_nice_go.open_barrier.call_count == 6 + assert mock_nice_go.close_barrier.call_count == 4 diff --git a/tests/components/nice_go/test_light.py b/tests/components/nice_go/test_light.py index b170a0ee3ab..2bc9de59b2b 100644 --- a/tests/components/nice_go/test_light.py +++ b/tests/components/nice_go/test_light.py @@ -3,7 +3,7 @@ from unittest.mock import AsyncMock from aiohttp import ClientError -from nice_go import ApiError +from nice_go import ApiError, AuthFailedError import pytest from syrupy import SnapshotAssertion @@ -160,3 +160,86 @@ async def test_unsupported_device_type( "Please create an issue with your device model in additional info" in caplog.text ) + + +async def test_auth_failed_error( + hass: HomeAssistant, + mock_nice_go: AsyncMock, + mock_config_entry: MockConfigEntry, +) -> None: + """Test that if an auth failed error occurs, the integration attempts a token refresh and a retry before throwing an error.""" + + await setup_integration(hass, mock_config_entry, [Platform.LIGHT]) + + def _on_side_effect(*args, **kwargs): + if mock_nice_go.light_on.call_count <= 3: + raise AuthFailedError + if mock_nice_go.light_on.call_count == 5: + raise AuthFailedError + if mock_nice_go.light_on.call_count == 6: + raise ApiError + + def _off_side_effect(*args, **kwargs): + if mock_nice_go.light_off.call_count <= 3: + raise AuthFailedError + if mock_nice_go.light_off.call_count == 4: + raise ApiError + + mock_nice_go.light_on.side_effect = _on_side_effect + mock_nice_go.light_off.side_effect = _off_side_effect + + with pytest.raises(HomeAssistantError, match="Error while turning on the light"): + await hass.services.async_call( + LIGHT_DOMAIN, + SERVICE_TURN_ON, + {ATTR_ENTITY_ID: "light.test_garage_1_light"}, + blocking=True, + ) + + assert mock_nice_go.authenticate.call_count == 1 + assert mock_nice_go.light_on.call_count == 2 + + with pytest.raises(HomeAssistantError, match="Error while turning off the light"): + await hass.services.async_call( + LIGHT_DOMAIN, + SERVICE_TURN_OFF, + {ATTR_ENTITY_ID: "light.test_garage_2_light"}, + blocking=True, + ) + + assert mock_nice_go.authenticate.call_count == 2 + assert mock_nice_go.light_off.call_count == 2 + + # Try again, but this time the auth failed error should not be raised + + await hass.services.async_call( + LIGHT_DOMAIN, + SERVICE_TURN_ON, + {ATTR_ENTITY_ID: "light.test_garage_1_light"}, + blocking=True, + ) + + assert mock_nice_go.authenticate.call_count == 3 + assert mock_nice_go.light_on.call_count == 4 + + # One more time but with an ApiError instead of AuthFailed + + with pytest.raises(HomeAssistantError, match="Error while turning on the light"): + await hass.services.async_call( + LIGHT_DOMAIN, + SERVICE_TURN_ON, + {ATTR_ENTITY_ID: "light.test_garage_1_light"}, + blocking=True, + ) + + with pytest.raises(HomeAssistantError, match="Error while turning off the light"): + await hass.services.async_call( + LIGHT_DOMAIN, + SERVICE_TURN_OFF, + {ATTR_ENTITY_ID: "light.test_garage_2_light"}, + blocking=True, + ) + + assert mock_nice_go.authenticate.call_count == 5 + assert mock_nice_go.light_on.call_count == 6 + assert mock_nice_go.light_off.call_count == 4 diff --git a/tests/components/nice_go/test_switch.py b/tests/components/nice_go/test_switch.py index d3a2141eb2b..cab009c5b94 100644 --- a/tests/components/nice_go/test_switch.py +++ b/tests/components/nice_go/test_switch.py @@ -3,7 +3,7 @@ from unittest.mock import AsyncMock from aiohttp import ClientError -from nice_go import ApiError +from nice_go import ApiError, AuthFailedError import pytest from homeassistant.components.switch import ( @@ -88,3 +88,86 @@ async def test_error( {ATTR_ENTITY_ID: entity_id}, blocking=True, ) + + +async def test_auth_failed_error( + hass: HomeAssistant, + mock_nice_go: AsyncMock, + mock_config_entry: MockConfigEntry, +) -> None: + """Test that if an auth failed error occurs, the integration attempts a token refresh and a retry before throwing an error.""" + + await setup_integration(hass, mock_config_entry, [Platform.SWITCH]) + + def _on_side_effect(*args, **kwargs): + if mock_nice_go.vacation_mode_on.call_count <= 3: + raise AuthFailedError + if mock_nice_go.vacation_mode_on.call_count == 5: + raise AuthFailedError + if mock_nice_go.vacation_mode_on.call_count == 6: + raise ApiError + + def _off_side_effect(*args, **kwargs): + if mock_nice_go.vacation_mode_off.call_count <= 3: + raise AuthFailedError + if mock_nice_go.vacation_mode_off.call_count == 4: + raise ApiError + + mock_nice_go.vacation_mode_on.side_effect = _on_side_effect + mock_nice_go.vacation_mode_off.side_effect = _off_side_effect + + with pytest.raises(HomeAssistantError, match="Error while turning on the switch"): + await hass.services.async_call( + SWITCH_DOMAIN, + SERVICE_TURN_ON, + {ATTR_ENTITY_ID: "switch.test_garage_1_vacation_mode"}, + blocking=True, + ) + + assert mock_nice_go.authenticate.call_count == 1 + assert mock_nice_go.vacation_mode_on.call_count == 2 + + with pytest.raises(HomeAssistantError, match="Error while turning off the switch"): + await hass.services.async_call( + SWITCH_DOMAIN, + SERVICE_TURN_OFF, + {ATTR_ENTITY_ID: "switch.test_garage_2_vacation_mode"}, + blocking=True, + ) + + assert mock_nice_go.authenticate.call_count == 2 + assert mock_nice_go.vacation_mode_off.call_count == 2 + + # Try again, but this time the auth failed error should not be raised + + await hass.services.async_call( + SWITCH_DOMAIN, + SERVICE_TURN_ON, + {ATTR_ENTITY_ID: "switch.test_garage_1_vacation_mode"}, + blocking=True, + ) + + assert mock_nice_go.authenticate.call_count == 3 + assert mock_nice_go.vacation_mode_on.call_count == 4 + + # One more time but with an ApiError instead of AuthFailed + + with pytest.raises(HomeAssistantError, match="Error while turning on the switch"): + await hass.services.async_call( + SWITCH_DOMAIN, + SERVICE_TURN_ON, + {ATTR_ENTITY_ID: "switch.test_garage_1_vacation_mode"}, + blocking=True, + ) + + with pytest.raises(HomeAssistantError, match="Error while turning off the switch"): + await hass.services.async_call( + SWITCH_DOMAIN, + SERVICE_TURN_OFF, + {ATTR_ENTITY_ID: "switch.test_garage_2_vacation_mode"}, + blocking=True, + ) + + assert mock_nice_go.authenticate.call_count == 5 + assert mock_nice_go.vacation_mode_on.call_count == 6 + assert mock_nice_go.vacation_mode_off.call_count == 4