From 98fea46eea0d9ab5e11d0e12003cd8644b8beebf Mon Sep 17 00:00:00 2001 From: Luke Lashley Date: Tue, 30 Dec 2025 10:26:18 -0500 Subject: [PATCH] Add support for vacuum entity for Roborock Q7 (#159966) --- .../components/roborock/coordinator.py | 1 + homeassistant/components/roborock/vacuum.py | 175 +++++++++++++- tests/components/roborock/conftest.py | 29 ++- tests/components/roborock/test_vacuum.py | 219 ++++++++++++++++++ 4 files changed, 420 insertions(+), 4 deletions(-) diff --git a/homeassistant/components/roborock/coordinator.py b/homeassistant/components/roborock/coordinator.py index 025a5c37186a..8e0a26d49166 100644 --- a/homeassistant/components/roborock/coordinator.py +++ b/homeassistant/components/roborock/coordinator.py @@ -552,6 +552,7 @@ class RoborockB01Q7UpdateCoordinator(RoborockDataUpdateCoordinatorB01): RoborockB01Props.CLEANING_TIME, RoborockB01Props.REAL_CLEAN_TIME, RoborockB01Props.HYPA, + RoborockB01Props.WIND, ] async def _async_update_data( diff --git a/homeassistant/components/roborock/vacuum.py b/homeassistant/components/roborock/vacuum.py index 836149f0a056..cb0425b1518b 100644 --- a/homeassistant/components/roborock/vacuum.py +++ b/homeassistant/components/roborock/vacuum.py @@ -3,7 +3,7 @@ import logging from typing import Any -from roborock.data import RoborockStateCode +from roborock.data import RoborockStateCode, SCWindMapping, WorkStatusMapping from roborock.exceptions import RoborockException from roborock.roborock_typing import RoborockCommand import voluptuous as vol @@ -24,8 +24,12 @@ from .const import ( GET_VACUUM_CURRENT_POSITION_SERVICE_NAME, SET_VACUUM_GOTO_POSITION_SERVICE_NAME, ) -from .coordinator import RoborockConfigEntry, RoborockDataUpdateCoordinator -from .entity import RoborockCoordinatedEntityV1 +from .coordinator import ( + RoborockB01Q7UpdateCoordinator, + RoborockConfigEntry, + RoborockDataUpdateCoordinator, +) +from .entity import RoborockCoordinatedEntityB01, RoborockCoordinatedEntityV1 _LOGGER = logging.getLogger(__name__) @@ -57,6 +61,20 @@ STATE_CODE_TO_STATE = { RoborockStateCode.device_offline: VacuumActivity.ERROR, # "Device offline" } +Q7_STATE_CODE_TO_STATE = { + WorkStatusMapping.SLEEPING: VacuumActivity.IDLE, + WorkStatusMapping.WAITING_FOR_ORDERS: VacuumActivity.IDLE, + WorkStatusMapping.PAUSED: VacuumActivity.PAUSED, + WorkStatusMapping.DOCKING: VacuumActivity.RETURNING, + WorkStatusMapping.CHARGING: VacuumActivity.DOCKED, + WorkStatusMapping.SWEEP_MOPING: VacuumActivity.CLEANING, + WorkStatusMapping.SWEEP_MOPING_2: VacuumActivity.CLEANING, + WorkStatusMapping.MOPING: VacuumActivity.CLEANING, + WorkStatusMapping.UPDATING: VacuumActivity.DOCKED, + WorkStatusMapping.MOP_CLEANING: VacuumActivity.DOCKED, + WorkStatusMapping.MOP_AIRDRYING: VacuumActivity.DOCKED, +} + PARALLEL_UPDATES = 0 @@ -69,6 +87,11 @@ async def async_setup_entry( async_add_entities( RoborockVacuum(coordinator) for coordinator in config_entry.runtime_data.v1 ) + async_add_entities( + RoborockQ7Vacuum(coordinator) + for coordinator in config_entry.runtime_data.b01 + if isinstance(coordinator, RoborockB01Q7UpdateCoordinator) + ) platform = entity_platform.async_get_current_platform() platform.async_register_entity_service( @@ -241,3 +264,149 @@ class RoborockVacuum(RoborockCoordinatedEntityV1, StateVacuumEntity): "x": robot_position.x, "y": robot_position.y, } + + +class RoborockQ7Vacuum(RoborockCoordinatedEntityB01, StateVacuumEntity): + """General Representation of a Roborock vacuum.""" + + _attr_icon = "mdi:robot-vacuum" + _attr_supported_features = ( + VacuumEntityFeature.PAUSE + | VacuumEntityFeature.STOP + | VacuumEntityFeature.RETURN_HOME + | VacuumEntityFeature.FAN_SPEED + | VacuumEntityFeature.SEND_COMMAND + | VacuumEntityFeature.LOCATE + | VacuumEntityFeature.STATE + | VacuumEntityFeature.START + ) + _attr_translation_key = DOMAIN + _attr_name = None + coordinator: RoborockB01Q7UpdateCoordinator + + def __init__( + self, + coordinator: RoborockB01Q7UpdateCoordinator, + ) -> None: + """Initialize a vacuum.""" + StateVacuumEntity.__init__(self) + RoborockCoordinatedEntityB01.__init__( + self, + coordinator.duid_slug, + coordinator, + ) + + @property + def fan_speed_list(self) -> list[str]: + """Get the list of available fan speeds.""" + return SCWindMapping.keys() + + @property + def activity(self) -> VacuumActivity | None: + """Return the status of the vacuum cleaner.""" + if self.coordinator.data.status is not None: + return Q7_STATE_CODE_TO_STATE.get(self.coordinator.data.status) + return None + + @property + def fan_speed(self) -> str | None: + """Return the fan speed of the vacuum cleaner.""" + return self.coordinator.data.wind_name + + async def async_start(self) -> None: + """Start the vacuum.""" + try: + await self.coordinator.api.start_clean() + except RoborockException as err: + raise HomeAssistantError( + translation_domain=DOMAIN, + translation_key="command_failed", + translation_placeholders={ + "command": "start_clean", + }, + ) from err + + async def async_pause(self) -> None: + """Pause the vacuum.""" + try: + await self.coordinator.api.pause_clean() + except RoborockException as err: + raise HomeAssistantError( + translation_domain=DOMAIN, + translation_key="command_failed", + translation_placeholders={ + "command": "pause_clean", + }, + ) from err + + async def async_stop(self, **kwargs: Any) -> None: + """Stop the vacuum.""" + try: + await self.coordinator.api.stop_clean() + except RoborockException as err: + raise HomeAssistantError( + translation_domain=DOMAIN, + translation_key="command_failed", + translation_placeholders={ + "command": "stop_clean", + }, + ) from err + + async def async_return_to_base(self, **kwargs: Any) -> None: + """Send vacuum back to base.""" + try: + await self.coordinator.api.return_to_dock() + except RoborockException as err: + raise HomeAssistantError( + translation_domain=DOMAIN, + translation_key="command_failed", + translation_placeholders={ + "command": "return_to_dock", + }, + ) from err + + async def async_locate(self, **kwargs: Any) -> None: + """Locate vacuum.""" + try: + await self.coordinator.api.find_me() + except RoborockException as err: + raise HomeAssistantError( + translation_domain=DOMAIN, + translation_key="command_failed", + translation_placeholders={ + "command": "find_me", + }, + ) from err + + async def async_set_fan_speed(self, fan_speed: str, **kwargs: Any) -> None: + """Set vacuum fan speed.""" + try: + await self.coordinator.api.set_fan_speed( + SCWindMapping.from_value(fan_speed) + ) + except RoborockException as err: + raise HomeAssistantError( + translation_domain=DOMAIN, + translation_key="command_failed", + translation_placeholders={ + "command": "set_fan_speed", + }, + ) from err + + async def async_send_command( + self, + command: str, + params: dict[str, Any] | list[Any] | None = None, + **kwargs: Any, + ) -> None: + """Send a command to a vacuum cleaner.""" + try: + await self.coordinator.api.send(command, params) + except RoborockException as err: + raise HomeAssistantError( + translation_domain=DOMAIN, + translation_key="command_failed", + translation_placeholders={ + "command": command, + }, + ) from err diff --git a/tests/components/roborock/conftest.py b/tests/components/roborock/conftest.py index 0378ad98cba8..e48031400291 100644 --- a/tests/components/roborock/conftest.py +++ b/tests/components/roborock/conftest.py @@ -22,6 +22,7 @@ from roborock.data import ( RoborockBase, RoborockDyadStateCode, ValleyElectricityTimer, + WorkStatusMapping, ZeoError, ZeoState, ) @@ -110,7 +111,33 @@ def create_zeo_trait() -> Mock: def create_b01_q7_trait() -> Mock: """Create B01 Q7 trait for B01 devices.""" b01_trait = AsyncMock() - b01_trait.query_values.return_value = Q7_B01_PROPS + b01_trait._props_data = deepcopy(Q7_B01_PROPS) + + async def query_values_side_effect(protocols): + return b01_trait._props_data + + b01_trait.query_values = AsyncMock(side_effect=query_values_side_effect) + + # Add API methods that update the state when called + async def start_clean_side_effect(): + b01_trait._props_data.status = WorkStatusMapping.SWEEP_MOPING + + async def pause_clean_side_effect(): + b01_trait._props_data.status = WorkStatusMapping.PAUSED + + async def stop_clean_side_effect(): + b01_trait._props_data.status = WorkStatusMapping.WAITING_FOR_ORDERS + + async def return_to_dock_side_effect(): + b01_trait._props_data.status = WorkStatusMapping.DOCKING + + b01_trait.start_clean = AsyncMock(side_effect=start_clean_side_effect) + b01_trait.pause_clean = AsyncMock(side_effect=pause_clean_side_effect) + b01_trait.stop_clean = AsyncMock(side_effect=stop_clean_side_effect) + b01_trait.return_to_dock = AsyncMock(side_effect=return_to_dock_side_effect) + b01_trait.find_me = AsyncMock() + b01_trait.set_fan_speed = AsyncMock() + b01_trait.send = AsyncMock() return b01_trait diff --git a/tests/components/roborock/test_vacuum.py b/tests/components/roborock/test_vacuum.py index 56895fab7c08..0f331402aad8 100644 --- a/tests/components/roborock/test_vacuum.py +++ b/tests/components/roborock/test_vacuum.py @@ -37,6 +37,8 @@ from tests.common import MockConfigEntry ENTITY_ID = "vacuum.roborock_s7_maxv" DEVICE_ID = "abc123" +Q7_ENTITY_ID = "vacuum.roborock_q7" +Q7_DEVICE_ID = "q7_duid" @pytest.fixture @@ -270,3 +272,220 @@ async def test_get_current_position_no_robot_position( blocking=True, return_response=True, ) + + +# Tests for RoborockQ7Vacuum + + +@pytest.fixture +def fake_q7_vacuum(fake_devices: list[FakeDevice]) -> FakeDevice: + """Get the fake Q7 vacuum device.""" + # The Q7 is the fourth device in the list (index 3) based on HOME_DATA + return fake_devices[3] + + +@pytest.fixture(name="q7_vacuum_api", autouse=False) +def fake_q7_vacuum_api_fixture( + fake_q7_vacuum: FakeDevice, + send_message_exception: Exception | None, +) -> Mock: + """Get the fake Q7 vacuum device API for asserting that commands happened.""" + assert fake_q7_vacuum.b01_q7_properties is not None + api = fake_q7_vacuum.b01_q7_properties + if send_message_exception is not None: + # For exception tests, override side effects to raise the exception + api.start_clean.side_effect = send_message_exception + api.pause_clean.side_effect = send_message_exception + api.stop_clean.side_effect = send_message_exception + 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.send.side_effect = send_message_exception + return api + + +async def test_q7_registry_entries( + hass: HomeAssistant, + entity_registry: er.EntityRegistry, + device_registry: dr.DeviceRegistry, + setup_entry: MockConfigEntry, +) -> None: + """Tests Q7 devices are registered in the entity registry.""" + entity_entry = entity_registry.async_get(Q7_ENTITY_ID) + assert entity_entry.unique_id == Q7_DEVICE_ID + + device_entry = device_registry.async_get(entity_entry.device_id) + assert device_entry is not None + + +@pytest.mark.parametrize( + ("service", "api_method", "service_params", "expected_activity"), + [ + (SERVICE_START, "start_clean", None, "cleaning"), + (SERVICE_PAUSE, "pause_clean", None, "paused"), + (SERVICE_STOP, "stop_clean", None, "idle"), + (SERVICE_RETURN_TO_BASE, "return_to_dock", None, "returning"), + ], +) +async def test_q7_state_changing_commands( + hass: HomeAssistant, + setup_entry: MockConfigEntry, + service: str, + api_method: str, + service_params: dict[str, Any] | None, + expected_activity: str, + q7_vacuum_api: Mock, + fake_q7_vacuum: FakeDevice, +) -> None: + """Test sending state-changing commands to the Q7 vacuum.""" + vacuum = hass.states.get(Q7_ENTITY_ID) + assert vacuum + + data = {ATTR_ENTITY_ID: Q7_ENTITY_ID, **(service_params or {})} + await hass.services.async_call( + Platform.VACUUM, + service, + data, + blocking=True, + ) + api_call = getattr(q7_vacuum_api, api_method) + assert api_call.call_count == 1 + assert api_call.call_args[0] == () + + # Verify the entity state was updated + assert fake_q7_vacuum.b01_q7_properties is not None + # Force coordinator refresh to get updated state + coordinator = setup_entry.runtime_data.b01[0] + + await coordinator.async_refresh() + await hass.async_block_till_done() + vacuum = hass.states.get(Q7_ENTITY_ID) + assert vacuum + assert vacuum.state == expected_activity + + +async def test_q7_locate_command( + hass: HomeAssistant, + setup_entry: MockConfigEntry, + q7_vacuum_api: Mock, +) -> None: + """Test sending locate command to the Q7 vacuum.""" + vacuum = hass.states.get(Q7_ENTITY_ID) + assert vacuum + + await hass.services.async_call( + Platform.VACUUM, + SERVICE_LOCATE, + {ATTR_ENTITY_ID: Q7_ENTITY_ID}, + blocking=True, + ) + assert q7_vacuum_api.find_me.call_count == 1 + assert q7_vacuum_api.find_me.call_args[0] == () + + +async def test_q7_set_fan_speed_command( + hass: HomeAssistant, + setup_entry: MockConfigEntry, + q7_vacuum_api: Mock, +) -> None: + """Test sending set_fan_speed command to the Q7 vacuum.""" + vacuum = hass.states.get(Q7_ENTITY_ID) + assert vacuum + + await hass.services.async_call( + Platform.VACUUM, + SERVICE_SET_FAN_SPEED, + {ATTR_ENTITY_ID: Q7_ENTITY_ID, "fan_speed": "quiet"}, + blocking=True, + ) + assert q7_vacuum_api.set_fan_speed.call_count == 1 + # set_fan_speed is called with the fan speed value as first argument + assert len(q7_vacuum_api.set_fan_speed.call_args[0]) == 1 + + +async def test_q7_send_command( + hass: HomeAssistant, + setup_entry: MockConfigEntry, + q7_vacuum_api: Mock, +) -> None: + """Test sending custom command to the Q7 vacuum.""" + vacuum = hass.states.get(Q7_ENTITY_ID) + assert vacuum + + await hass.services.async_call( + Platform.VACUUM, + SERVICE_SEND_COMMAND, + {ATTR_ENTITY_ID: Q7_ENTITY_ID, "command": "test_command"}, + blocking=True, + ) + assert q7_vacuum_api.send.call_count == 1 + # send is called with command as first argument and params as second + assert q7_vacuum_api.send.call_args[0] == ("test_command", None) + + +@pytest.mark.parametrize( + ("service", "api_method", "service_params"), + [ + (SERVICE_START, "start_clean", None), + (SERVICE_PAUSE, "pause_clean", None), + (SERVICE_STOP, "stop_clean", None), + (SERVICE_RETURN_TO_BASE, "return_to_dock", None), + (SERVICE_LOCATE, "find_me", None), + (SERVICE_SET_FAN_SPEED, "set_fan_speed", {"fan_speed": "quiet"}), + (SERVICE_SEND_COMMAND, "send", {"command": "test_command"}), + ], +) +@pytest.mark.parametrize("send_message_exception", [RoborockException()]) +async def test_q7_failed_commands( + hass: HomeAssistant, + setup_entry: MockConfigEntry, + service: str, + api_method: str, + service_params: dict[str, Any] | None, + q7_vacuum_api: Mock, +) -> None: + """Test that when Q7 commands fail, we raise HomeAssistantError.""" + vacuum = hass.states.get(Q7_ENTITY_ID) + assert vacuum + # Store the original state to verify it doesn't change on error + original_state = vacuum.state + + data = {ATTR_ENTITY_ID: Q7_ENTITY_ID, **(service_params or {})} + command_name = ( + service_params.get("command", api_method) if service_params else api_method + ) + + with pytest.raises(HomeAssistantError, match=f"Error while calling {command_name}"): + await hass.services.async_call( + Platform.VACUUM, + service, + data, + blocking=True, + ) + + # Verify the entity state remains unchanged after failed command + await hass.async_block_till_done() + vacuum = hass.states.get(Q7_ENTITY_ID) + assert vacuum + assert vacuum.state == original_state + + +async def test_q7_activity_none_status( + hass: HomeAssistant, + setup_entry: MockConfigEntry, + fake_q7_vacuum: FakeDevice, +) -> None: + """Test that activity returns None when status is None.""" + assert fake_q7_vacuum.b01_q7_properties is not None + # Set status to None + fake_q7_vacuum.b01_q7_properties._props_data.status = None + + # Force coordinator refresh to get updated state + coordinator = setup_entry.runtime_data.b01[0] + await coordinator.async_refresh() + await hass.async_block_till_done() + + # Verify the entity state is unknown when status is None + vacuum = hass.states.get(Q7_ENTITY_ID) + assert vacuum + assert vacuum.state == "unknown"