Make V02 authentication failures recoverable

Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>

Author: hugo-brito

custom_components/bestway/__init__.py

@@ -5,7 +5,7 @@
 from datetime import datetime, timedelta
 from logging import getLogger
 
-from aiohttp import ClientSession
+from aiohttp import ClientError, ClientSession
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import Platform
 from homeassistant.core import HomeAssistant
@@ -158,7 +158,11 @@ async def _async_setup_aws_iot(
     hass: HomeAssistant, entry: ConfigEntry, session: ClientSession
 ) -> bool:
     """Set up AWS IoT V02 backend."""
-    from .aws_iot.api import AwsIotApi, AwsIotAuthException
+    from .aws_iot.api import (
+        AwsIotApi,
+        AwsIotAuthException,
+        AwsIotConnectionError,
+    )
     from .aws_iot.websocket import AwsIotWebSocket
 
     visitor_id = entry.data["visitor_id"]
@@ -193,7 +197,10 @@ async def _async_setup_aws_iot(
         hass.config_entries.async_update_entry(
             entry, data={**entry.data, "token": token}
         )
-        api._token = token
+        api.update_token(token)
+    except (AwsIotConnectionError, TimeoutError, ClientError) as ex:
+        _LOGGER.warning("AWS IoT authentication service unavailable: %s", ex)
+        raise ConfigEntryNotReady from ex
     except AwsIotAuthException as ex:
         _LOGGER.error("AWS IoT authentication failed: %s", ex)
         raise ConfigEntryAuthFailed from ex
@@ -212,10 +219,12 @@ async def token_refresh_callback() -> str:
                     new_token = await AwsIotApi.authenticate(
                         session, visitor_id, location, api_base
                     )
-                    api._token = new_token
+                    api.update_token(new_token)
                     hass.config_entries.async_update_entry(
                         entry, data={**entry.data, "token": new_token}
                     )
+                    for websocket in coordinator.websockets:
+                        websocket.update_token(new_token)
                     return new_token
 
                 ws = AwsIotWebSocket(

custom_components/bestway/aws_iot/api.py

@@ -18,7 +18,7 @@
 from time import time
 from typing import Any
 
-from aiohttp import ClientSession
+from aiohttp import ClientError, ClientSession
 
 from .encryption import encrypt_command_payload
 from ..bestway.model import (
@@ -34,7 +34,7 @@
 DEFAULT_API_BASE = "https://smarthub-eu.bestwaycorp.com"  # EU endpoint
 APP_ID = "AhFLL54HnChhrxcl9ZUJL6QNfolTIB"
 APP_SECRET = "4ECvVs13enL5AiYSmscNjvlaisklQDz7vWPCCWXcEFjhWfTmLT"
-TIMEOUT = 10
+TIMEOUT = 20
 
 # Regional API endpoints (from ServiceConfig.java)
 API_ENDPOINTS = {
@@ -53,6 +53,10 @@ class AwsIotAuthException(AwsIotException):
     """Authentication error."""
 
 
+class AwsIotConnectionError(AwsIotException):
+    """Transient connection error."""
+
+
 class AwsIotApi:
     """AWS IoT API client matching Gizwits BestwayApi interface.
 
@@ -215,7 +219,8 @@ async def authenticate(
             Authentication token
 
         Raises:
-            AwsIotAuthException: If authentication fails
+            AwsIotAuthException: If authentication is rejected
+            AwsIotConnectionError: If the authentication service cannot be reached
         """
         import random
         import string
@@ -271,20 +276,32 @@ async def authenticate(
         _LOGGER.debug("Sign in headers: %s", "sign" in headers)
         _LOGGER.debug("All header keys: %s", list(headers.keys()))
 
-        async with asyncio.timeout(TIMEOUT):
-            async with session.post(
-                url, headers=headers, json=payload, ssl=False
-            ) as resp:
-                data = await resp.json()
-                _LOGGER.debug("Auth response: %s", data)
-                _LOGGER.debug("Response status: %s", resp.status)
-                token = data.get("data", {}).get("token")
-
-                if not token:
-                    _LOGGER.error("No token in response. Full response: %s", data)
-                    raise AwsIotAuthException("No token in authentication response")
-
-                return str(token)
+        try:
+            async with asyncio.timeout(TIMEOUT):
+                async with session.post(
+                    url, headers=headers, json=payload, ssl=False
+                ) as resp:
+                    data = await resp.json()
+                    _LOGGER.debug("Auth response: %s", data)
+                    _LOGGER.debug("Response status: %s", resp.status)
+
+                    if resp.status in (401, 403):
+                        raise AwsIotAuthException("Authentication rejected")
+
+                    token = data.get("data", {}).get("token")
+                    if not token:
+                        _LOGGER.error("No token in response. Full response: %s", data)
+                        raise AwsIotAuthException("No token in authentication response")
+
+                    return str(token)
+        except (TimeoutError, ClientError) as err:
+            raise AwsIotConnectionError(
+                "Unable to reach the authentication service"
+            ) from err
+
+    def update_token(self, token: str) -> None:
+        """Replace the token used by subsequent API requests."""
+        self._token = token
 
     @staticmethod
     async def bind_qr_code(

custom_components/bestway/aws_iot/websocket.py

@@ -93,6 +93,10 @@ def ws_url(self) -> str:
             )
         return endpoint
 
+    def update_token(self, token: str) -> None:
+        """Replace the token used by the next connection attempt."""
+        self._token = token
+
     async def connect(self) -> None:
         """Connect to region-specific WebSocket endpoint.
 

custom_components/bestway/config_flow.py

@@ -19,7 +19,12 @@
 from homeassistant.helpers.aiohttp_client import async_get_clientsession
 import voluptuous as vol
 
-from .aws_iot.api import AwsIotAuthException
+from .aws_iot.api import (
+    API_ENDPOINTS,
+    AwsIotApi,
+    AwsIotAuthException,
+    AwsIotConnectionError,
+)
 from .bestway.api import (
     BestwayApi,
     BestwayIncorrectPasswordException,
@@ -137,6 +142,37 @@ async def async_step_user(
         else:
             return await self.async_step_aws_iot_auth()
 
+    async def async_step_reauth(self, entry_data: dict[str, Any]) -> ConfigFlowResult:
+        """Refresh credentials for an existing AWS IoT entry."""
+        entry = self.hass.config_entries.async_get_entry(self.context["entry_id"])
+        if entry is None or entry.data.get("backend") != BACKEND_AWS_IOT:
+            return self.async_abort(reason="reauth_unsuccessful")
+
+        visitor_id = entry.data["visitor_id"]
+        location = entry.data.get("location", "GB")
+        api_base = entry.data.get("api_base")
+        if not api_base:
+            api_base = API_ENDPOINTS.get(
+                entry.data.get("region", "EU"), API_ENDPOINTS["EU"]
+            )
+
+        try:
+            token = await AwsIotApi.authenticate(
+                async_get_clientsession(self.hass),
+                visitor_id,
+                location,
+                api_base,
+            )
+        except AwsIotConnectionError:
+            return self.async_abort(reason="cannot_connect")
+        except AwsIotAuthException:
+            return self.async_abort(reason="reauth_unsuccessful")
+
+        return self.async_update_reload_and_abort(
+            entry,
+            data_updates={"token": token},
+        )
+
     async def async_step_gizwits_auth(
         self, user_input: dict[str, Any] | None = None
     ) -> ConfigFlowResult:
@@ -239,8 +275,6 @@ async def async_step_aws_iot_auth(
             )
 
         try:
-            from .aws_iot.api import AwsIotApi, API_ENDPOINTS
-
             session = async_get_clientsession(self.hass)
 
             # Map region to API endpoint
@@ -340,6 +374,9 @@ async def async_step_aws_iot_auth(
         except AwsIotAuthException as auth_err:
             _LOGGER.error("AWS IoT authentication failed: %s", auth_err)
             errors["base"] = "auth_failed"
+        except AwsIotConnectionError as connection_err:
+            _LOGGER.error("AWS IoT connection failed: %s", connection_err)
+            errors["base"] = "cannot_connect"
         except Exception:  # pylint: disable=broad-except
             _LOGGER.exception("AWS IoT setup failed")
             errors["base"] = "unknown"

custom_components/bestway/translations/en.json

@@ -27,6 +27,11 @@
         }
       }
     },
+    "abort": {
+      "cannot_connect": "Could not connect to the Bestway API. Home Assistant will retry setup.",
+      "reauth_successful": "Authentication refreshed successfully.",
+      "reauth_unsuccessful": "Authentication could not be refreshed."
+    },
     "error": {
       "cannot_connect": "Could not connect to the Bestway API",
       "user_does_not_exist": "User account does not exist",

tests/test_aws_iot_api.py

@@ -6,6 +6,7 @@
 from custom_components.bestway.aws_iot.api import (
     AwsIotApi,
     AwsIotAuthException,
+    AwsIotConnectionError,
 )
 
 
@@ -19,6 +20,26 @@ def create_mock_response(status: int, json_data: dict):
     return response
 
 
+@pytest.mark.asyncio
+async def test_authenticate_wraps_timeout_as_connection_error(mock_session):
+    """Authentication timeouts are classified as transient connection errors."""
+    mock_session.post = MagicMock(side_effect=TimeoutError)
+
+    with pytest.raises(AwsIotConnectionError):
+        await AwsIotApi.authenticate(mock_session, "test_visitor")
+
+
+@pytest.mark.asyncio
+async def test_authenticate_rejects_missing_token(mock_session):
+    """A successful response without a token is an authentication failure."""
+    mock_session.post = MagicMock(
+        return_value=create_mock_response(200, {"code": 1, "data": {}})
+    )
+
+    with pytest.raises(AwsIotAuthException):
+        await AwsIotApi.authenticate(mock_session, "test_visitor")
+
+
 @pytest.fixture
 def mock_session():
     """Create mock aiohttp ClientSession."""

tests/test_aws_iot_websocket.py

@@ -224,6 +224,13 @@ async def test_region_fallback_to_eu(aws_websocket):
     assert "eu-central-1" in url
 
 
+def test_update_token_changes_next_connection_token(aws_websocket):
+    """Token updates are applied to subsequent WebSocket connections."""
+    aws_websocket.update_token("fresh_token")
+
+    assert aws_websocket._token == "fresh_token"
+
+
 @pytest.mark.asyncio
 async def test_heartbeat_loop_sends_messages(aws_websocket):
     """Test heartbeat loop sends JSON heartbeat and WebSocket ping."""

tests/test_config_flow.py

@@ -7,9 +7,11 @@
 from homeassistant import config_entries
 from homeassistant.data_entry_flow import FlowResultType
 import pytest
+from pytest_homeassistant_custom_component.common import MockConfigEntry
 
 from custom_components.bestway.bestway.model import BestwayUserToken
 from custom_components.bestway.const import (
+    BACKEND_AWS_IOT,
     CONF_API_ROOT,
     CONF_API_ROOT_EU,
     CONF_PASSWORD,
@@ -210,3 +212,39 @@ async def test_backend_selection_shows_both_options(hass):
     # Schema should have backend field with options
     schema_keys = list(result["data_schema"].schema.keys())
     assert any("backend" in str(key) for key in schema_keys)
+
+
+async def test_aws_iot_reauth_refreshes_token(hass):
+    """AWS IoT reauth uses the stored visitor ID without user input."""
+    entry = MockConfigEntry(
+        version=2,
+        domain=DOMAIN,
+        title="Bestway Spa",
+        data={
+            "backend": BACKEND_AWS_IOT,
+            "visitor_id": "test_visitor",
+            "token": "old_token",
+            "region": "EU",
+            "api_base": "https://example.test",
+        },
+        source=config_entries.SOURCE_USER,
+    )
+    entry.add_to_hass(hass)
+
+    with patch(
+        "custom_components.bestway.config_flow.AwsIotApi.authenticate",
+        return_value="new_token",
+    ) as authenticate:
+        result = await hass.config_entries.flow.async_init(
+            DOMAIN,
+            context={
+                "source": config_entries.SOURCE_REAUTH,
+                "entry_id": entry.entry_id,
+            },
+            data=dict(entry.data),
+        )
+
+    assert result["type"] is FlowResultType.ABORT
+    assert result["reason"] == "reauth_successful"
+    assert entry.data["token"] == "new_token"
+    authenticate.assert_awaited_once()