diff --git a/src/pyvesync/vesync.py b/src/pyvesync/vesync.py index 63ffc436..e9298a9c 100644 --- a/src/pyvesync/vesync.py +++ b/src/pyvesync/vesync.py @@ -442,6 +442,8 @@ async def async_call_api( json_object: dict | None | DataClassORJSONMixin = None, headers: dict | None = None, device: VeSyncBaseDevice | None = None, + *, + _reauth_retry: bool = False, ) -> tuple[dict | None, int]: """Make API calls by passing endpoint, header and body. @@ -505,7 +507,13 @@ async def async_call_api( request_body=req_dict, ) resp_dict, status_code = await self._api_response_wrapper( - resp_bytes, resp_status, api, req_dict, device=device + resp_bytes, + resp_status, + api, + req_dict, + device=device, + headers=headers, + reauth_retry=_reauth_retry, ) except ClientResponseError as e: @@ -513,15 +521,40 @@ async def async_call_api( raise return resp_dict, status_code - async def _api_response_wrapper( + def _with_current_credentials(self, request_part: dict | None) -> dict | None: + """Return a copy of a request body or headers with current credentials. + + Requests carry the token and account ID that were current when they + were built, so they must be refreshed before a retry after re-login. + """ + if request_part is None: + return None + refreshed = dict(request_part) + for key in ('token', 'tk'): + if key in refreshed: + refreshed[key] = self.auth.token + for key in ('accountID', 'accountId'): + if key in refreshed: + refreshed[key] = self.auth.account_id + return refreshed + + async def _api_response_wrapper( # noqa: PLR0913 self, response_bytes: bytes | None, status_code: int, endpoint: str, request_body: dict | None, device: VeSyncBaseDevice | None = None, + *, + headers: dict | None = None, + reauth_retry: bool = False, ) -> tuple[dict | None, int]: - """Internal wrapper used by async_call_api.""" + """Internal wrapper used by async_call_api. + + A rejected token triggers one re-login and one retry of the request + with the new credentials. If the retry is rejected as well, + VeSyncTokenError is raised instead of logging in again. + """ resp_dict = LibraryLogger.try_json_loads(response_bytes) if not isinstance(resp_dict, dict): return None, status_code @@ -529,9 +562,14 @@ async def _api_response_wrapper( error_info = Helpers.parse_error_code(resp_dict) if error_info.error_type == ErrorTypes.TOKEN_ERROR: self.enabled = False - if await self._reauthenticate(): + if not reauth_retry and await self._reauthenticate(): return await self.async_call_api( - endpoint, 'post', request_body, device=device + endpoint, + 'post', + self._with_current_credentials(request_body), + headers=self._with_current_credentials(headers), + device=device, + _reauth_retry=True, ) raise VeSyncTokenError(error_info.message) raise_api_errors(error_info) diff --git a/src/tests/test_reauth_retry.py b/src/tests/test_reauth_retry.py new file mode 100644 index 00000000..4bde6bcf --- /dev/null +++ b/src/tests/test_reauth_retry.py @@ -0,0 +1,136 @@ +"""Tests for the re-login and retry that follows a rejected token. + +A rejected token must cause exactly one re-login and one retry, and the retry +must carry the new credentials and the original headers. +""" + +import asyncio + +import orjson +import pytest + +from pyvesync import VeSync +from pyvesync.utils.errors import VeSyncTokenError + +TOKEN_EXPIRED = {'code': -11001000, 'msg': 'token expired'} +SUCCESS = {'code': 0, 'result': {}} +ENDPOINT = '/cloud/v2/deviceManaged/bypassV2' +OLD_TOKEN = 'old-token' +NEW_TOKEN = 'new-token' +NEW_ACCOUNT = 'new-account' + + +class FakeResponse: + """Minimal aiohttp response for async_call_api.""" + + status = 200 + + def __init__(self, body: dict) -> None: + self._body = orjson.dumps(body) + + async def read(self) -> bytes: + return self._body + + async def __aenter__(self): + return self + + async def __aexit__(self, *exc_info: object) -> bool: + return False + + +class FakeSession: + """Session that rejects any request still carrying the old token.""" + + def __init__(self, *, always_reject: bool = False) -> None: + self.always_reject = always_reject + self.requests: list[dict] = [] + + def request(self, method, url, json=None, headers=None, **kwargs): + del method, url, kwargs + self.requests.append({'json': json, 'headers': headers}) + sent = {(json or {}).get('token'), (headers or {}).get('tk')} + if self.always_reject or OLD_TOKEN in sent: + return FakeResponse(TOKEN_EXPIRED) + return FakeResponse(SUCCESS) + + async def close(self) -> None: + return None + + +def make_manager(session: FakeSession): + """Return a manager with old credentials, a re-login log and a fake re-login.""" + manager = VeSync('user@example.com', 'password', session=session) + manager.set_credentials(OLD_TOKEN, 'old-account', 'US', 'US') + logins: list[int] = [] + + async def fake_reauthenticate(self) -> bool: + logins.append(1) + self.set_credentials(NEW_TOKEN, NEW_ACCOUNT, 'US', 'US') + return True + + return manager, logins, fake_reauthenticate + + +@pytest.fixture(name='patch_reauth') +def patch_reauth_fixture(monkeypatch): + """Replace VeSyncAuth.reauthenticate for the duration of a test.""" + + def apply(manager: VeSync, replacement) -> None: + monkeypatch.setattr(type(manager.auth), 'reauthenticate', replacement) + + return apply + + +def test_rejected_token_retries_once_with_new_credentials(patch_reauth): + session = FakeSession() + manager, logins, fake_reauthenticate = make_manager(session) + patch_reauth(manager, fake_reauthenticate) + body = {'token': manager.token, 'accountID': manager.account_id, 'method': 'bypassV2'} + headers = {'User-Agent': 'okhttp/3.12.1'} + + response, status = asyncio.run( + manager.async_call_api(ENDPOINT, 'post', json_object=body, headers=headers) + ) + + assert (response, status) == (SUCCESS, 200) + assert len(logins) == 1 + assert len(session.requests) == 2 + retry = session.requests[1] + assert retry['json']['token'] == NEW_TOKEN + assert retry['json']['accountID'] == NEW_ACCOUNT + assert retry['json']['method'] == 'bypassV2' + assert retry['headers'] == headers + # The caller's dict is left as it was built. + assert body['token'] == OLD_TOKEN + + +def test_credentials_in_headers_are_refreshed(patch_reauth): + session = FakeSession() + manager, logins, fake_reauthenticate = make_manager(session) + patch_reauth(manager, fake_reauthenticate) + headers = {'tk': manager.token, 'accountId': manager.account_id, 'tz': 'UTC'} + + response, _ = asyncio.run( + manager.async_call_api(ENDPOINT, 'post', json_object={}, headers=headers) + ) + + assert response == SUCCESS + assert len(logins) == 1 + assert session.requests[1]['headers'] == { + 'tk': NEW_TOKEN, + 'accountId': NEW_ACCOUNT, + 'tz': 'UTC', + } + + +def test_token_rejected_after_retry_raises_without_another_login(patch_reauth): + session = FakeSession(always_reject=True) + manager, logins, fake_reauthenticate = make_manager(session) + patch_reauth(manager, fake_reauthenticate) + body = {'token': manager.token, 'accountID': manager.account_id} + + with pytest.raises(VeSyncTokenError): + asyncio.run(manager.async_call_api(ENDPOINT, 'post', json_object=body)) + + assert len(logins) == 1 + assert len(session.requests) == 2