From 659c1f544427acef4f9c2813d4002b8bf593a4c6 Mon Sep 17 00:00:00 2001 From: Roman Dmytrenko Date: Wed, 30 Sep 2026 12:27:49 +0100 Subject: [PATCH] Add map, segment, and clean_segments support for Roborock Q7 (#176496) Signed-off-by: Roman Dmytrenko --- homeassistant/components/roborock/__init__.py | 1 + .../components/roborock/coordinator.py | 15 + homeassistant/components/roborock/image.py | 53 ++- homeassistant/components/roborock/vacuum.py | 114 ++++- tests/components/roborock/conftest.py | 41 ++ .../roborock/snapshots/test_vacuum.ambr | 4 +- tests/components/roborock/test_image.py | 67 ++- tests/components/roborock/test_vacuum.py | 403 +++++++++++++++++- 8 files changed, 690 insertions(+), 8 deletions(-) diff --git a/homeassistant/components/roborock/__init__.py b/homeassistant/components/roborock/__init__.py index 92ae0d7500c4..ddeaf3849c98 100644 --- a/homeassistant/components/roborock/__init__.py +++ b/homeassistant/components/roborock/__init__.py @@ -238,6 +238,7 @@ async def async_setup_device( RoborockDataUpdateCoordinator | RoborockDataUpdateCoordinatorA01 | RoborockDataUpdateCoordinatorB01 + | RoborockB01Q7UpdateCoordinator | RoborockB01Q10UpdateCoordinator | None ) = None diff --git a/homeassistant/components/roborock/coordinator.py b/homeassistant/components/roborock/coordinator.py index c2bebddd979d..cf621229294d 100644 --- a/homeassistant/components/roborock/coordinator.py +++ b/homeassistant/components/roborock/coordinator.py @@ -711,6 +711,21 @@ class RoborockB01Q7UpdateCoordinator(RoborockDataUpdateCoordinatorB01): ) return data + async def async_refresh_q7_map(self) -> bool: + """Refresh the Q7 map list and map content traits. + + Returns True when the traits were refreshed without errors. + Never raises; callers fall back to the cached trait values on failure. + """ + try: + await self.api.map.refresh() + if self.api.map.current_map_id is not None: + await self.api.map_content.refresh() + except RoborockException as ex: + _LOGGER.debug("Failed to refresh Q7 map: %s", ex) + return False + return True + class RoborockB01Q10UpdateCoordinator(DataUpdateCoordinator[None]): """Coordinator for B01 Q10 devices. diff --git a/homeassistant/components/roborock/image.py b/homeassistant/components/roborock/image.py index 52b41236c1e7..be3ea5629b3b 100644 --- a/homeassistant/components/roborock/image.py +++ b/homeassistant/components/roborock/image.py @@ -18,12 +18,17 @@ from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback from homeassistant.util import dt as dt_util from .coordinator import ( + RoborockB01Q7UpdateCoordinator, RoborockB01Q10UpdateCoordinator, RoborockConfigEntry, RoborockCoordinatorType, RoborockDataUpdateCoordinator, ) -from .entity import RoborockCoordinatedEntityB01Q10, RoborockCoordinatedEntityV1 +from .entity import ( + RoborockCoordinatedEntityB01Q7, + RoborockCoordinatedEntityB01Q10, + RoborockCoordinatedEntityV1, +) _LOGGER = logging.getLogger(__name__) @@ -85,6 +90,8 @@ async def async_setup_entry( async_update_map_entities ) ) + elif isinstance(coordinator, RoborockB01Q7UpdateCoordinator): + async_add_entities([RoborockMapQ7(coordinator)]) elif isinstance(coordinator, RoborockB01Q10UpdateCoordinator): async_add_entities([RoborockMapQ10(coordinator)]) @@ -213,3 +220,47 @@ class RoborockMapQ10(RoborockCoordinatedEntityB01Q10, ImageEntity): async def async_image(self) -> bytes | None: """Get the cached image.""" return self._cached_map + + +class RoborockMapQ7(RoborockCoordinatedEntityB01Q7, ImageEntity): + """A class to let you visualize the current map of a Q7 device.""" + + _attr_content_type = "image/png" + _attr_entity_category = EntityCategory.DIAGNOSTIC + _attr_translation_key = "map" + + def __init__(self, coordinator: RoborockB01Q7UpdateCoordinator) -> None: + """Initialize the Roborock Q7 map.""" + RoborockCoordinatedEntityB01Q7.__init__( + self, f"map_{coordinator.duid_slug}", coordinator + ) + ImageEntity.__init__(self, coordinator.hass) + self._map_content_trait = coordinator.api.map_content + self._cached_map: bytes | None = None + self._attr_image_last_updated = None + + @override + async def async_added_to_hass(self) -> None: + """When entity is added to hass, fetch and cache the current map image.""" + await super().async_added_to_hass() + await self.coordinator.async_refresh_q7_map() + image_content = self._map_content_trait.image_content + if image_content is not None: + self._cached_map = image_content + self._attr_image_last_updated = dt_util.utcnow() + self.async_write_ha_state() + + @callback + @override + def _handle_coordinator_update(self) -> None: + """Handle updated data from the coordinator.""" + image_content = self._map_content_trait.image_content + if image_content is not None and image_content != self._cached_map: + self._cached_map = image_content + self._attr_image_last_updated = dt_util.utcnow() + super()._handle_coordinator_update() + + @override + async def async_image(self) -> bytes | None: + """Get the cached image.""" + return self._cached_map diff --git a/homeassistant/components/roborock/vacuum.py b/homeassistant/components/roborock/vacuum.py index bd4607d9449d..baf1c0214b31 100644 --- a/homeassistant/components/roborock/vacuum.py +++ b/homeassistant/components/roborock/vacuum.py @@ -400,6 +400,7 @@ class RoborockQ7Vacuum(RoborockCoordinatedEntityB01Q7, StateVacuumEntity): | VacuumEntityFeature.LOCATE | VacuumEntityFeature.STATE | VacuumEntityFeature.START + | VacuumEntityFeature.CLEAN_AREA ) _attr_translation_key = DOMAIN _attr_name = None @@ -550,9 +551,120 @@ class RoborockQ7Vacuum(RoborockCoordinatedEntityB01Q7, StateVacuumEntity): }, ) from err + @staticmethod + def _get_room_names(map_data: Any) -> dict[str, str]: + """Safely extract and normalize the room_names mapping.""" + if map_data is None: + return {} + + additional_parameters = getattr(map_data, "additional_parameters", None) + if not isinstance(additional_parameters, dict): + return {} + + raw_room_names = additional_parameters.get("room_names") + if not isinstance(raw_room_names, dict): + return {} + + room_names: dict[str, str] = {} + for raw_key, raw_value in raw_room_names.items(): + try: + room_id = int(raw_key) + except TypeError, ValueError, OverflowError: + _LOGGER.error("Skipping invalid room id %r in room_names data", raw_key) + continue + if not isinstance(raw_value, str): + _LOGGER.error( + "Skipping invalid room name for id %r: expected str, got %s", + raw_key, + type(raw_value).__name__, + ) + continue + room_names[str(room_id)] = raw_value + + return room_names + async def get_maps(self) -> ServiceResponse: """Get map information such as map id and room ids.""" - raise ServiceNotSupported(DOMAIN, "get_maps", self.entity_id) + map_trait = self.coordinator.api.map + map_content_trait = self.coordinator.api.map_content + if not await self.coordinator.async_refresh_q7_map(): + raise HomeAssistantError( + translation_domain=DOMAIN, + translation_key="map_failure", + ) + + room_names = self._get_room_names(map_content_trait.map_data) + + return { + "maps": [ + { + "flag": entry.id, + "name": f"Map {entry.id}", + "rooms": ( + dict(room_names) if entry.id == map_trait.current_map_id else {} + ), + } + for entry in (map_trait.map_list or []) + if entry.id is not None + ] + } + + @override + async def async_get_segments(self) -> list[Segment]: + """Get the segments that can be cleaned.""" + map_trait = self.coordinator.api.map + map_content_trait = self.coordinator.api.map_content + current_map_id = map_trait.current_map_id + map_data = map_content_trait.map_data + if await self.coordinator.async_refresh_q7_map(): + current_map_id = map_trait.current_map_id + map_data = map_content_trait.map_data + if current_map_id is None: + return [] + room_names = self._get_room_names(map_data) + if not room_names: + return [] + map_name = f"Map {current_map_id}" + return [ + Segment(id=f"{current_map_id}_{room_id}", name=name, group=map_name) + for room_id, name in room_names.items() + ] + + @override + async def async_clean_segments(self, segment_ids: list[str], **kwargs: Any) -> None: + """Clean the specified segments.""" + if not await self.coordinator.async_refresh_q7_map(): + raise HomeAssistantError( + translation_domain=DOMAIN, + translation_key="map_failure", + ) + parsed: list[tuple[int, int]] = [] + for seg_id in segment_ids: + try: + map_id_str, room_id_str = seg_id.split("_", maxsplit=1) + parsed.append((int(map_id_str), int(room_id_str))) + except ValueError as err: + raise HomeAssistantError( + translation_domain=DOMAIN, + translation_key="segment_id_parse_error", + translation_placeholders={"segment_id": seg_id}, + ) from err + current_map_id = self.coordinator.api.map.current_map_id + current_map_segments = [ + room_id for map_id, room_id in parsed if map_id == current_map_id + ] + if not current_map_segments: + return + try: + await self.coordinator.api.clean_segments(current_map_segments) + except RoborockException as err: + raise HomeAssistantError( + translation_domain=DOMAIN, + translation_key="command_failed", + translation_placeholders={ + "command": "clean_segments", + }, + ) from err async def get_vacuum_current_position(self) -> ServiceResponse: """Get the current position of the vacuum from the map.""" diff --git a/tests/components/roborock/conftest.py b/tests/components/roborock/conftest.py index c994c053cbe3..9e567b192af9 100644 --- a/tests/components/roborock/conftest.py +++ b/tests/components/roborock/conftest.py @@ -35,10 +35,12 @@ from roborock.data import ( ZeoError, ZeoState, ) +from roborock.data.b01_q7.b01_q7_containers import Q7MapListEntry from roborock.data.v1.v1_containers import StatusV2 from roborock.device_features import RoborockDockFeatures from roborock.devices.device import RoborockDevice from roborock.devices.device_manager import DeviceManager +from roborock.devices.traits.b01 import Q7PropertiesApi from roborock.devices.traits.b01.q10.status import StatusTrait as Q10StatusTrait from roborock.devices.traits.common import DpsDataConverter from roborock.devices.traits.v1 import PropertiesApi @@ -67,6 +69,7 @@ from roborock.roborock_message import ( RoborockDyadDataProtocol, RoborockZeoProtocol, ) +from vacuum_map_parser_base.map_data import MapData from homeassistant.components.roborock.const import ( CONF_BASE_URL, @@ -178,9 +181,47 @@ def create_b01_q7_trait() -> Mock: b01_trait.set_clean_path_preference = AsyncMock() b01_trait.set_water_level = AsyncMock() b01_trait.send = AsyncMock() + + b01_trait.map = AsyncMock() + b01_trait.map.refresh = AsyncMock() + b01_trait.map.map_list = [] + b01_trait.map.current_map_id = None + + b01_trait.map_content = AsyncMock() + b01_trait.map_content.refresh = AsyncMock() + b01_trait.map_content.image_content = None + b01_trait.map_content.map_data = None + + b01_trait.clean_segments = AsyncMock() + return b01_trait +def seed_q7_map( + api: Q7PropertiesApi, + entries: list[tuple[int | None, bool | None]], + room_names: dict[int, str] | None = None, + image_content: bytes | None = None, +) -> None: + """Seed Q7 map traits, deriving the current map like the library does. + + Mirrors Q7MapList.current_map_id: prefers the entry marked current, + otherwise falls back to the first entry. + """ + api.map.map_list = [Q7MapListEntry(id=map_id, cur=cur) for map_id, cur in entries] + current_map_id = next( + (map_id for map_id, cur in entries if cur), + entries[0][0] if entries else None, + ) + api.map.current_map_id = current_map_id if isinstance(current_map_id, int) else None + if room_names is not None: + map_data = MapData() + map_data.additional_parameters = {"room_names": room_names} + api.map_content.map_data = map_data + if image_content is not None: + api.map_content.image_content = image_content + + def attach_update_listeners(trait: Mock) -> Callable[[], None]: """Give a mock Q10 trait working add_update_listener support. diff --git a/tests/components/roborock/snapshots/test_vacuum.ambr b/tests/components/roborock/snapshots/test_vacuum.ambr index 3e8901e8d5c4..4a9d85e7ae44 100644 --- a/tests/components/roborock/snapshots/test_vacuum.ambr +++ b/tests/components/roborock/snapshots/test_vacuum.ambr @@ -133,7 +133,7 @@ 'platform': 'roborock', 'previous_unique_id': None, 'suggested_object_id': None, - 'supported_features': , + 'supported_features': , 'translation_key': 'roborock', 'unique_id': 'q7_duid', 'unit_of_measurement': None, @@ -151,7 +151,7 @@ 'max_plus', ]), : 'Roborock Q7', - : , + : , }), 'context': , 'entity_id': 'vacuum.roborock_q7', diff --git a/tests/components/roborock/test_image.py b/tests/components/roborock/test_image.py index eb04ab3952f2..22502b53a72a 100644 --- a/tests/components/roborock/test_image.py +++ b/tests/components/roborock/test_image.py @@ -46,7 +46,7 @@ async def test_floorplan_image( fake_devices: list[FakeDevice], ) -> None: """Test floor plan map image is correctly set up.""" - assert len(hass.states.async_all("image")) == 5 + assert len(hass.states.async_all("image")) == 6 assert hass.states.get("image.roborock_s7_maxv_upstairs") is not None # Load the image on demand @@ -132,7 +132,7 @@ async def test_map_status_change( fake_vacuum: FakeDevice, ) -> None: """Test floor plan map image is correctly updated on status change.""" - assert len(hass.states.async_all("image")) == 5 + assert len(hass.states.async_all("image")) == 6 assert hass.states.get("image.roborock_s7_maxv_upstairs") is not None client = await hass_client() @@ -182,6 +182,7 @@ async def test_map_status_change( "image.roborock_s7_2_upstairs", "image.roborock_s7_maxv_downstairs", "image.roborock_s7_maxv_upstairs", + "image.roborock_q7_map", "image.roborock_q10_s5_map", }, ), @@ -193,6 +194,7 @@ async def test_map_status_change( # Expect default names based on map flags "image.roborock_s7_maxv_map_0", "image.roborock_s7_maxv_map_1", + "image.roborock_q7_map", "image.roborock_q10_s5_map", }, ), @@ -358,3 +360,64 @@ async def test_map_load_delayed( # Assert first map entity is removed, second still exists assert hass.states.get("image.roborock_s7_maxv_main_floor") is None assert hass.states.get("image.roborock_s7_maxv_upstairs") is not None + + +async def test_q7_map_image( + hass: HomeAssistant, + setup_entry: MockConfigEntry, + hass_client: ClientSessionGenerator, + fake_q7_vacuum: FakeDevice, + freezer: FrozenDateTimeFactory, +) -> None: + """Test the coordinator-driven Q7 map image.""" + entity_id = "image.roborock_q7_map" + assert hass.states.get(entity_id) is not None + + assert fake_q7_vacuum.b01_q7_properties is not None + api = fake_q7_vacuum.b01_q7_properties + map_content_trait = api.map_content + + api.map.refresh.assert_awaited() + + map_content_trait.image_content = b"\x89PNG-q7" + freezer.tick(timedelta(seconds=65)) + async_fire_time_changed(hass) + await hass.async_block_till_done() + + # The initial image is served via the proxy + client = await hass_client() + resp = await client.get(f"/api/image_proxy/{entity_id}") + assert resp.status == HTTPStatus.OK + assert await resp.read() == b"\x89PNG-q7" + + # A coordinator tick that does not change image_content must not update the entity + state = hass.states.get(entity_id) + assert state is not None + last_updated = state.state + freezer.tick(timedelta(seconds=65)) + async_fire_time_changed(hass) + await hass.async_block_till_done() + state = hass.states.get(entity_id) + assert state is not None + assert state.state == last_updated + + map_content_trait.image_content = b"\x89PNG-q7-new" + freezer.tick(timedelta(seconds=65)) + async_fire_time_changed(hass) + await hass.async_block_till_done() + + state = hass.states.get(entity_id) + assert state is not None + assert state.state != last_updated + resp = await client.get(f"/api/image_proxy/{entity_id}") + assert resp.status == HTTPStatus.OK + assert await resp.read() == b"\x89PNG-q7-new" + + # Setting image_content to None keeps the last cached image + map_content_trait.image_content = None + freezer.tick(timedelta(seconds=65)) + async_fire_time_changed(hass) + await hass.async_block_till_done() + resp = await client.get(f"/api/image_proxy/{entity_id}") + assert resp.status == HTTPStatus.OK + assert await resp.read() == b"\x89PNG-q7-new" diff --git a/tests/components/roborock/test_vacuum.py b/tests/components/roborock/test_vacuum.py index 796d13d7081f..42095cfdaf62 100644 --- a/tests/components/roborock/test_vacuum.py +++ b/tests/components/roborock/test_vacuum.py @@ -51,7 +51,7 @@ from homeassistant.helpers import ( from homeassistant.setup import async_setup_component from homeassistant.util import dt as dt_util -from .conftest import FakeDevice, set_trait_attributes +from .conftest import FakeDevice, seed_q7_map, set_trait_attributes from .mock_data import STATUS from tests.common import MockConfigEntry, async_fire_time_changed, snapshot_platform @@ -252,7 +252,6 @@ async def test_get_maps( @pytest.mark.parametrize( "entity_id", [ - Q7_ENTITY_ID, Q10_ENTITY_ID, ], ) @@ -700,6 +699,7 @@ def fake_q7_vacuum_api_fixture( api.return_to_dock.side_effect = send_message_exception api.find_me.side_effect = send_message_exception api.set_fan_speed.side_effect = send_message_exception + api.clean_segments.side_effect = send_message_exception api.send.side_effect = send_message_exception return api @@ -911,6 +911,405 @@ async def test_q7_working_sleep_is_paused( assert vacuum.state == "paused" +async def test_q7_get_maps( + hass: HomeAssistant, + setup_entry: MockConfigEntry, + fake_q7_vacuum: FakeDevice, +) -> None: + """Test that get_maps returns map data for Q7.""" + assert fake_q7_vacuum.b01_q7_properties is not None + api = fake_q7_vacuum.b01_q7_properties + + seed_q7_map(api, [(1, True), (2, False)], {10: "Living room", 11: "Bedroom"}) + api.map.refresh.reset_mock() + api.map_content.refresh.reset_mock() + + response = await hass.services.async_call( + DOMAIN, + GET_MAPS_SERVICE_NAME, + {ATTR_ENTITY_ID: Q7_ENTITY_ID}, + blocking=True, + return_response=True, + ) + assert response == { + Q7_ENTITY_ID: { + "maps": [ + { + "flag": 1, + "name": "Map 1", + "rooms": {"10": "Living room", "11": "Bedroom"}, + }, + {"flag": 2, "name": "Map 2", "rooms": {}}, + ] + } + } + api.map.refresh.assert_awaited() + api.map_content.refresh.assert_awaited() + + +async def test_q7_get_maps_empty( + hass: HomeAssistant, + setup_entry: MockConfigEntry, + fake_q7_vacuum: FakeDevice, +) -> None: + """Test that get_maps returns empty list when no maps.""" + assert fake_q7_vacuum.b01_q7_properties is not None + api = fake_q7_vacuum.b01_q7_properties + seed_q7_map(api, []) + api.map.refresh.reset_mock() + api.map_content.refresh.reset_mock() + + response = await hass.services.async_call( + DOMAIN, + GET_MAPS_SERVICE_NAME, + {ATTR_ENTITY_ID: Q7_ENTITY_ID}, + blocking=True, + return_response=True, + ) + assert response == {Q7_ENTITY_ID: {"maps": []}} + api.map.refresh.assert_awaited() + api.map_content.refresh.assert_not_awaited() + + +async def test_q7_get_maps_exception( + hass: HomeAssistant, + setup_entry: MockConfigEntry, + fake_q7_vacuum: FakeDevice, +) -> None: + """Test that get_maps raises HomeAssistantError on failure.""" + assert fake_q7_vacuum.b01_q7_properties is not None + api = fake_q7_vacuum.b01_q7_properties + api.map.refresh.side_effect = RoborockException() + + with pytest.raises( + HomeAssistantError, match="Something went wrong creating the map" + ): + await hass.services.async_call( + DOMAIN, + GET_MAPS_SERVICE_NAME, + {ATTR_ENTITY_ID: Q7_ENTITY_ID}, + blocking=True, + return_response=True, + ) + + +async def test_q7_get_segments( + hass: HomeAssistant, + setup_entry: MockConfigEntry, + fake_q7_vacuum: FakeDevice, + hass_ws_client: WebSocketGenerator, +) -> None: + """Test that async_get_segments returns rooms from the Q7 map.""" + assert fake_q7_vacuum.b01_q7_properties is not None + api = fake_q7_vacuum.b01_q7_properties + + seed_q7_map(api, [(1, True), (2, False)], {10: "Living room", 11: "Bedroom"}) + api.map.refresh.reset_mock() + api.map_content.refresh.reset_mock() + + client = await hass_ws_client(hass) + await client.send_json_auto_id( + {"type": "vacuum/get_segments", "entity_id": Q7_ENTITY_ID} + ) + msg = await client.receive_json() + assert msg["success"] + assert msg["result"] == { + "segments": [ + {"id": "1_10", "name": "Living room", "group": "Map 1"}, + {"id": "1_11", "name": "Bedroom", "group": "Map 1"}, + ] + } + api.map.refresh.assert_awaited() + api.map_content.refresh.assert_awaited() + + +async def test_q7_get_segments_refresh_failure( + hass: HomeAssistant, + setup_entry: MockConfigEntry, + fake_q7_vacuum: FakeDevice, + hass_ws_client: WebSocketGenerator, +) -> None: + """Test that async_get_segments falls back to cached rooms on refresh failure.""" + assert fake_q7_vacuum.b01_q7_properties is not None + api = fake_q7_vacuum.b01_q7_properties + + seed_q7_map(api, [(1, True)], {10: "Living room"}) + api.map.refresh.side_effect = RoborockException() + + client = await hass_ws_client(hass) + await client.send_json_auto_id( + {"type": "vacuum/get_segments", "entity_id": Q7_ENTITY_ID} + ) + msg = await client.receive_json() + assert msg["success"] + assert msg["result"] == { + "segments": [{"id": "1_10", "name": "Living room", "group": "Map 1"}] + } + + +async def test_q7_get_segments_partial_refresh_failure( + hass: HomeAssistant, + setup_entry: MockConfigEntry, + fake_q7_vacuum: FakeDevice, + hass_ws_client: WebSocketGenerator, +) -> None: + """Test that a partial refresh keeps the previous map id and rooms paired.""" + assert fake_q7_vacuum.b01_q7_properties is not None + api = fake_q7_vacuum.b01_q7_properties + + seed_q7_map(api, [(1, True)], {10: "Living room"}) + + async def switch_map() -> None: + api.map.current_map_id = 2 + + api.map.refresh.side_effect = switch_map + api.map_content.refresh.side_effect = RoborockException() + + client = await hass_ws_client(hass) + await client.send_json_auto_id( + {"type": "vacuum/get_segments", "entity_id": Q7_ENTITY_ID} + ) + msg = await client.receive_json() + assert msg["success"] + assert msg["result"] == { + "segments": [{"id": "1_10", "name": "Living room", "group": "Map 1"}] + } + + +async def test_q7_get_segments_no_current_map( + hass: HomeAssistant, + setup_entry: MockConfigEntry, + fake_q7_vacuum: FakeDevice, + hass_ws_client: WebSocketGenerator, +) -> None: + """Test that async_get_segments returns empty list without a current map.""" + assert fake_q7_vacuum.b01_q7_properties is not None + api = fake_q7_vacuum.b01_q7_properties + + seed_q7_map(api, []) + api.map.refresh.reset_mock() + api.map_content.refresh.reset_mock() + + client = await hass_ws_client(hass) + await client.send_json_auto_id( + {"type": "vacuum/get_segments", "entity_id": Q7_ENTITY_ID} + ) + msg = await client.receive_json() + assert msg["success"] + assert msg["result"] == {"segments": []} + api.map.refresh.assert_awaited() + api.map_content.refresh.assert_not_awaited() + + +async def test_q7_get_segments_first_map_fallback( + hass: HomeAssistant, + setup_entry: MockConfigEntry, + fake_q7_vacuum: FakeDevice, + hass_ws_client: WebSocketGenerator, +) -> None: + """Test that rooms attach to the first map when none is marked current.""" + assert fake_q7_vacuum.b01_q7_properties is not None + api = fake_q7_vacuum.b01_q7_properties + + seed_q7_map(api, [(2, False), (1, False)], {10: "Living room"}) + + client = await hass_ws_client(hass) + await client.send_json_auto_id( + {"type": "vacuum/get_segments", "entity_id": Q7_ENTITY_ID} + ) + msg = await client.receive_json() + assert msg["success"] + assert msg["result"] == { + "segments": [{"id": "2_10", "name": "Living room", "group": "Map 2"}] + } + + +async def test_q7_get_segments_no_data( + hass: HomeAssistant, + setup_entry: MockConfigEntry, + fake_q7_vacuum: FakeDevice, + hass_ws_client: WebSocketGenerator, +) -> None: + """Test that async_get_segments returns empty list when no map data.""" + assert fake_q7_vacuum.b01_q7_properties is not None + api = fake_q7_vacuum.b01_q7_properties + seed_q7_map(api, [(1, True)]) + + client = await hass_ws_client(hass) + await client.send_json_auto_id( + {"type": "vacuum/get_segments", "entity_id": Q7_ENTITY_ID} + ) + msg = await client.receive_json() + assert msg["success"] + assert msg["result"] == {"segments": []} + + +async def test_q7_clean_segments( + hass: HomeAssistant, + setup_entry: MockConfigEntry, + entity_registry: er.EntityRegistry, + fake_q7_vacuum: FakeDevice, +) -> None: + """Test that clean_area service calls clean_segments with correct room ids.""" + assert fake_q7_vacuum.b01_q7_properties is not None + api = fake_q7_vacuum.b01_q7_properties + seed_q7_map(api, [(1, True), (2, False)]) + api.map.refresh.reset_mock() + + entity_registry.async_update_entity_options( + Q7_ENTITY_ID, + VACUUM_DOMAIN, + { + "area_mapping": {"living_room": ["1_10"]}, + "last_seen_segments": [ + {"id": "1_10", "name": "Living room", "group": "Map 1"}, + ], + }, + ) + + await hass.services.async_call( + VACUUM_DOMAIN, + SERVICE_CLEAN_AREA, + {ATTR_ENTITY_ID: Q7_ENTITY_ID, "cleaning_area_id": ["living_room"]}, + blocking=True, + ) + + assert fake_q7_vacuum.b01_q7_properties.clean_segments.call_count == 1 + assert fake_q7_vacuum.b01_q7_properties.clean_segments.call_args == call([10]) + api.map.refresh.assert_awaited() + + +@pytest.mark.parametrize("send_message_exception", [RoborockException()]) +async def test_q7_clean_segments_failed( + hass: HomeAssistant, + setup_entry: MockConfigEntry, + entity_registry: er.EntityRegistry, + fake_q7_vacuum: FakeDevice, + q7_vacuum_api: Mock, +) -> None: + """Test that a clean_segments failure raises HomeAssistantError.""" + assert fake_q7_vacuum.b01_q7_properties is not None + seed_q7_map(fake_q7_vacuum.b01_q7_properties, [(1, True)]) + fake_q7_vacuum.b01_q7_properties.map.refresh.reset_mock() + entity_registry.async_update_entity_options( + Q7_ENTITY_ID, + VACUUM_DOMAIN, + { + "area_mapping": {"living_room": ["1_10"]}, + "last_seen_segments": [ + {"id": "1_10", "name": "Living room", "group": "Map 1"}, + ], + }, + ) + + with pytest.raises(HomeAssistantError, match="Error while calling clean_segments"): + await hass.services.async_call( + VACUUM_DOMAIN, + SERVICE_CLEAN_AREA, + {ATTR_ENTITY_ID: Q7_ENTITY_ID, "cleaning_area_id": ["living_room"]}, + blocking=True, + ) + fake_q7_vacuum.b01_q7_properties.map.refresh.assert_awaited() + + +async def test_q7_clean_segments_refresh_failure( + hass: HomeAssistant, + setup_entry: MockConfigEntry, + entity_registry: er.EntityRegistry, + fake_q7_vacuum: FakeDevice, +) -> None: + """Test that clean_area fails when the map refresh fails.""" + assert fake_q7_vacuum.b01_q7_properties is not None + api = fake_q7_vacuum.b01_q7_properties + seed_q7_map(api, [(1, True)], {10: "Living room"}) + api.map.refresh.side_effect = RoborockException() + + entity_registry.async_update_entity_options( + Q7_ENTITY_ID, + VACUUM_DOMAIN, + { + "area_mapping": {"living_room": ["1_10"]}, + "last_seen_segments": [ + {"id": "1_10", "name": "Living room", "group": "Map 1"}, + ], + }, + ) + + with pytest.raises( + HomeAssistantError, match="Something went wrong creating the map" + ): + await hass.services.async_call( + VACUUM_DOMAIN, + SERVICE_CLEAN_AREA, + {ATTR_ENTITY_ID: Q7_ENTITY_ID, "cleaning_area_id": ["living_room"]}, + blocking=True, + ) + + api.clean_segments.assert_not_called() + + +async def test_q7_clean_segments_ignores_other_map( + hass: HomeAssistant, + setup_entry: MockConfigEntry, + entity_registry: er.EntityRegistry, + fake_q7_vacuum: FakeDevice, +) -> None: + """Test that segments from another map are not sent to the device.""" + assert fake_q7_vacuum.b01_q7_properties is not None + api = fake_q7_vacuum.b01_q7_properties + seed_q7_map(api, [(1, True), (2, False)]) + + entity_registry.async_update_entity_options( + Q7_ENTITY_ID, + VACUUM_DOMAIN, + { + "area_mapping": {"living_room": ["2_10"]}, + "last_seen_segments": [ + {"id": "2_10", "name": "Living room", "group": "Map 2"}, + ], + }, + ) + + await hass.services.async_call( + VACUUM_DOMAIN, + SERVICE_CLEAN_AREA, + {ATTR_ENTITY_ID: Q7_ENTITY_ID, "cleaning_area_id": ["living_room"]}, + blocking=True, + ) + + api.clean_segments.assert_not_called() + + +async def test_q7_clean_segments_invalid_id( + hass: HomeAssistant, + setup_entry: MockConfigEntry, + entity_registry: er.EntityRegistry, + fake_q7_vacuum: FakeDevice, +) -> None: + """Test that a malformed segment id raises HomeAssistantError.""" + assert fake_q7_vacuum.b01_q7_properties is not None + seed_q7_map(fake_q7_vacuum.b01_q7_properties, [(1, True)]) + + entity_registry.async_update_entity_options( + Q7_ENTITY_ID, + VACUUM_DOMAIN, + { + "area_mapping": {"living_room": ["bogus"]}, + "last_seen_segments": [ + {"id": "bogus", "name": "Living room", "group": "Map 1"}, + ], + }, + ) + + with pytest.raises(HomeAssistantError, match="Invalid segment ID format"): + await hass.services.async_call( + VACUUM_DOMAIN, + SERVICE_CLEAN_AREA, + {ATTR_ENTITY_ID: Q7_ENTITY_ID, "cleaning_area_id": ["living_room"]}, + blocking=True, + ) + + @pytest.fixture(name="q10_vacuum_api", autouse=False) def fake_q10_vacuum_api_fixture( fake_q10_vacuum: FakeDevice,