Add wrapper class for integer values in Tuya models (home-assistant#156039)

Co-authored-by: Copilot <[email protected]>

Author: epenet

homeassistant/components/tuya/models.py

@@ -120,15 +120,16 @@ def from_json(cls, dpcode: DPCode, data: str) -> Self | None:
 }
 
 
-@dataclass
 class DPCodeWrapper:
     """Base DPCode wrapper.
 
     Used as a common interface for referring to a DPCode, and
     access read conversion routines.
     """
 
-    dpcode: str
+    def __init__(self, dpcode: str) -> None:
+        """Init DPCodeWrapper."""
+        self.dpcode = dpcode
 
     def _read_device_status_raw(self, device: CustomerDevice) -> Any | None:
         """Read the raw device status for the DPCode.
@@ -139,10 +140,9 @@ def _read_device_status_raw(self, device: CustomerDevice) -> Any | None:
 
     def read_device_status(self, device: CustomerDevice) -> Any | None:
         """Read the device value for the dpcode."""
-        raise NotImplementedError("read_device_value must be implemented")
+        raise NotImplementedError("read_device_status must be implemented")
 
 
-@dataclass
 class DPCodeBooleanWrapper(DPCodeWrapper):
     """Simple wrapper for boolean values.
 
@@ -156,11 +156,39 @@ def read_device_status(self, device: CustomerDevice) -> bool | None:
         return None
 
 
-@dataclass(kw_only=True)
-class DPCodeEnumWrapper(DPCodeWrapper):
+class DPCodeTypeInformationWrapper[T: TypeInformation](DPCodeWrapper):
+    """Base DPCode wrapper with Type Information."""
+
+    DPTYPE: DPType
+    type_information: T
+
+    def __init__(self, dpcode: str, type_information: T) -> None:
+        """Init DPCodeWrapper."""
+        super().__init__(dpcode)
+        self.type_information = type_information
+
+    @classmethod
+    def find_dpcode(
+        cls,
+        device: CustomerDevice,
+        dpcodes: str | DPCode | tuple[DPCode, ...],
+        *,
+        prefer_function: bool = False,
+    ) -> Self | None:
+        """Find and return a DPCodeTypeInformationWrapper for the given DP codes."""
+        if type_information := find_dpcode(  # type: ignore[call-overload]
+            device, dpcodes, dptype=cls.DPTYPE, prefer_function=prefer_function
+        ):
+            return cls(
+                dpcode=type_information.dpcode, type_information=type_information
+            )
+        return None
+
+
+class DPCodeEnumWrapper(DPCodeTypeInformationWrapper[EnumTypeData]):
     """Simple wrapper for EnumTypeData values."""
 
