From c8df218528daa44a4ce7acfef8d8d1094d9ce781 Mon Sep 17 00:00:00 2001 From: Arjun Ray Date: Sun, 21 Sep 2025 02:33:32 -0400 Subject: [PATCH 1/2] feat: add method for returning country and country code --- roborock/web_api.py | 65 +++++++++++++++++++++++++++---------------- tests/conftest.py | 4 +-- tests/test_web_api.py | 7 +++++ 3 files changed, 50 insertions(+), 26 deletions(-) diff --git a/roborock/web_api.py b/roborock/web_api.py index 79012a23..1124d362 100644 --- a/roborock/web_api.py +++ b/roborock/web_api.py @@ -57,32 +57,41 @@ def __init__(self, username: str, base_url=None, session: aiohttp.ClientSession self._device_identifier = secrets.token_urlsafe(16) self.session = session + async def _get_info_by_email(self, url: str) -> dict: + url_request = PreparedRequest(url, self.session) + response = await url_request.request( + "post", + "/api/v1/getUrlByEmail", + params={"email": self._username, "needtwostepauth": "false"}, + ) + if response is None: + raise RoborockUrlException("Get url by email returned None") + response_code = response.get("code") + if response_code != 200: + _LOGGER.info("Get base url failed for %s with the following context: %s", self._username, response) + if response_code == 2003: + raise RoborockInvalidEmail("Your email was incorrectly formatted.") + elif response_code == 1001: + raise RoborockMissingParameters( + "You are missing parameters for this request, are you sure you entered your username?" + ) + elif response_code == 9002: + raise RoborockTooManyRequest("Please temporarily disable making requests and try again later.") + raise RoborockUrlException(f"error code: {response_code} msg: {response.get('error')}") + response_data = response.get("data") + if response_data is None: + raise RoborockUrlException("response does not have 'data'") + + return { + "base_url": response_data.get("url"), + "country": response_data.get("country"), + "countrycode": response_data.get("countrycode"), + } + async def _get_base_url(self) -> str: if not self.base_url: - url_request = PreparedRequest(self._default_url, self.session) - response = await url_request.request( - "post", - "/api/v1/getUrlByEmail", - params={"email": self._username, "needtwostepauth": "false"}, - ) - if response is None: - raise RoborockUrlException("get url by email returned None") - response_code = response.get("code") - if response_code != 200: - _LOGGER.info("Get base url failed for %s with the following context: %s", self._username, response) - if response_code == 2003: - raise RoborockInvalidEmail("Your email was incorrectly formatted.") - elif response_code == 1001: - raise RoborockMissingParameters( - "You are missing parameters for this request, are you sure you entered your username?" - ) - elif response_code == 9002: - raise RoborockTooManyRequest("Please temporarily disable making requests and try again later.") - raise RoborockUrlException(f"error code: {response_code} msg: {response.get('error')}") - response_data = response.get("data") - if response_data is None: - raise RoborockUrlException("response does not have 'data'") - self.base_url = response_data.get("url") + data = await self._get_info_by_email(self._default_url) + self.base_url = data["base_url"] return self.base_url def _get_header_client_id(self): @@ -161,6 +170,14 @@ async def add_device(self, user_data: UserData, s: str, t: str) -> dict: return add_device_response["result"] + async def get_country_code_and_country(self) -> dict: + """Get country and country code for user.""" + + base_url = await self._get_base_url() + data = await self._get_info_by_email(base_url) + + return {"country": data["country"], "countrycode": data["countrycode"]} + async def request_code(self) -> None: try: self._login_limiter.try_acquire("login") diff --git a/tests/conftest.py b/tests/conftest.py index a38d429f..ff4a6cb8 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -16,7 +16,7 @@ from roborock.roborock_message import RoborockMessage from roborock.version_1_apis.roborock_local_client_v1 import RoborockLocalClientV1 from roborock.version_1_apis.roborock_mqtt_client_v1 import RoborockMqttClientV1 -from tests.mock_data import HOME_DATA_RAW, HOME_DATA_SCENES_RAW, TEST_LOCAL_API_HOST, USER_DATA +from tests.mock_data import BASE_URL, HOME_DATA_RAW, HOME_DATA_SCENES_RAW, TEST_LOCAL_API_HOST, USER_DATA _LOGGER = logging.getLogger(__name__) @@ -177,7 +177,7 @@ def mock_rest() -> aioresponses: status=200, payload={ "code": 200, - "data": {"country": "US", "countrycode": "1", "url": "https://usiot.roborock.com"}, + "data": {"country": USER_DATA["country"], "countrycode": USER_DATA["countrycode"], "url": BASE_URL}, "msg": "success", }, ) diff --git a/tests/test_web_api.py b/tests/test_web_api.py index 1a11f200..b409a879 100644 --- a/tests/test_web_api.py +++ b/tests/test_web_api.py @@ -71,3 +71,10 @@ async def test_code_login_v4_flow(mock_rest) -> None: await api.request_code_v4() ud = await api.code_login_v4(4123, "US", 1) assert ud == UserData.from_dict(USER_DATA) + + +async def test_get_country_code_and_country(mock_rest) -> None: + """Test that we can login with a code and we get back the correct userdata object.""" + api = RoborockApiClient(username="test_user@gmail.com", base_url="https://euiot.roborock.com") + res = await api.get_country_code_and_country() + assert res == {"country": USER_DATA["country"], "countrycode": USER_DATA["countrycode"]} From 6ca24b0c8397b995dee94eb5e73da72fe50a0415 Mon Sep 17 00:00:00 2001 From: Arjun Ray Date: Sun, 21 Sep 2025 18:06:11 -0400 Subject: [PATCH 2/2] fix: add return type --- roborock/web_api.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/roborock/web_api.py b/roborock/web_api.py index 1124d362..52e307f4 100644 --- a/roborock/web_api.py +++ b/roborock/web_api.py @@ -170,7 +170,7 @@ async def add_device(self, user_data: UserData, s: str, t: str) -> dict: return add_device_response["result"] - async def get_country_code_and_country(self) -> dict: + async def get_country_code_and_country(self) -> dict[str, str]: """Get country and country code for user.""" base_url = await self._get_base_url()