feat: add the ability to set q7 mode

roborock/data/b01_q7/b01_q7_containers.py

@@ -3,7 +3,9 @@
 from ..containers import RoborockBase
 from .b01_q7_code_mappings import (
     B01Fault,
+    CleanTypeMapping,
     SCWindMapping,
+    WaterLevelMapping,
     WorkModeMapping,
     WorkStatusMapping,
 )
@@ -77,8 +79,8 @@ class B01Props(RoborockBase):
     status: WorkStatusMapping | None = None
     fault: B01Fault | None = None
     wind: SCWindMapping | None = None
-    water: int | None = None
-    mode: int | None = None
+    water: WaterLevelMapping | None = None
+    mode: CleanTypeMapping | None = None
     quantity: int | None = None
     alarm: int | None = None
     volume: int | None = None

roborock/devices/traits/b01/q7/__init__.py

@@ -6,6 +6,7 @@
 from roborock import B01Props
 from roborock.data.b01_q7.b01_q7_code_mappings import (
     CleanTaskTypeMapping,
+    CleanTypeMapping,
     SCDeviceCleanParam,
     SCWindMapping,
     WaterLevelMapping,
@@ -54,6 +55,10 @@ async def set_water_level(self, water_level: WaterLevelMapping) -> None:
         """Set the water level (water)."""
         await self.set_prop(RoborockB01Props.WATER, water_level.code)
 
+    async def set_mode(self, mode: CleanTypeMapping) -> None:
+        """Set the cleaning mode (vacuum, mop, or vacuum and mop)."""
+        await self.set_prop(RoborockB01Props.MODE, mode.code)
+
     async def start_clean(self) -> None:
         """Start cleaning."""
         await self.send(

tests/devices/traits/b01/q7/test_init.py

@@ -11,6 +11,7 @@
 
 from roborock.data.b01_q7 import (
     CleanTaskTypeMapping,
+    CleanTypeMapping,
     SCDeviceCleanParam,
     SCWindMapping,
     WaterLevelMapping,
@@ -217,6 +218,32 @@ async def test_q7_api_set_water_level(
     assert payload_data["dps"]["10000"]["params"] == {RoborockB01Props.WATER: WaterLevelMapping.HIGH.code}
 
 
+@pytest.mark.parametrize(
+    ("mode", "expected_code"),
+    [
+        (CleanTypeMapping.VACUUM, 0),
+        (CleanTypeMapping.VAC_AND_MOP, 1),
+        (CleanTypeMapping.MOP, 2),
+    ],
+)
+async def test_q7_api_set_mode(
+    mode: CleanTypeMapping,
+    expected_code: int,
+    q7_api: Q7PropertiesApi,
+    fake_channel: FakeChannel,
+    message_builder: B01MessageBuilder,
+):
+    """Test setting cleaning mode (vacuum, mop, or both)."""
+    fake_channel.response_queue.append(message_builder.build({"result": "ok"}))
+    await q7_api.set_mode(mode)
+
+    assert len(fake_channel.published_messages) == 1
+    message = fake_channel.published_messages[0]
+    payload_data = json.loads(unpad(message.payload, AES.block_size))
+    assert payload_data["dps"]["10000"]["method"] == "prop.set"
+    assert payload_data["dps"]["10000"]["params"] == {RoborockB01Props.MODE: expected_code}
+
+
 async def test_q7_api_start_clean(
     q7_api: Q7PropertiesApi, fake_channel: FakeChannel, message_builder: B01MessageBuilder
 ):