Migrate Tuya vacuum to use wrapper class (home-assistant#156569)

Author: epenet

homeassistant/components/tuya/vacuum.py

@@ -16,9 +16,9 @@
 from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
 
 from . import TuyaConfigEntry
-from .const import TUYA_DISCOVERY_NEW, DeviceCategory, DPCode, DPType
+from .const import TUYA_DISCOVERY_NEW, DeviceCategory, DPCode
 from .entity import TuyaEntity
-from .models import EnumTypeData, find_dpcode
+from .models import DPCodeBooleanWrapper, DPCodeEnumWrapper
 from .util import get_dpcode
 
 TUYA_MODE_RETURN_HOME = "chargego"
@@ -64,7 +64,27 @@ def async_discover_device(device_ids: list[str]) -> None:
         for device_id in device_ids:
             device = manager.device_map[device_id]
             if device.category == DeviceCategory.SD:
-                entities.append(TuyaVacuumEntity(device, manager))
+                entities.append(
+                    TuyaVacuumEntity(
+                        device,
+                        manager,
+                        charge_wrapper=DPCodeBooleanWrapper.find_dpcode(
+                            device, DPCode.SWITCH_CHARGE, prefer_function=True
+                        ),
+                        fan_speed_wrapper=DPCodeEnumWrapper.find_dpcode(
+                            device, DPCode.SUCTION, prefer_function=True
+                        ),
+                        locate_wrapper=DPCodeBooleanWrapper.find_dpcode(
+                            device, DPCode.SEEK, prefer_function=True
+                        ),
+                        mode_wrapper=DPCodeEnumWrapper.find_dpcode(
+                            device, DPCode.MODE, prefer_function=True
+                        ),
+                        switch_wrapper=DPCodeBooleanWrapper.find_dpcode(
+                            device, DPCode.POWER_GO, prefer_function=True
+                        ),
+                    )
+                )
         async_add_entities(entities)
 
     async_discover_device([*manager.device_map])
@@ -77,51 +97,56 @@ def async_discover_device(device_ids: list[str]) -> None:
 class TuyaVacuumEntity(TuyaEntity, StateVacuumEntity):
     """Tuya Vacuum Device."""
 
-    _fan_speed: EnumTypeData | None = None
     _attr_name = None
 
-    def __init__(self, device: CustomerDevice, device_manager: Manager) -> None:
+    def __init__(
+        self,
+        device: CustomerDevice,
+        device_manager: Manager,
+        *,
+        charge_wrapper: DPCodeBooleanWrapper | None,
+        fan_speed_wrapper: DPCodeEnumWrapper | None,
+        locate_wrapper: DPCodeBooleanWrapper | None,
+        mode_wrapper: DPCodeEnumWrapper | None,
+        switch_wrapper: DPCodeBooleanWrapper | None,
+    ) -> None:
         """Init Tuya vacuum."""
         super().__init__(device, device_manager)
+        self._charge_wrapper = charge_wrapper
+        self._fan_speed_wrapper = fan_speed_wrapper
+        self._locate_wrapper = locate_wrapper
+        self._mode_wrapper = mode_wrapper
+        self._switch_wrapper = switch_wrapper
 
         self._attr_fan_speed_list = []
-
         self._attr_supported_features = (
             VacuumEntityFeature.SEND_COMMAND | VacuumEntityFeature.STATE
         )
         if get_dpcode(self.device, DPCode.PAUSE):
             self._attr_supported_features |= VacuumEntityFeature.PAUSE
 
-        self._return_home_use_switch_charge = False
-        if get_dpcode(self.device, DPCode.SWITCH_CHARGE):
-            self._attr_supported_features |= VacuumEntityFeature.RETURN_HOME
-            self._return_home_use_switch_charge = True
-        elif (
-            enum_type := find_dpcode(
-                self.device, DPCode.MODE, dptype=DPType.ENUM, prefer_function=True
-            )
-        ) and TUYA_MODE_RETURN_HOME in enum_type.range:
+        if charge_wrapper or (
+            mode_wrapper
+            and TUYA_MODE_RETURN_HOME in mode_wrapper.type_information.range
+        ):
             self._attr_supported_features |= VacuumEntityFeature.RETURN_HOME
 
-        if get_dpcode(self.device, DPCode.SEEK):
+        if locate_wrapper:
             self._attr_supported_features |= VacuumEntityFeature.LOCATE
 
-        if get_dpcode(self.device, DPCode.POWER_GO):
+        if switch_wrapper:
             self._attr_supported_features |= (
                 VacuumEntityFeature.STOP | VacuumEntityFeature.START
             )
 
-        if enum_type := find_dpcode(
-            self.device, DPCode.SUCTION, dptype=DPType.ENUM, prefer_function=True
-        ):
-            self._fan_speed = enum_type
-            self._attr_fan_speed_list = enum_type.range
+        if fan_speed_wrapper:
+            self._attr_fan_speed_list = fan_speed_wrapper.type_information.range
             self._attr_supported_features |= VacuumEntityFeature.FAN_SPEED
 
     @property
     def fan_speed(self) -> str | None:
         """Return the fan speed of the vacuum cleaner."""
-        return self.device.status.get(DPCode.SUCTION)
+        return self._read_wrapper(self._fan_speed_wrapper)
 
     @property
     def activity(self) -> VacuumActivity | None:
@@ -134,32 +159,34 @@ def activity(self) -> VacuumActivity | None:
             return None
         return TUYA_STATUS_TO_HA.get(status)
 
-    def start(self, **kwargs: Any) -> None:
+    async def async_start(self, **kwargs: Any) -> None:
         """Start the device."""
-        self._send_command([{"code": DPCode.POWER_GO, "value": True}])
+        await self._async_send_dpcode_update(self._switch_wrapper, True)
 
-    def stop(self, **kwargs: Any) -> None:
+    async def async_stop(self, **kwargs: Any) -> None:
         """Stop the device."""
-        self._send_command([{"code": DPCode.POWER_GO, "value": False}])
+        await self._async_send_dpcode_update(self._switch_wrapper, False)
 
     def pause(self, **kwargs: Any) -> None:
         """Pause the device."""
         self._send_command([{"code": DPCode.POWER_GO, "value": False}])
 
-    def return_to_base(self, **kwargs: Any) -> None:
+    async def async_return_to_base(self, **kwargs: Any) -> None:
         """Return device to dock."""
-        if self._return_home_use_switch_charge:
-            self._send_command([{"code": DPCode.SWITCH_CHARGE, "value": True}])
+        if self._charge_wrapper:
+            await self._async_send_dpcode_update(self._charge_wrapper, True)
         else:
-            self._send_command([{"code": DPCode.MODE, "value": TUYA_MODE_RETURN_HOME}])
+            await self._async_send_dpcode_update(
+                self._mode_wrapper, TUYA_MODE_RETURN_HOME
+            )
 
-    def locate(self, **kwargs: Any) -> None:
+    async def async_locate(self, **kwargs: Any) -> None:
         """Locate the device."""
-        self._send_command([{"code": DPCode.SEEK, "value": True}])
+        await self._async_send_dpcode_update(self._locate_wrapper, True)
 
-    def set_fan_speed(self, fan_speed: str, **kwargs: Any) -> None:
+    async def async_set_fan_speed(self, fan_speed: str, **kwargs: Any) -> None:
         """Set fan speed."""
-        self._send_command([{"code": DPCode.SUCTION, "value": fan_speed}])
+        await self._async_send_dpcode_update(self._fan_speed_wrapper, fan_speed)
 
     def send_command(
         self,

tests/components/tuya/test_vacuum.py

@@ -2,15 +2,21 @@
 
 from __future__ import annotations
 
+from typing import Any
 from unittest.mock import patch
 
 import pytest
 from syrupy.assertion import SnapshotAssertion
 from tuya_sharing import CustomerDevice, Manager
 
 from homeassistant.components.vacuum import (
+    ATTR_FAN_SPEED,
     DOMAIN as VACUUM_DOMAIN,
+    SERVICE_LOCATE,
     SERVICE_RETURN_TO_BASE,
+    SERVICE_SET_FAN_SPEED,
+    SERVICE_START,
+    SERVICE_STOP,
 )
 from homeassistant.const import ATTR_ENTITY_ID, Platform
 from homeassistant.core import HomeAssistant
@@ -37,30 +43,77 @@ async def test_platform_setup_and_discovery(
 
 
 @pytest.mark.parametrize(
-    "mock_device_code",
-    ["sd_lr33znaodtyarrrz"],
+    ("mock_device_code", "entity_id", "service", "service_data", "expected_command"),
+    [
+        (
+            "sd_i6hyjg3af7doaswm",
+            "vacuum.hoover",
+            SERVICE_RETURN_TO_BASE,
+            {},
+            {"code": "mode", "value": "chargego"},
+        ),
+        (
+            # Based on #141278
+            "sd_lr33znaodtyarrrz",
+            "vacuum.v20",
+            SERVICE_RETURN_TO_BASE,
+            {},
+            {"code": "switch_charge", "value": True},
+        ),
+        (
+            "sd_lr33znaodtyarrrz",
+            "vacuum.v20",
+            SERVICE_SET_FAN_SPEED,
+            {ATTR_FAN_SPEED: "gentle"},
+            {"code": "suction", "value": "gentle"},
+        ),
+        (
+            "sd_i6hyjg3af7doaswm",
+            "vacuum.hoover",
+            SERVICE_LOCATE,
+            {},
+            {"code": "seek", "value": True},
+        ),
+        (
+            "sd_i6hyjg3af7doaswm",
+            "vacuum.hoover",
+            SERVICE_START,
+            {},
+            {"code": "power_go", "value": True},
+        ),
+        (
+            "sd_i6hyjg3af7doaswm",
+            "vacuum.hoover",
+            SERVICE_STOP,
+            {},
+            {"code": "power_go", "value": False},
+        ),
+    ],
 )
-async def test_return_home(
+async def test_action(
     hass: HomeAssistant,
     mock_manager: Manager,
     mock_config_entry: MockConfigEntry,
     mock_device: CustomerDevice,
+    entity_id: str,
+    service: str,
+    service_data: dict[str, Any],
+    expected_command: dict[str, Any],
 ) -> None:
-    """Test return home service."""
-    # Based on #141278
-    entity_id = "vacuum.v20"
+    """Test service action."""
     await initialize_entry(hass, mock_manager, mock_config_entry, mock_device)
 
     state = hass.states.get(entity_id)
     assert state is not None, f"{entity_id} does not exist"
     await hass.services.async_call(
         VACUUM_DOMAIN,
-        SERVICE_RETURN_TO_BASE,
+        service,
         {
             ATTR_ENTITY_ID: entity_id,
+            **service_data,
         },
         blocking=True,
     )
     mock_manager.send_commands.assert_called_once_with(
-        mock_device.id, [{"code": "switch_charge", "value": True}]
+        mock_device.id, [expected_command]
     )