Improved auth failure handling in Nice G.O. (#136607)

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:

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
 
+    @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

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
 
+    @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

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
 
+    @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

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

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

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

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