Move Roborock MapParser to coordinator (#140750)

Move MapParser to coordinator

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

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 (
+from .const import DOMAIN, IMAGE_CACHE_INTERVAL, MAP_SLEEP
-    DEFAULT_DRAWABLES,
-    DOMAIN,
-    DRAWABLES,
-    IMAGE_CACHE_INTERVAL,
-    MAP_FILE_FORMAT,
-    MAP_SCALE,
-    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(

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"

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(

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(

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"),