Move Roborock MapParser to coordinator (#140750)

Move MapParser to coordinator
This commit is contained in:
Luke Lashley 2025-03-16 15:55:00 -04:00 committed by GitHub
parent 784381a25f
commit b0db7b432e
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
6 changed files with 48 additions and 53 deletions

View File

@ -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

View File

@ -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(

View File

@ -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"

View File

@ -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(

View File

@ -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(

View File

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