-    type_information: EnumTypeData
+    DPTYPE = DPType.ENUM
 
     def read_device_status(self, device: CustomerDevice) -> str | None:
         """Read the device value for the dpcode.
@@ -174,20 +202,20 @@ def read_device_status(self, device: CustomerDevice) -> str | None:
             return raw_value
         return None
 
-    @classmethod
-    def find_dpcode(
-        cls,
-        device: CustomerDevice,
-        dpcodes: str | DPCode | tuple[DPCode, ...],
-        *,
-        prefer_function: bool = False,
-    ) -> Self | None:
-        """Find and return a DPCodeEnumWrapper for the given DP codes."""
-        if enum_type := find_dpcode(
-            device, dpcodes, dptype=DPType.ENUM, prefer_function=prefer_function
-        ):
-            return cls(dpcode=enum_type.dpcode, type_information=enum_type)
-        return None
+
+class DPCodeIntegerWrapper(DPCodeTypeInformationWrapper[IntegerTypeData]):
+    """Simple wrapper for IntegerTypeData values."""
+
+    DPTYPE = DPType.INTEGER
+
+    def read_device_status(self, device: CustomerDevice) -> float | None:
+        """Read the device value for the dpcode.
+
+        Value will be scaled based on the Integer type information.
+        """
+        if (raw_value := self._read_device_status_raw(device)) is None:
+            return None
+        return raw_value / (10**self.type_information.scale)
 
 
 @overload

homeassistant/components/tuya/number.py

@@ -23,11 +23,9 @@
     TUYA_DISCOVERY_NEW,
     DeviceCategory,
     DPCode,
-    DPType,
 )
 from .entity import TuyaEntity
-from .models import IntegerTypeData, find_dpcode
-from .util import ActionDPCodeNotFoundError
+from .models import DPCodeIntegerWrapper, IntegerTypeData
 
 NUMBERS: dict[DeviceCategory, tuple[NumberEntityDescription, ...]] = {
     DeviceCategory.BH: (
@@ -456,9 +454,13 @@ def async_discover_device(device_ids: list[str]) -> None:
             device = manager.device_map[device_id]
             if descriptions := NUMBERS.get(device.category):
                 entities.extend(
-                    TuyaNumberEntity(device, manager, description)
+                    TuyaNumberEntity(device, manager, description, dpcode_wrapper)
                     for description in descriptions
-                    if description.key in device.status
+                    if (
+                        dpcode_wrapper := DPCodeIntegerWrapper.find_dpcode(
+                            device, description.key, prefer_function=True
+                        )
+                    )
                 )
 
         async_add_entities(entities)
@@ -480,21 +482,19 @@ def __init__(
         device: CustomerDevice,
         device_manager: Manager,
         description: NumberEntityDescription,
+        dpcode_wrapper: DPCodeIntegerWrapper,
     ) -> None:
         """Init Tuya sensor."""
         super().__init__(device, device_manager)
         self.entity_description = description
         self._attr_unique_id = f"{super().unique_id}{description.key}"
+        self._dpcode_wrapper = dpcode_wrapper
 
-        if int_type := find_dpcode(
-            self.device, description.key, dptype=DPType.INTEGER, prefer_function=True
-        ):
-            self._number = int_type
-            self._attr_native_max_value = self._number.max_scaled
-            self._attr_native_min_value = self._number.min_scaled
-            self._attr_native_step = self._number.step_scaled
-            if description.native_unit_of_measurement is None:
-                self._attr_native_unit_of_measurement = int_type.unit
+        self._attr_native_max_value = dpcode_wrapper.type_information.max_scaled
+        self._attr_native_min_value = dpcode_wrapper.type_information.min_scaled
+        self._attr_native_step = dpcode_wrapper.type_information.step_scaled
+        if description.native_unit_of_measurement is None:
+            self._attr_native_unit_of_measurement = dpcode_wrapper.type_information.unit
 
         # Logic to ensure the set device class and API received Unit Of Measurement
         # match Home Assistants requirements.
@@ -538,26 +538,17 @@ def __init__(
     @property
     def native_value(self) -> float | None:
         """Return the entity value to represent the entity state."""
-        # Unknown or unsupported data type
-        if self._number is None:
-            return None
-
-        # Raw value
-        if (value := self.device.status.get(self.entity_description.key)) is None:
-            return None
-
-        return self._number.scale_value(value)
+        return self._dpcode_wrapper.read_device_status(self.device)
 
     def set_native_value(self, value: float) -> None:
         """Set new value."""
-        if self._number is None:
-            raise ActionDPCodeNotFoundError(self.device, self.entity_description.key)
-
         self._send_command(
             [
                 {
-                    "code": self.entity_description.key,
-                    "value": self._number.scale_value_back(value),
+                    "code": self._dpcode_wrapper.dpcode,
+                    "value": (
+                        self._dpcode_wrapper.type_information.scale_value_back(value)
+                    ),
                 }
             ]
         )

tests/components/tuya/test_number.py

@@ -15,7 +15,6 @@
 )
 from homeassistant.const import ATTR_ENTITY_ID, Platform
 from homeassistant.core import HomeAssistant
-from homeassistant.exceptions import ServiceValidationError
 from homeassistant.helpers import entity_registry as er
 
 from . import initialize_entry
@@ -66,46 +65,3 @@ async def test_set_value(
     mock_manager.send_commands.assert_called_once_with(
         mock_device.id, [{"code": "delay_set", "value": 18}]
     )
-
-
-@pytest.mark.parametrize(
-    "mock_device_code",
-    ["mal_gyitctrjj1kefxp2"],
-)
-async def test_set_value_no_function(
-    hass: HomeAssistant,
-    mock_manager: Manager,
-    mock_config_entry: MockConfigEntry,
-    mock_device: CustomerDevice,
-) -> None:
-    """Test set value when no function available."""
-
-    # Mock a device with delay_set in status but not in function or status_range
-    mock_device.function.pop("delay_set")
-    mock_device.status_range.pop("delay_set")
-
-    entity_id = "number.multifunction_alarm_arm_delay"
-    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"
-    with pytest.raises(ServiceValidationError) as err:
-        await hass.services.async_call(
-            NUMBER_DOMAIN,
-            SERVICE_SET_VALUE,
-            {
-                ATTR_ENTITY_ID: entity_id,
-                ATTR_VALUE: 18,
-            },
-            blocking=True,
-        )
-    assert err.value.translation_key == "action_dpcode_not_found"
-    assert err.value.translation_placeholders == {
-        "expected": "['delay_set']",
-        "available": (
-            "['alarm_delay_time', 'alarm_time', 'master_mode', 'master_state', "
-            "'muffling', 'sub_admin', 'sub_class', 'switch_alarm_light', "
-            "'switch_alarm_propel', 'switch_alarm_sound', 'switch_kb_light', "
-            "'switch_kb_sound', 'switch_mode_sound']"
-        ),
-    }