diff --git a/pyproject.toml b/pyproject.toml index 1e9694f4..73042dd8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -26,7 +26,7 @@ dependencies = [ "paho-mqtt>=1.6.1,<3.0.0", "construct>=2.10.57,<3", "vacuum-map-parser-roborock", - "pyrate-limiter>=3.7.0,<4", + "pyrate-limiter>=4.0.0,<5", "aiomqtt>=2.5.0,<3", "click-shell~=2.1", ] diff --git a/roborock/web_api.py b/roborock/web_api.py index e0231766..578e6e72 100644 --- a/roborock/web_api.py +++ b/roborock/web_api.py @@ -12,7 +12,7 @@ import aiohttp from aiohttp import ContentTypeError, FormData -from pyrate_limiter import BucketFullException, Duration, Limiter, Rate +from pyrate_limiter import Duration, Limiter, Rate from roborock import HomeDataSchedule from roborock.data import HomeData, HomeDataRoom, HomeDataScene, ProductResponse, RRiot, UserData @@ -62,7 +62,7 @@ class RoborockApiClient: Rate(40, Duration.DAY), ] - _login_limiter = Limiter(_LOGIN_RATES, max_delay=1000) + _login_limiter = Limiter(_LOGIN_RATES) _home_data_limiter = Limiter(_HOME_DATA_RATES) def __init__( @@ -204,11 +204,8 @@ async def add_device(self, user_data: UserData, s: str, t: str) -> dict: return add_device_response["result"] async def request_code(self) -> None: - try: - await self._login_limiter.try_acquire_async("login") - except BucketFullException as ex: - _LOGGER.info(ex.meta_info) - raise RoborockRateLimit("Reached maximum requests for login. Please try again later.") from ex + if not await self._login_limiter.try_acquire_async("login", blocking=True, timeout=1): + raise RoborockRateLimit("Reached maximum requests for login. Please try again later.") base_url = await self.base_url header_clientid = self._get_header_client_id() code_request = PreparedRequest(base_url, self.session, {"header_clientid": header_clientid}) @@ -238,11 +235,8 @@ async def request_code_v4(self) -> None: if await self.country_code is None or await self.country is None: _LOGGER.info("No country code or country found, trying old version of request code.") return await self.request_code() - try: - await self._login_limiter.try_acquire_async("login") - except BucketFullException as ex: - _LOGGER.info(ex.meta_info) - raise RoborockRateLimit("Reached maximum requests for login. Please try again later.") from ex + if not await self._login_limiter.try_acquire_async("login", blocking=True, timeout=1): + raise RoborockRateLimit("Reached maximum requests for login. Please try again later.") base_url = await self.base_url header_clientid = self._get_header_client_id() code_request = PreparedRequest( @@ -370,11 +364,8 @@ async def code_login_v4( return UserData.from_dict(user_data) async def pass_login(self, password: str) -> UserData: - try: - await self._login_limiter.try_acquire_async("login") - except BucketFullException as ex: - _LOGGER.info(ex.meta_info) - raise RoborockRateLimit("Reached maximum requests for login. Please try again later.") from ex + if not await self._login_limiter.try_acquire_async("login", blocking=True, timeout=1): + raise RoborockRateLimit("Reached maximum requests for login. Please try again later.") base_url = await self.base_url header_clientid = self._get_header_client_id() @@ -468,11 +459,8 @@ async def _get_home_id(self, user_data: UserData): return home_id_response["data"]["rrHomeId"] async def get_home_data(self, user_data: UserData) -> HomeData: - try: - self._home_data_limiter.try_acquire("home_data") - except BucketFullException as ex: - _LOGGER.info(ex.meta_info) - raise RoborockRateLimit("Reached maximum requests for home data. Please try again later.") from ex + if not self._home_data_limiter.try_acquire("home_data", blocking=False): + raise RoborockRateLimit("Reached maximum requests for home data. Please try again later.") rriot = user_data.rriot if rriot is None: raise RoborockException("rriot is none") @@ -497,11 +485,8 @@ async def get_home_data(self, user_data: UserData) -> HomeData: async def get_home_data_v2(self, user_data: UserData) -> HomeData: """This is the same as get_home_data, but uses a different endpoint and includes non-robotic vacuums.""" - try: - self._home_data_limiter.try_acquire("home_data") - except BucketFullException as ex: - _LOGGER.info(ex.meta_info) - raise RoborockRateLimit("Reached maximum requests for home data. Please try again later.") from ex + if not self._home_data_limiter.try_acquire("home_data", blocking=False): + raise RoborockRateLimit("Reached maximum requests for home data. Please try again later.") rriot = user_data.rriot if rriot is None: raise RoborockException("rriot is none") @@ -526,11 +511,8 @@ async def get_home_data_v2(self, user_data: UserData) -> HomeData: async def get_home_data_v3(self, user_data: UserData) -> HomeData: """This is the same as get_home_data, but uses a different endpoint and includes non-robotic vacuums.""" - try: - self._home_data_limiter.try_acquire("home_data") - except BucketFullException as ex: - _LOGGER.info(ex.meta_info) - raise RoborockRateLimit("Reached maximum requests for home data. Please try again later.") from ex + if not self._home_data_limiter.try_acquire("home_data", blocking=False): + raise RoborockRateLimit("Reached maximum requests for home data. Please try again later.") rriot = user_data.rriot home_id = await self._get_home_id(user_data) if rriot.r.a is None: diff --git a/tests/fixtures/web_api_fixtures.py b/tests/fixtures/web_api_fixtures.py index 69a32b3e..f90d5b9f 100644 --- a/tests/fixtures/web_api_fixtures.py +++ b/tests/fixtures/web_api_fixtures.py @@ -1,7 +1,7 @@ import re from collections.abc import Generator from typing import Any -from unittest.mock import patch +from unittest.mock import AsyncMock, patch import pytest from aioresponses import aioresponses @@ -13,8 +13,11 @@ def skip_rate_limit() -> Generator[None, None, None]: """Don't rate limit tests as they aren't actually hitting the api.""" with ( - patch("roborock.web_api.RoborockApiClient._login_limiter.try_acquire"), - patch("roborock.web_api.RoborockApiClient._home_data_limiter.try_acquire"), + patch( + "roborock.web_api.RoborockApiClient._login_limiter.try_acquire_async", + new=AsyncMock(return_value=True), + ), + patch("roborock.web_api.RoborockApiClient._home_data_limiter.try_acquire", return_value=True), ): yield diff --git a/uv.lock b/uv.lock index 06463df6..681a882f 100644 --- a/uv.lock +++ b/uv.lock @@ -1229,11 +1229,11 @@ wheels = [ [[package]] name = "pyrate-limiter" -version = "3.9.0" +version = "4.0.2" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/ec/da/f682c5c5f9f0a5414363eb4397e6b07d84a02cde69c4ceadcbf32c85537c/pyrate_limiter-3.9.0.tar.gz", hash = "sha256:6b882e2c77cda07a241d3730975daea4258344b39c878f1dd8849df73f70b0ce", size = 289308, upload-time = "2025-07-30T14:36:58.659Z" } +sdist = { url = "https://files.pythonhosted.org/packages/03/98/2b3dc1ba6bdf2efaeaa3e102124cbd2636a4ccec241ffeb8a1de207f5cd4/pyrate_limiter-4.0.2.tar.gz", hash = "sha256:b678841e2215f114ef6f98c7093755ca3b466de83cb5a881231fd6e321fa14b5", size = 301304, upload-time = "2026-01-23T09:37:33.612Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/04/af/d8bf0959ece9bc4679bd203908c31019556a421d76d8143b0c6871c7f614/pyrate_limiter-3.9.0-py3-none-any.whl", hash = "sha256:77357840c8cf97a36d67005d4e090787043f54000c12c2b414ff65657653e378", size = 33628, upload-time = "2025-07-30T14:36:57.71Z" }, + { url = "https://files.pythonhosted.org/packages/13/b9/80ffe3f2c34d3247186d74b1d08c1fed1e3ad4127ff6a8a5501b7bf16a97/pyrate_limiter-4.0.2-py3-none-any.whl", hash = "sha256:35ec42b9bb9cfabcafab14d0c5c6523f48378c3da2949e534ce3cbdfea71eadd", size = 36439, upload-time = "2026-01-23T09:37:32.097Z" }, ] [[package]] @@ -1363,7 +1363,7 @@ requires-dist = [ { name = "paho-mqtt", specifier = ">=1.6.1,<3.0.0" }, { name = "pycryptodome", specifier = "~=3.18" }, { name = "pycryptodomex", marker = "sys_platform == 'darwin'", specifier = "~=3.18" }, - { name = "pyrate-limiter", specifier = ">=3.7.0,<4" }, + { name = "pyrate-limiter", specifier = ">=4.0.0,<5" }, { name = "vacuum-map-parser-roborock" }, ]