Add support for more switchbot cloud vacuum models (#146637)

homeassistant/components/switchbot_cloud/__init__.py

@@ -142,12 +142,15 @@ async def make_device_data(
             hass, entry, api, device, coordinators_by_id
         )
         devices_data.sensors.append((device, coordinator))
-
     if isinstance(device, Device) and device.device_type in [
         "K10+",
         "K10+ Pro",
         "Robot Vacuum Cleaner S1",
         "Robot Vacuum Cleaner S1 Plus",
+        "K20+ Pro",
+        "Robot Vacuum Cleaner K10+ Pro Combo",
+        "Robot Vacuum Cleaner S10",
+        "S20",
     ]:
         coordinator = await coordinator_for_device(
             hass, entry, api, device, coordinators_by_id, True

homeassistant/components/switchbot_cloud/vacuum.py

@@ -2,7 +2,15 @@
 
 from typing import Any
 
-from switchbot_api import Device, Remote, SwitchBotAPI, VacuumCommands
+from switchbot_api import (
+    Device,
+    Remote,
+    SwitchBotAPI,
+    VacuumCleanerV2Commands,
+    VacuumCleanerV3Commands,
+    VacuumCleanMode,
+    VacuumCommands,
+)
 
 from homeassistant.components.vacuum import (
     StateVacuumEntity,
@@ -63,6 +71,11 @@ VACUUM_FAN_SPEED_TO_SWITCHBOT_FAN_SPEED: dict[str, str] = {
 class SwitchBotCloudVacuum(SwitchBotCloudEntity, StateVacuumEntity):
     """Representation of a SwitchBot vacuum."""
 
+    # "K10+"
+    # "K10+ Pro"
+    # "Robot Vacuum Cleaner S1"
+    # "Robot Vacuum Cleaner S1 Plus"
+
     _attr_supported_features: VacuumEntityFeature = (
         VacuumEntityFeature.BATTERY
         | VacuumEntityFeature.FAN_SPEED
@@ -85,23 +98,26 @@ class SwitchBotCloudVacuum(SwitchBotCloudEntity, StateVacuumEntity):
                 VacuumCommands.POW_LEVEL,
                 parameters=VACUUM_FAN_SPEED_TO_SWITCHBOT_FAN_SPEED[fan_speed],
             )
-        self.async_write_ha_state()
+            await self.coordinator.async_request_refresh()
 
     async def async_pause(self) -> None:
         """Pause the cleaning task."""
         await self.send_api_command(VacuumCommands.STOP)
+        self.async_write_ha_state()
 
     async def async_return_to_base(self, **kwargs: Any) -> None:
         """Set the vacuum cleaner to return to the dock."""
         await self.send_api_command(VacuumCommands.DOCK)
+        await self.coordinator.async_request_refresh()
 
     async def async_start(self) -> None:
         """Start or resume the cleaning task."""
         await self.send_api_command(VacuumCommands.START)
+        await self.coordinator.async_request_refresh()
 
     def _set_attributes(self) -> None:
         """Set attributes from coordinator data."""
-        if not self.coordinator.data:
+        if self.coordinator.data is None:
             return
 
         self._attr_battery_level = self.coordinator.data.get("battery")
@@ -109,11 +125,127 @@ class SwitchBotCloudVacuum(SwitchBotCloudEntity, StateVacuumEntity):
 
         switchbot_state = str(self.coordinator.data.get("workingStatus"))
         self._attr_activity = VACUUM_SWITCHBOT_STATE_TO_HA_STATE.get(switchbot_state)
+        if self._attr_fan_speed is None:
+            self._attr_fan_speed = VACUUM_FAN_SPEED_QUIET
+
+
+class SwitchBotCloudVacuumK20PlusPro(SwitchBotCloudVacuum):
+    """Representation of a SwitchBot K20+ Pro."""
+
+    async def async_set_fan_speed(self, fan_speed: str, **kwargs: Any) -> None:
+        """Set fan speed."""
+        self._attr_fan_speed = fan_speed
+        await self.send_api_command(
+            VacuumCleanerV2Commands.CHANGE_PARAM,
+            parameters={
+                "fanLevel": int(VACUUM_FAN_SPEED_TO_SWITCHBOT_FAN_SPEED[fan_speed]) + 1,
+                "waterLevel": 1,
+                "times": 1,
+            },
+        )
+        await self.coordinator.async_request_refresh()
+
+    async def async_pause(self) -> None:
+        """Pause the cleaning task."""
+        await self.send_api_command(VacuumCleanerV2Commands.PAUSE)
+        await self.coordinator.async_request_refresh()
+
+    async def async_return_to_base(self, **kwargs: Any) -> None:
+        """Set the vacuum cleaner to return to the dock."""
+        await self.send_api_command(VacuumCleanerV2Commands.DOCK)
+        await self.coordinator.async_request_refresh()
+
+    async def async_start(self) -> None:
+        """Start or resume the cleaning task."""
+        fan_level = (
+            VACUUM_FAN_SPEED_TO_SWITCHBOT_FAN_SPEED.get(self.fan_speed)
+            if self.fan_speed
+            else None
+        )
+        await self.send_api_command(
+            VacuumCleanerV2Commands.START_CLEAN,
+            parameters={
+                "action": VacuumCleanMode.SWEEP.value,
+                "param": {
+                    "fanLevel": int(fan_level if fan_level else VACUUM_FAN_SPEED_QUIET)
+                    + 1,
+                    "times": 1,
+                },
+            },
+        )
+        await self.coordinator.async_request_refresh()
+
+
+class SwitchBotCloudVacuumK10PlusProCombo(SwitchBotCloudVacuumK20PlusPro):
+    """Representation of a SwitchBot vacuum K10+ Pro Combo."""
+
+    async def async_set_fan_speed(self, fan_speed: str, **kwargs: Any) -> None:
+        """Set fan speed."""
+        self._attr_fan_speed = fan_speed
+        if fan_speed in VACUUM_FAN_SPEED_TO_SWITCHBOT_FAN_SPEED:
+            await self.send_api_command(
+                VacuumCleanerV2Commands.CHANGE_PARAM,
+                parameters={
+                    "fanLevel": int(VACUUM_FAN_SPEED_TO_SWITCHBOT_FAN_SPEED[fan_speed])
+                    + 1,
+                    "times": 1,
+                },
+            )
+        await self.coordinator.async_request_refresh()
+
+
+class SwitchBotCloudVacuumV3(SwitchBotCloudVacuumK20PlusPro):
+    """Representation of a SwitchBot vacuum Robot Vacuum Cleaner S10 & S20."""
+
+    async def async_set_fan_speed(self, fan_speed: str, **kwargs: Any) -> None:
+        """Set fan speed."""
+        self._attr_fan_speed = fan_speed
+        await self.send_api_command(
+            VacuumCleanerV3Commands.CHANGE_PARAM,
+            parameters={
+                "fanLevel": int(VACUUM_FAN_SPEED_TO_SWITCHBOT_FAN_SPEED[fan_speed]) + 1,
+                "waterLevel": 1,
+                "times": 1,
+            },
+        )
+        await self.coordinator.async_request_refresh()
+
+    async def async_start(self) -> None:
+        """Start or resume the cleaning task."""
+        fan_level = (
+            VACUUM_FAN_SPEED_TO_SWITCHBOT_FAN_SPEED.get(self.fan_speed)
+            if self.fan_speed
+            else None
+        )
+        await self.send_api_command(
+            VacuumCleanerV3Commands.START_CLEAN,
+            parameters={
+                "action": VacuumCleanMode.SWEEP.value,
+                "param": {
+                    "fanLevel": int(fan_level if fan_level else VACUUM_FAN_SPEED_QUIET),
+                    "waterLevel": 1,
+                    "times": 1,
+                },
+            },
+        )
+        await self.coordinator.async_request_refresh()
 
 
 @callback
 def _async_make_entity(
     api: SwitchBotAPI, device: Device | Remote, coordinator: SwitchBotCoordinator
-) -> SwitchBotCloudVacuum:
+) -> (
+    SwitchBotCloudVacuum
+    | SwitchBotCloudVacuumK20PlusPro
+    | SwitchBotCloudVacuumV3
+    | SwitchBotCloudVacuumK10PlusProCombo
+):
     """Make a SwitchBotCloudVacuum."""
+    if device.device_type in VacuumCleanerV2Commands.get_supported_devices():
+        if device.device_type == "K20+ Pro":
+            return SwitchBotCloudVacuumK20PlusPro(api, device, coordinator)
+        return SwitchBotCloudVacuumK10PlusProCombo(api, device, coordinator)
+
+    if device.device_type in VacuumCleanerV3Commands.get_supported_devices():
+        return SwitchBotCloudVacuumV3(api, device, coordinator)
     return SwitchBotCloudVacuum(api, device, coordinator)

tests/components/switchbot_cloud/test_vacuum.py

@@ -0,0 +1,522 @@
+"""Test for the switchbot_cloud vacuum."""
+
+from unittest.mock import patch
+
+from switchbot_api import (
+    Device,
+    VacuumCleanerV2Commands,
+    VacuumCleanerV3Commands,
+    VacuumCleanMode,
+    VacuumCommands,
+)
+
+from homeassistant.components.switchbot_cloud import SwitchBotAPI
+from homeassistant.components.switchbot_cloud.const import VACUUM_FAN_SPEED_QUIET
+from homeassistant.components.vacuum import (
+    ATTR_FAN_SPEED,
+    DOMAIN as VACUUM_DOMAIN,
+    SERVICE_PAUSE,
+    SERVICE_RETURN_TO_BASE,
+    SERVICE_SET_FAN_SPEED,
+    SERVICE_START,
+    VacuumActivity,
+)
+from homeassistant.const import ATTR_ENTITY_ID, STATE_UNKNOWN
+from homeassistant.core import HomeAssistant
+
+from . import configure_integration
+
+
+async def test_coordinator_data_is_none(
+    hass: HomeAssistant, mock_list_devices, mock_get_status
+) -> None:
+    """Test coordinator data is none."""
+    mock_list_devices.return_value = [
+        Device(
+            version="V1.0",
+            deviceId="vacuum-id-1",
+            deviceName="vacuum-1",
+            deviceType="K10+",
+            hubDeviceId="test-hub-id",
+        ),
+    ]
+    mock_get_status.side_effect = [
+        None,
+    ]
+    await configure_integration(hass)
+    entity_id = "vacuum.vacuum_1"
+    state = hass.states.get(entity_id)
+
+    assert state.state == STATE_UNKNOWN
+
+
+async def test_k10_plus_set_fan_speed(
+    hass: HomeAssistant, mock_list_devices, mock_get_status
+) -> None:
+    """Test K10 plus set fan speed."""
+
+    mock_list_devices.side_effect = [
+        [
+            Device(
+                version="V1.0",
+                deviceId="vacuum-id-1",
+                deviceName="vacuum-1",
+                deviceType="K10+",
+                hubDeviceId="test-hub-id",
+            )
+        ]
+    ]
+    mock_get_status.side_effect = [
+        {
+            "deviceType": "K10+",
+            "workingStatus": "Cleaning",
+            "battery": 50,
+            "onlineStatus": "online",
+        },
+    ]
+
+    await configure_integration(hass)
+    entity_id = "vacuum.vacuum_1"
+    with patch.object(SwitchBotAPI, "send_command") as mock_send_command:
+        await hass.services.async_call(
+            VACUUM_DOMAIN,
+            SERVICE_SET_FAN_SPEED,
+            {ATTR_ENTITY_ID: entity_id, ATTR_FAN_SPEED: VACUUM_FAN_SPEED_QUIET},
+            blocking=True,
+        )
+        mock_send_command.assert_called_once_with(
+            "vacuum-id-1", VacuumCommands.POW_LEVEL, "command", "0"
+        )
+
+
+async def test_k10_plus_return_to_base(
+    hass: HomeAssistant, mock_list_devices, mock_get_status
+) -> None:
+    """Test k10 plus return to base."""
+    mock_list_devices.return_value = [
+        Device(
+            version="V1.0",
+            deviceId="vacuum-id-1",
+            deviceName="vacuum-1",
+            deviceType="K10+",
+            hubDeviceId="test-hub-id",
+        ),
+    ]
+
+    mock_get_status.side_effect = [
+        {
+            "deviceType": "K10+",
+            "workingStatus": "Charging",
+            "battery": 50,
+            "onlineStatus": "online",
+        }
+    ]
+
+    await configure_integration(hass)
+    entity_id = "vacuum.vacuum_1"
+    state = hass.states.get(entity_id)
+
+    assert state.state == VacuumActivity.DOCKED.value
+
+    with patch.object(SwitchBotAPI, "send_command") as mock_send_command:
+        await hass.services.async_call(
+            VACUUM_DOMAIN,
+            SERVICE_RETURN_TO_BASE,
+            {ATTR_ENTITY_ID: entity_id},
+            blocking=True,
+        )
+        mock_send_command.assert_called_once_with(
+            "vacuum-id-1", VacuumCommands.DOCK, "command", "default"
+        )
+
+
+async def test_k10_plus_pause(
+    hass: HomeAssistant, mock_list_devices, mock_get_status
+) -> None:
+    """Test k10 plus pause."""
+    mock_list_devices.return_value = [
+        Device(
+            version="V1.0",
+            deviceId="vacuum-id-1",
+            deviceName="vacuum-1",
+            deviceType="K10+",
+            hubDeviceId="test-hub-id",
+        ),
+    ]
+
+    mock_get_status.side_effect = [
+        {
+            "deviceType": "K10+",
+            "workingStatus": "Charging",
+            "battery": 50,
+            "onlineStatus": "online",
+        }
+    ]
+
+    await configure_integration(hass)
+    entity_id = "vacuum.vacuum_1"
+    state = hass.states.get(entity_id)
+
+    assert state.state == VacuumActivity.DOCKED.value
+
+    with patch.object(SwitchBotAPI, "send_command") as mock_send_command:
+        await hass.services.async_call(
+            VACUUM_DOMAIN, SERVICE_PAUSE, {ATTR_ENTITY_ID: entity_id}, blocking=True
+        )
+        mock_send_command.assert_called_once_with(
+            "vacuum-id-1", VacuumCommands.STOP, "command", "default"
+        )
+
+
+async def test_k10_plus_set_start(
+    hass: HomeAssistant, mock_list_devices, mock_get_status
+) -> None:
+    """Test K10 plus start."""
+
+    mock_list_devices.side_effect = [
+        [
+            Device(
+                version="V1.0",
+                deviceId="vacuum-id-1",
+                deviceName="vacuum-1",
+                deviceType="K10+",
+                hubDeviceId="test-hub-id",
+            )
+        ]
+    ]
+    mock_get_status.side_effect = [
+        {
+            "deviceType": "K10+",
+            "workingStatus": "Cleaning",
+            "battery": 50,
+            "onlineStatus": "online",
+        },
+    ]
+
+    await configure_integration(hass)
+    entity_id = "vacuum.vacuum_1"
+    with patch.object(SwitchBotAPI, "send_command") as mock_send_command:
+        await hass.services.async_call(
+            VACUUM_DOMAIN,
+            SERVICE_START,
+            {ATTR_ENTITY_ID: entity_id},
+            blocking=True,
+        )
+        mock_send_command.assert_called_once_with(
+            "vacuum-id-1", VacuumCommands.START, "command", "default"
+        )
+
+
+async def test_k20_plus_pro_set_fan_speed(
+    hass: HomeAssistant, mock_list_devices, mock_get_status
+) -> None:
+    """Test K10 plus set fan speed."""
+
+    mock_list_devices.side_effect = [
+        [
+            Device(
+                version="V1.0",
+                deviceId="vacuum-id-1",
+                deviceName="vacuum-1",
+                deviceType="K20+ Pro",
+                hubDeviceId="test-hub-id",
+            )
+        ]
+    ]
+    mock_get_status.side_effect = [
+        {
+            "deviceType": "K20+ Pro",
+            "workingStatus": "Cleaning",
+            "battery": 50,
+            "onlineStatus": "online",
+        },
+    ]
+
+    await configure_integration(hass)
+    entity_id = "vacuum.vacuum_1"
+    with patch.object(SwitchBotAPI, "send_command") as mock_send_command:
+        await hass.services.async_call(
+            VACUUM_DOMAIN,
+            SERVICE_SET_FAN_SPEED,
+            {ATTR_ENTITY_ID: entity_id, ATTR_FAN_SPEED: VACUUM_FAN_SPEED_QUIET},
+            blocking=True,
+        )
+        mock_send_command.assert_called_once_with(
+            "vacuum-id-1",
+            VacuumCleanerV2Commands.CHANGE_PARAM,
+            "command",
+            {
+                "fanLevel": 1,
+                "waterLevel": 1,
+                "times": 1,
+            },
+        )
+
+
+async def test_k20_plus_pro_return_to_base(
+    hass: HomeAssistant, mock_list_devices, mock_get_status
+) -> None:
+    """Test K20+ Pro return to base."""
+    mock_list_devices.side_effect = [
+        [
+            Device(
+                version="V1.0",
+                deviceId="vacuum-id-1",
+                deviceName="vacuum-1",
+                deviceType="K20+ Pro",
+                hubDeviceId="test-hub-id",
+            )
+        ]
+    ]
+    mock_get_status.side_effect = [
+        {
+            "deviceType": "K20+ Pro",
+            "workingStatus": "Charging",
+            "battery": 50,
+            "onlineStatus": "online",
+        },
+    ]
+
+    await configure_integration(hass)
+    entity_id = "vacuum.vacuum_1"
+    state = hass.states.get(entity_id)
+
+    assert state.state == VacuumActivity.DOCKED.value
+
+    with patch.object(SwitchBotAPI, "send_command") as mock_send_command:
+        await hass.services.async_call(
+            VACUUM_DOMAIN,
+            SERVICE_RETURN_TO_BASE,
+            {ATTR_ENTITY_ID: entity_id},
+            blocking=True,
+        )
+        mock_send_command.assert_called_once_with(
+            "vacuum-id-1", VacuumCleanerV2Commands.DOCK, "command", "default"
+        )
+
+
+async def test_k20_plus_pro_pause(
+    hass: HomeAssistant, mock_list_devices, mock_get_status
+) -> None:
+    """Test K20+ Pro pause."""
+    mock_list_devices.side_effect = [
+        [
+            Device(
+                version="V1.0",
+                deviceId="vacuum-id-1",
+                deviceName="vacuum-1",
+                deviceType="K20+ Pro",
+                hubDeviceId="test-hub-id",
+            )
+        ]
+    ]
+    mock_get_status.side_effect = [
+        {
+            "deviceType": "K20+ Pro",
+            "workingStatus": "Charging",
+            "battery": 50,
+            "onlineStatus": "online",
+        },
+    ]
+
+    await configure_integration(hass)
+    entity_id = "vacuum.vacuum_1"
+    state = hass.states.get(entity_id)
+
+    assert state.state == VacuumActivity.DOCKED.value
+
+    with patch.object(SwitchBotAPI, "send_command") as mock_send_command:
+        await hass.services.async_call(
+            VACUUM_DOMAIN, SERVICE_PAUSE, {ATTR_ENTITY_ID: entity_id}, blocking=True
+        )
+        mock_send_command.assert_called_once_with(
+            "vacuum-id-1", VacuumCleanerV2Commands.PAUSE, "command", "default"
+        )
+
+
+async def test_k20_plus_pro_start(
+    hass: HomeAssistant, mock_list_devices, mock_get_status
+) -> None:
+    """Test K20+ Pro start."""
+
+    mock_list_devices.side_effect = [
+        [
+            Device(
+                version="V1.0",
+                deviceId="vacuum-id-1",
+                deviceName="vacuum-1",
+                deviceType="K20+ Pro",
+                hubDeviceId="test-hub-id",
+            )
+        ]
+    ]
+    mock_get_status.side_effect = [
+        {
+            "deviceType": "K20+ Pro",
+            "workingStatus": "Cleaning",
+            "battery": 50,
+            "onlineStatus": "online",
+        },
+    ]
+
+    await configure_integration(hass)
+    entity_id = "vacuum.vacuum_1"
+    with patch.object(SwitchBotAPI, "send_command") as mock_send_command:
+        await hass.services.async_call(
+            VACUUM_DOMAIN,
+            SERVICE_START,
+            {ATTR_ENTITY_ID: entity_id},
+            blocking=True,
+        )
+        mock_send_command.assert_called_once_with(
+            "vacuum-id-1",
+            VacuumCleanerV2Commands.START_CLEAN,
+            "command",
+            {
+                "action": VacuumCleanMode.SWEEP.value,
+                "param": {
+                    "fanLevel": 1,
+                    "times": 1,
+                },
+            },
+        )
+
+
+async def test_k10_plus_pro_combo_set_fan_speed(
+    hass: HomeAssistant, mock_list_devices, mock_get_status
+) -> None:
+    """Test k10+ Pro Combo set fan speed."""
+
+    mock_list_devices.side_effect = [
+        [
+            Device(
+                version="V1.0",
+                deviceId="vacuum-id-1",
+                deviceName="vacuum-1",
+                deviceType="Robot Vacuum Cleaner K10+ Pro Combo",
+                hubDeviceId="test-hub-id",
+            )
+        ]
+    ]
+    mock_get_status.side_effect = [
+        {
+            "deviceType": "Robot Vacuum Cleaner K10+ Pro Combo",
+            "workingStatus": "Cleaning",
+            "battery": 50,
+            "onlineStatus": "online",
+        },
+    ]
+
+    await configure_integration(hass)
+    entity_id = "vacuum.vacuum_1"
+    with patch.object(SwitchBotAPI, "send_command") as mock_send_command:
+        await hass.services.async_call(
+            VACUUM_DOMAIN,
+            SERVICE_SET_FAN_SPEED,
+            {ATTR_ENTITY_ID: entity_id, ATTR_FAN_SPEED: VACUUM_FAN_SPEED_QUIET},
+            blocking=True,
+        )
+        mock_send_command.assert_called_once_with(
+            "vacuum-id-1",
+            VacuumCleanerV2Commands.CHANGE_PARAM,
+            "command",
+            {
+                "fanLevel": 1,
+                "times": 1,
+            },
+        )
+
+
+async def test_s20_start(
+    hass: HomeAssistant, mock_list_devices, mock_get_status
+) -> None:
+    """Test s20 start."""
+
+    mock_list_devices.side_effect = [
+        [
+            Device(
+                version="V1.0",
+                deviceId="vacuum-id-1",
+                deviceName="vacuum-1",
+                deviceType="S20",
+                hubDeviceId="test-hub-id",
+            )
+        ]
+    ]
+    mock_get_status.side_effect = [
+        {
+            "deviceType": "s20",
+            "workingStatus": "Cleaning",
+            "battery": 50,
+            "onlineStatus": "online",
+        },
+    ]
+
+    await configure_integration(hass)
+    entity_id = "vacuum.vacuum_1"
+    with patch.object(SwitchBotAPI, "send_command") as mock_send_command:
+        await hass.services.async_call(
+            VACUUM_DOMAIN,
+            SERVICE_START,
+            {ATTR_ENTITY_ID: entity_id},
+            blocking=True,
+        )
+        mock_send_command.assert_called_once_with(
+            "vacuum-id-1",
+            VacuumCleanerV3Commands.START_CLEAN,
+            "command",
+            {
+                "action": VacuumCleanMode.SWEEP.value,
+                "param": {
+                    "fanLevel": 0,
+                    "waterLevel": 1,
+                    "times": 1,
+                },
+            },
+        )
+
+
+async def test_s20set_fan_speed(
+    hass: HomeAssistant, mock_list_devices, mock_get_status
+) -> None:
+    """Test s20 set fan speed."""
+
+    mock_list_devices.side_effect = [
+        [
+            Device(
+                version="V1.0",
+                deviceId="vacuum-id-1",
+                deviceName="vacuum-1",
+                deviceType="S20",
+                hubDeviceId="test-hub-id",
+            )
+        ]
+    ]
+    mock_get_status.side_effect = [
+        {
+            "deviceType": "S20",
+            "workingStatus": "Cleaning",
+            "battery": 50,
+            "onlineStatus": "online",
+        },
+    ]
+
+    await configure_integration(hass)
+    entity_id = "vacuum.vacuum_1"
+    with patch.object(SwitchBotAPI, "send_command") as mock_send_command:
+        await hass.services.async_call(
+            VACUUM_DOMAIN,
+            SERVICE_SET_FAN_SPEED,
+            {ATTR_ENTITY_ID: entity_id, ATTR_FAN_SPEED: VACUUM_FAN_SPEED_QUIET},
+            blocking=True,
+        )
+        mock_send_command.assert_called_once_with(
+            "vacuum-id-1",
+            VacuumCleanerV3Commands.CHANGE_PARAM,
+            "command",
+            {
+                "fanLevel": 1,
+                "waterLevel": 1,
+                "times": 1,
+            },
+        )