Harden V02 runtime authentication recovery

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

Author: hugo-brito

custom_components/bestway/__init__.py

@@ -207,10 +207,21 @@ async def _async_setup_aws_iot(
 
     # Initialize coordinator
     coordinator = BestwayUpdateCoordinator(hass, entry, api)
+
+    def token_updated(new_token: str) -> None:
+        nonlocal token
+        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)
+
+    api.set_token_update_callback(token_updated)
     await coordinator.async_config_entry_first_refresh()
 
     # Initialize per-device WebSockets
-    websockets = []
+    websockets = coordinator.websockets
     if api.devices:
         for device_id, device in api.devices.items():
             try:
@@ -220,11 +231,6 @@ async def token_refresh_callback() -> str:
                         session, visitor_id, location, api_base
                     )
                     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(
@@ -257,8 +263,6 @@ async def token_refresh_callback() -> str:
         _LOGGER.warning("No devices found, WebSocket not initialized")
 
     # Store WebSockets list on coordinator
-    coordinator.websockets = websockets
-
     hass.data.setdefault(DOMAIN, {})[entry.entry_id] = coordinator
     await hass.config_entries.async_forward_entry_setups(entry, _PLATFORMS)
 

custom_components/bestway/aws_iot/api.py

@@ -15,6 +15,7 @@
 import hashlib
 import logging
 import secrets
+from collections.abc import Callable
 from time import time
 from typing import Any
 
@@ -98,6 +99,7 @@ def __init__(
 
         # State cache (matches Gizwits interface)
         self._state_cache: dict[str, BestwayDeviceStatus] = {}
+        self._token_update_callback: Callable[[str], None] | None = None
 
     @staticmethod
     def generate_visitor_id() -> str:
@@ -281,13 +283,13 @@ async def authenticate(
                 async with session.post(
                     url, headers=headers, json=payload, ssl=False
                 ) as resp:
+                    if resp.status in (401, 403):
+                        raise AwsIotAuthException("Authentication rejected")
+
                     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)
@@ -302,6 +304,12 @@ async def authenticate(
     def update_token(self, token: str) -> None:
         """Replace the token used by subsequent API requests."""
         self._token = token
+        if self._token_update_callback is not None:
+            self._token_update_callback(token)
+
+    def set_token_update_callback(self, callback: Callable[[str], None] | None) -> None:
+        """Set a callback for persisting and propagating refreshed tokens."""
+        self._token_update_callback = callback
 
     @staticmethod
     async def bind_qr_code(
@@ -420,15 +428,14 @@ async def _do_get(self, path: str) -> dict[str, Any]:
 
         async with asyncio.timeout(TIMEOUT):
             async with self._session.get(url, headers=headers, ssl=False) as response:
-                data = await response.json()
-
                 # Check for errors
-                if response.status in (400, 401):
+                if response.status in (400, 401, 403):
                     raise AwsIotAuthException("Token expired or invalid")
 
                 if response.status != 200:
                     raise AwsIotException(f"API error: {response.status}")
 
+                data = await response.json()
                 return dict(data)
 
     async def _do_post(self, path: str, data: dict[str, Any]) -> dict[str, Any]:
@@ -454,19 +461,17 @@ async def _do_post(self, path: str, data: dict[str, Any]) -> dict[str, Any]:
             async with self._session.post(
                 url, headers=headers, json=data, ssl=False
             ) as response:
-                result = await response.json()
-
-                _LOGGER.debug(
-                    "POST %s response (status=%d): %s", path, response.status, result
-                )
-
                 # Check for errors
-                if response.status in (400, 401):
+                if response.status in (400, 401, 403):
                     raise AwsIotAuthException("Token expired or invalid")
 
                 if response.status != 200:
                     raise AwsIotException(f"API error: {response.status}")
 
+                result = await response.json()
+                _LOGGER.debug(
+                    "POST %s response (status=%d): %s", path, response.status, result
+                )
                 return dict(result)
 
     async def refresh_bindings(self) -> None:
@@ -586,23 +591,10 @@ async def refresh_bindings(self) -> None:
 
             self.devices[device_id] = device
 
-    async def fetch_data(self) -> Any:  # Returns BestwayApiResults
-        """Fetch latest state for all devices.
-
-        Implements the same interface as Gizwits BestwayApi.fetch_data().
-
-        For each device:
-        1. POST /api/device/thing_shadow/ with device_id + product_id
-        2. Parse shadow.state.reported or shadow.state.desired
-        3. Return raw AWS field names (water_temperature, temperature_setting, etc.)
-        4. Store in state cache
-
-        Returns:
-            BestwayApiResults with devices dict
-        """
-        # Import here to avoid circular dependency
-        from ..bestway.api import BestwayApiResults
-
+    async def _poll_all_devices(self) -> tuple[int, bool]:
+        """Poll every device once and return success and auth-failure status."""
+        refreshed = 0
+        auth_failed = False
         for device_id in self.devices:
             try:
                 # Get device metadata
@@ -656,14 +648,22 @@ async def fetch_data(self) -> Any:  # Returns BestwayApiResults
                 self._state_cache[device_id] = BestwayDeviceStatus(
                     timestamp=int(time()), attrs=mapped
                 )
+                refreshed += 1
 
                 _LOGGER.debug(
                     "Fetched state for device %s: %d fields",
                     device_id[:12],
                     len(mapped),
                 )
 
-            except Exception as err:
+            except AwsIotAuthException as err:
+                auth_failed = True
+                _LOGGER.warning(
+                    "Authentication failure fetching device %s: %s",
+                    device_id[:12],
+                    err,
+                )
+            except Exception as err:  # pylint: disable=broad-except
                 _LOGGER.warning(
                     "Failed to fetch state for device %s: %s", device_id[:12], err
                 )
@@ -673,6 +673,36 @@ async def fetch_data(self) -> Any:  # Returns BestwayApiResults
                         timestamp=int(time()), attrs={}
                     )
 
+        return refreshed, auth_failed
+
+    async def fetch_data(self) -> Any:  # Returns BestwayApiResults
+        """Fetch state, refreshing an expired token once before failing."""
+        from homeassistant.exceptions import ConfigEntryAuthFailed
+        from homeassistant.helpers.update_coordinator import UpdateFailed
+
+        from ..bestway.api import BestwayApiResults
+
+        refreshed, auth_failed = await self._poll_all_devices()
+
+        if self.devices and refreshed == 0 and auth_failed:
+            _LOGGER.info("Re-authenticating after auth failure during poll")
+            try:
+                token = await self.authenticate(
+                    self._session, self._visitor_id, self._location, self._api_base
+                )
+            except AwsIotAuthException as err:
+                raise ConfigEntryAuthFailed from err
+            except AwsIotConnectionError as err:
+                raise UpdateFailed(
+                    "Unable to reach Bestway authentication service"
+                ) from err
+
+            self.update_token(token)
+            refreshed, _ = await self._poll_all_devices()
+
+        if self.devices and refreshed == 0:
+            raise UpdateFailed("Unable to refresh any Bestway device state")
+
         return BestwayApiResults(devices=self._state_cache)
 
     async def set_device_state(

custom_components/bestway/translations/en.json

@@ -28,7 +28,7 @@
       }
     },
     "abort": {
-      "cannot_connect": "Could not connect to the Bestway API. Home Assistant will retry setup.",
+      "cannot_connect": "Could not connect to the Bestway API. Try reloading the integration.",
       "reauth_successful": "Authentication refreshed successfully.",
       "reauth_unsuccessful": "Authentication could not be refreshed."
     },

tests/test_aws_iot_api.py

@@ -1,6 +1,6 @@
 """Tests for AWS IoT API client."""
 
-from unittest.mock import AsyncMock, MagicMock
+from unittest.mock import AsyncMock, MagicMock, patch
 import pytest
 
 from custom_components.bestway.aws_iot.api import (
@@ -40,6 +40,19 @@ async def test_authenticate_rejects_missing_token(mock_session):
         await AwsIotApi.authenticate(mock_session, "test_visitor")
 
 
+@pytest.mark.asyncio
+async def test_authenticate_checks_rejection_before_parsing_body(mock_session):
+    """A non-JSON 401 remains an authentication failure."""
+    response = create_mock_response(401, {})
+    response.json.side_effect = Exception("not JSON")
+    mock_session.post = MagicMock(return_value=response)
+
+    with pytest.raises(AwsIotAuthException):
+        await AwsIotApi.authenticate(mock_session, "test_visitor")
+
+    response.json.assert_not_awaited()
+
+
 @pytest.fixture
 def mock_session():
     """Create mock aiohttp ClientSession."""
@@ -323,6 +336,75 @@ async def test_fetch_data_returns_results(aws_api, mock_session):
     assert status.attrs["Tnow"] == 36
 
 
+def _make_aws_device(device_id="device1"):
+    """Build a real AWS IoT device for polling tests."""
+    from custom_components.bestway.bestway.model import BestwayDevice
+
+    return BestwayDevice(
+        protocol_version=2,
+        device_id=device_id,
+        product_name="AIRJET",
+        alias="Test Spa",
+        mcu_soft_version="unknown",
+        mcu_hard_version="unknown",
+        wifi_soft_version="unknown",
+        wifi_hard_version="unknown",
+        is_online=True,
+        backend="aws_iot",
+        product_id="T53NN8",
+    )
+
+
+@pytest.mark.asyncio
+async def test_fetch_data_reauthenticates_and_propagates_token(aws_api):
+    """A poll auth failure refreshes once and propagates the new token."""
+    aws_api.devices = {"device1": _make_aws_device()}
+    shadow = {"code": 0, "data": {"state": {"reported": {"power_state": 1}}}}
+    aws_api._do_post = AsyncMock(side_effect=[AwsIotAuthException("expired"), shadow])
+    token_updated = MagicMock()
+    aws_api.set_token_update_callback(token_updated)
+
+    with patch.object(
+        AwsIotApi, "authenticate", new=AsyncMock(return_value="fresh_token")
+    ):
+        results = await aws_api.fetch_data()
+
+    assert results.devices["device1"].attrs["power"] is True
+    assert aws_api._token == "fresh_token"
+    token_updated.assert_called_once_with("fresh_token")
+
+
+@pytest.mark.asyncio
+async def test_fetch_data_raises_auth_failed_when_reauth_rejected(aws_api):
+    """A rejected runtime reauth starts Home Assistant's reauth handling."""
+    from homeassistant.exceptions import ConfigEntryAuthFailed
+
+    aws_api.devices = {"device1": _make_aws_device()}
+    aws_api._do_post = AsyncMock(side_effect=AwsIotAuthException("expired"))
+
+    with (
+        patch.object(
+            AwsIotApi,
+            "authenticate",
+            new=AsyncMock(side_effect=AwsIotAuthException("rejected")),
+        ),
+        pytest.raises(ConfigEntryAuthFailed),
+    ):
+        await aws_api.fetch_data()
+
+
+@pytest.mark.asyncio
+async def test_fetch_data_raises_update_failed_on_total_connection_failure(aws_api):
+    """A total poll failure makes entities unavailable instead of stale."""
+    from homeassistant.helpers.update_coordinator import UpdateFailed
+
+    aws_api.devices = {"device1": _make_aws_device()}
+    aws_api._do_post = AsyncMock(side_effect=ConnectionError("offline"))
+
+    with pytest.raises(UpdateFailed):
+        await aws_api.fetch_data()
+
+
 @pytest.mark.asyncio
 async def test_set_device_state_sends_command(aws_api, mock_session):
     """Test control command sends encrypted payload."""
@@ -366,3 +448,16 @@ async def test_do_get_handles_401(aws_api, mock_session):
 
     with pytest.raises(AwsIotAuthException):
         await aws_api._do_get("/test")
+
+
+@pytest.mark.asyncio
+async def test_do_post_checks_auth_status_before_parsing_body(aws_api, mock_session):
+    """A non-JSON auth rejection still triggers runtime reauthentication."""
+    response = create_mock_response(401, {})
+    response.json.side_effect = Exception("not JSON")
+    mock_session.post = MagicMock(return_value=response)
+
+    with pytest.raises(AwsIotAuthException):
+        await aws_api._do_post("/test", {})
+
+    response.json.assert_not_awaited()