From b0db7b432e2c9590a51004b881608fc7d9dfe386 Mon Sep 17 00:00:00 2001 From: Luke Lashley Date: Sun, 16 Mar 2025 15:55:00 -0400 Subject: [PATCH] Move Roborock MapParser to coordinator (#140750) Move MapParser to coordinator --- .../components/roborock/coordinator.py | 34 ++++++++++++++ homeassistant/components/roborock/image.py | 46 +------------------ homeassistant/components/roborock/vacuum.py | 5 +- tests/components/roborock/conftest.py | 2 +- tests/components/roborock/test_image.py | 10 ++-- tests/components/roborock/test_vacuum.py | 4 +- 6 files changed, 48 insertions(+), 53 deletions(-) diff --git a/homeassistant/components/roborock/coordinator.py b/homeassistant/components/roborock/coordinator.py index bf06387b377..2f156545929 100644 --- a/homeassistant/components/roborock/coordinator.py +++ b/homeassistant/components/roborock/coordinator.py @@ -5,6 +5,7 @@ from __future__ import annotations import asyncio from dataclasses import dataclass from datetime import timedelta +import io import logging from propcache.api import cached_property @@ -25,6 +26,10 @@ from roborock.version_1_apis.roborock_local_client_v1 import RoborockLocalClient from roborock.version_1_apis.roborock_mqtt_client_v1 import RoborockMqttClientV1 from roborock.version_a01_apis import RoborockClientA01 from roborock.web_api import RoborockApiClient +from vacuum_map_parser_base.config.color import ColorsPalette +from vacuum_map_parser_base.config.image_config import ImageConfig +from vacuum_map_parser_base.config.size import Sizes +from vacuum_map_parser_roborock.map_data_parser import RoborockMapDataParser from homeassistant.config_entries import ConfigEntry from homeassistant.const import ATTR_CONNECTIONS @@ -38,7 +43,11 @@ from homeassistant.util import slugify from .const import ( A01_UPDATE_INTERVAL, + DEFAULT_DRAWABLES, DOMAIN, + DRAWABLES, + MAP_FILE_FORMAT, + MAP_SCALE, V1_CLOUD_IN_CLEANING_INTERVAL, V1_CLOUD_NOT_CLEANING_INTERVAL, V1_LOCAL_IN_CLEANING_INTERVAL, @@ -127,6 +136,18 @@ class RoborockDataUpdateCoordinator(DataUpdateCoordinator[DeviceProp]): self._user_data = user_data self._api_client = api_client self._is_cloud_api = False + drawables = [ + drawable + for drawable, default_value in DEFAULT_DRAWABLES.items() + if config_entry.options.get(DRAWABLES, {}).get(drawable, default_value) + ] + self.map_parser = RoborockMapDataParser( + ColorsPalette(), + Sizes({k: v * MAP_SCALE for k, v in Sizes.SIZES.items()}), + drawables, + ImageConfig(scale=MAP_SCALE), + [], + ) @cached_property def dock_device_info(self) -> DeviceInfo: @@ -145,6 +166,19 @@ class RoborockDataUpdateCoordinator(DataUpdateCoordinator[DeviceProp]): sw_version=self.roborock_device_info.device.fv, ) + def parse_image(self, map_bytes: bytes) -> bytes | None: + """Parse map_bytes and store it as image bytes.""" + try: + parsed_map = self.map_parser.parse(map_bytes) + except (IndexError, ValueError) as err: + _LOGGER.debug("Exception when parsing map contents: %s", err) + return None + if parsed_map.image is None: + return None + img_byte_arr = io.BytesIO() + parsed_map.image.data.save(img_byte_arr, format=MAP_FILE_FORMAT) + return img_byte_arr.getvalue() + async def _async_setup(self) -> None: """Set up the coordinator.""" # Verify we can communicate locally - if we can't, switch to cloud api diff --git a/homeassistant/components/roborock/image.py b/homeassistant/components/roborock/image.py index 2fb5d644826..b56abaeebdb 100644 --- a/homeassistant/components/roborock/image.py +++ b/homeassistant/components/roborock/image.py @@ -1,16 +1,10 @@ """Support for Roborock image.""" import asyncio -from collections.abc import Callable from datetime import datetime -import io import logging from roborock import RoborockCommand -from vacuum_map_parser_base.config.color import ColorsPalette -from vacuum_map_parser_base.config.image_config import ImageConfig -from vacuum_map_parser_base.config.size import Sizes -from vacuum_map_parser_roborock.map_data_parser import RoborockMapDataParser from homeassistant.components.image import ImageEntity from homeassistant.config_entries import ConfigEntry @@ -20,15 +14,7 @@ from homeassistant.exceptions import HomeAssistantError from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback from homeassistant.util import dt as dt_util -from .const import ( - DEFAULT_DRAWABLES, - DOMAIN, - DRAWABLES, - IMAGE_CACHE_INTERVAL, - MAP_FILE_FORMAT, - MAP_SCALE, - MAP_SLEEP, -) +from .const import DOMAIN, IMAGE_CACHE_INTERVAL, MAP_SLEEP from .coordinator import RoborockConfigEntry, RoborockDataUpdateCoordinator from .entity import RoborockCoordinatedEntityV1 @@ -42,31 +28,6 @@ async def async_setup_entry( ) -> None: """Set up Roborock image platform.""" - drawables = [ - drawable - for drawable, default_value in DEFAULT_DRAWABLES.items() - if config_entry.options.get(DRAWABLES, {}).get(drawable, default_value) - ] - parser = RoborockMapDataParser( - ColorsPalette(), - Sizes({k: v * MAP_SCALE for k, v in Sizes.SIZES.items()}), - drawables, - ImageConfig(scale=MAP_SCALE), - [], - ) - - def parse_image(map_bytes: bytes) -> bytes | None: - try: - parsed_map = parser.parse(map_bytes) - except (IndexError, ValueError) as err: - _LOGGER.debug("Exception when parsing map contents: %s", err) - return None - if parsed_map.image is None: - return None - img_byte_arr = io.BytesIO() - parsed_map.image.data.save(img_byte_arr, format=MAP_FILE_FORMAT) - return img_byte_arr.getvalue() - await asyncio.gather( *(refresh_coordinators(hass, coord) for coord in config_entry.runtime_data.v1) ) @@ -78,7 +39,6 @@ async def async_setup_entry( coord, map_info.flag, map_info.name, - parse_image, ) for coord in config_entry.runtime_data.v1 for map_info in coord.maps.values() @@ -100,14 +60,12 @@ class RoborockMap(RoborockCoordinatedEntityV1, ImageEntity): coordinator: RoborockDataUpdateCoordinator, map_flag: int, map_name: str, - parser: Callable[[bytes], bytes | None], ) -> None: """Initialize a Roborock map.""" RoborockCoordinatedEntityV1.__init__(self, unique_id, coordinator) ImageEntity.__init__(self, coordinator.hass) self.config_entry = config_entry self._attr_name = map_name - self.parser = parser self.map_flag = map_flag self.cached_map = b"" self._attr_entity_category = EntityCategory.DIAGNOSTIC @@ -154,7 +112,7 @@ class RoborockMap(RoborockCoordinatedEntityV1, ImageEntity): ) if ( not isinstance(response[0], bytes) - or (content := self.parser(response[0])) is None + or (content := self.coordinator.parse_image(response[0])) is None ): _LOGGER.debug("Failed to parse map contents: %s", response[0]) raise HomeAssistantError( diff --git a/homeassistant/components/roborock/vacuum.py b/homeassistant/components/roborock/vacuum.py index 59abc888673..db201ff06d2 100644 --- a/homeassistant/components/roborock/vacuum.py +++ b/homeassistant/components/roborock/vacuum.py @@ -6,6 +6,10 @@ from typing import Any from roborock.code_mappings import RoborockStateCode from roborock.roborock_message import RoborockDataProtocol from roborock.roborock_typing import RoborockCommand +from vacuum_map_parser_base.config.color import ColorsPalette +from vacuum_map_parser_base.config.image_config import ImageConfig +from vacuum_map_parser_base.config.size import Sizes +from vacuum_map_parser_roborock.map_data_parser import RoborockMapDataParser import voluptuous as vol from homeassistant.components.vacuum import ( @@ -26,7 +30,6 @@ from .const import ( ) from .coordinator import RoborockConfigEntry, RoborockDataUpdateCoordinator from .entity import RoborockCoordinatedEntityV1 -from .image import ColorsPalette, ImageConfig, RoborockMapDataParser, Sizes STATE_CODE_TO_STATE = { RoborockStateCode.starting: VacuumActivity.IDLE, # "Starting" diff --git a/tests/components/roborock/conftest.py b/tests/components/roborock/conftest.py index 9b3a6633c62..cafac280620 100644 --- a/tests/components/roborock/conftest.py +++ b/tests/components/roborock/conftest.py @@ -110,7 +110,7 @@ def bypass_api_fixture(bypass_api_client_fixture: Any) -> None: return_value=MULTI_MAP_LIST, ), patch( - "homeassistant.components.roborock.image.RoborockMapDataParser.parse", + "homeassistant.components.roborock.coordinator.RoborockMapDataParser.parse", return_value=MAP_DATA, ), patch( diff --git a/tests/components/roborock/test_image.py b/tests/components/roborock/test_image.py index 7d79cf4f6ab..d81f1289fe3 100644 --- a/tests/components/roborock/test_image.py +++ b/tests/components/roborock/test_image.py @@ -65,7 +65,7 @@ async def test_floorplan_image( "homeassistant.components.roborock.image.dt_util.utcnow", return_value=now ), patch( - "homeassistant.components.roborock.image.RoborockMapDataParser.parse", + "homeassistant.components.roborock.coordinator.RoborockMapDataParser.parse", return_value=new_map_data, ) as parse_map, ): @@ -94,7 +94,7 @@ async def test_floorplan_image_failed_parse( # Update image, but get none for parse image. with ( patch( - "homeassistant.components.roborock.image.RoborockMapDataParser.parse", + "homeassistant.components.roborock.coordinator.RoborockMapDataParser.parse", return_value=map_data, ), patch( @@ -148,7 +148,7 @@ async def test_fail_to_load_image( """Test that we gracefully handle failing to load an image.""" with ( patch( - "homeassistant.components.roborock.image.RoborockMapDataParser.parse", + "homeassistant.components.roborock.coordinator.RoborockMapDataParser.parse", ) as parse_map, patch( "homeassistant.components.roborock.roborock_storage.Path.exists", @@ -178,7 +178,7 @@ async def test_fail_parse_on_startup( map_data = copy.deepcopy(MAP_DATA) map_data.image = None with patch( - "homeassistant.components.roborock.image.RoborockMapDataParser.parse", + "homeassistant.components.roborock.coordinator.RoborockMapDataParser.parse", return_value=map_data, ): await async_setup_component(hass, DOMAIN, {}) @@ -226,7 +226,7 @@ async def test_fail_updating_image( # Update image, but get none for parse image. with ( patch( - "homeassistant.components.roborock.image.RoborockMapDataParser.parse", + "homeassistant.components.roborock.coordinator.RoborockMapDataParser.parse", return_value=map_data, ), patch( diff --git a/tests/components/roborock/test_vacuum.py b/tests/components/roborock/test_vacuum.py index 15fdeb4767c..2a2d9f210ed 100644 --- a/tests/components/roborock/test_vacuum.py +++ b/tests/components/roborock/test_vacuum.py @@ -261,7 +261,7 @@ async def test_get_current_position( return_value=b"", ), patch( - "homeassistant.components.roborock.image.RoborockMapDataParser.parse", + "homeassistant.components.roborock.coordinator.RoborockMapDataParser.parse", return_value=map_data, ), ): @@ -316,7 +316,7 @@ async def test_get_current_position_no_robot_position( return_value=b"", ), patch( - "homeassistant.components.roborock.image.RoborockMapDataParser.parse", + "homeassistant.components.roborock.coordinator.RoborockMapDataParser.parse", return_value=map_data, ), pytest.raises(HomeAssistantError, match="Robot position not found"),