Add support for dual head WiZ devices (#66955)

homeassistant/components/wiz/manifest.json

@@ -10,7 +10,7 @@
   "dependencies": ["network"],
   "quality_scale": "platinum",
   "documentation": "https://www.home-assistant.io/integrations/wiz",
-  "requirements": ["pywizlight==0.5.10"],
+  "requirements": ["pywizlight==0.5.11"],
   "iot_class": "local_push",
   "codeowners": ["@sbidy"]
 }

homeassistant/components/wiz/number.py

@@ -1,9 +1,17 @@
 """Support for WiZ effect speed numbers."""
 from __future__ import annotations
 
-from pywizlight.bulblibrary import BulbClass
+from collections.abc import Callable, Coroutine
+from dataclasses import dataclass
+from typing import Optional, cast
 
-from homeassistant.components.number import NumberEntity, NumberMode
+from pywizlight import wizlight
+
+from homeassistant.components.number import (
+    NumberEntity,
+    NumberEntityDescription,
+    NumberMode,
+)
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
@@ -12,7 +20,55 @@ from .const import DOMAIN
 from .entity import WizEntity
 from .models import WizData
 
-EFFECT_SPEED_UNIQUE_ID = "{}_effect_speed"
+
+@dataclass
+class WizNumberEntityDescriptionMixin:
+    """Mixin to describe a WiZ number entity."""
+
+    value_fn: Callable[[wizlight], int | None]
+    set_value_fn: Callable[[wizlight, int], Coroutine[None, None, None]]
+    required_feature: str
+
+
+@dataclass
+class WizNumberEntityDescription(
+    NumberEntityDescription, WizNumberEntityDescriptionMixin
+):
+    """Class to describe a WiZ number entity."""
+
+
+async def _async_set_speed(device: wizlight, speed: int) -> None:
+    await device.set_speed(speed)
+
+
+async def _async_set_ratio(device: wizlight, ratio: int) -> None:
+    await device.set_ratio(ratio)
+
+
+NUMBERS: tuple[WizNumberEntityDescription, ...] = (
+    WizNumberEntityDescription(
+        key="effect_speed",
+        min_value=10,
+        max_value=200,
+        step=1,
+        icon="mdi:speedometer",
+        name="Effect Speed",
+        value_fn=lambda device: cast(Optional[int], device.state.get_speed()),
+        set_value_fn=_async_set_speed,
+        required_feature="effect",
+    ),
+    WizNumberEntityDescription(
+        key="dual_head_ratio",
+        min_value=0,
+        max_value=100,
+        step=1,
+        icon="mdi:floor-lamp-dual",
+        name="Dual Head Ratio",
+        value_fn=lambda device: cast(Optional[int], device.state.get_ratio()),
+        set_value_fn=_async_set_ratio,
+        required_feature="dual_head",
+    ),
+)
 
 
 async def async_setup_entry(
@@ -22,37 +78,44 @@ async def async_setup_entry(
 ) -> None:
     """Set up the wiz speed number."""
     wiz_data: WizData = hass.data[DOMAIN][entry.entry_id]
-    if wiz_data.bulb.bulbtype.bulb_type != BulbClass.SOCKET:
+    async_add_entities(
-        async_add_entities([WizSpeedNumber(wiz_data, entry.title)])
+        WizSpeedNumber(wiz_data, entry.title, description)
+        for description in NUMBERS
+        if getattr(wiz_data.bulb.bulbtype.features, description.required_feature)
+    )
 
 
 class WizSpeedNumber(WizEntity, NumberEntity):
     """Defines a WiZ speed number."""
 
-    _attr_min_value = 10
+    entity_description: WizNumberEntityDescription
-    _attr_max_value = 200
-    _attr_step = 1
     _attr_mode = NumberMode.SLIDER
-    _attr_icon = "mdi:speedometer"
 
-    def __init__(self, wiz_data: WizData, name: str) -> None:
+    def __init__(
+        self, wiz_data: WizData, name: str, description: WizNumberEntityDescription
+    ) -> None:
         """Initialize an WiZ device."""
         super().__init__(wiz_data, name)
-        self._attr_unique_id = EFFECT_SPEED_UNIQUE_ID.format(self._device.mac)
+        self.entity_description = description
-        self._attr_name = f"{name} Effect Speed"
+        self._attr_unique_id = f"{self._device.mac}_{description.key}"
+        self._attr_name = f"{name} {description.name}"
         self._async_update_attrs()
 
     @property
     def available(self) -> bool:
         """Return if entity is available."""
-        return super().available and self._device.state.get_speed() is not None
+        return (
+            super().available
+            and self.entity_description.value_fn(self._device) is not None
+        )
 
     @callback
     def _async_update_attrs(self) -> None:
         """Handle updating _attr values."""
-        self._attr_value = self._device.state.get_speed()
+        if (value := self.entity_description.value_fn(self._device)) is not None:
+            self._attr_value = float(value)
 
     async def async_set_value(self, value: float) -> None:
         """Set the speed value."""
-        await self._device.set_speed(int(value))
+        await self.entity_description.set_value_fn(self._device, int(value))
         await self.coordinator.async_request_refresh()

requirements_all.txt

@@ -2057,7 +2057,7 @@ pywemo==0.7.0
 pywilight==0.0.70
 
 # homeassistant.components.wiz
-pywizlight==0.5.10
+pywizlight==0.5.11
 
 # homeassistant.components.xeoma
 pyxeoma==1.4.1

requirements_test_all.txt

@@ -1288,7 +1288,7 @@ pywemo==0.7.0
 pywilight==0.0.70
 
 # homeassistant.components.wiz
-pywizlight==0.5.10
+pywizlight==0.5.11
 
 # homeassistant.components.zerproc
 pyzerproc==0.4.8

tests/components/wiz/__init__.py

@@ -7,7 +7,7 @@ from typing import Callable
 from unittest.mock import AsyncMock, MagicMock, patch
 
 from pywizlight import SCENES, BulbType, PilotParser, wizlight
-from pywizlight.bulblibrary import FEATURE_MAP, BulbClass, KelvinRange
+from pywizlight.bulblibrary import BulbClass, Features, KelvinRange
 from pywizlight.discovery import DiscoveredBulb
 
 from homeassistant.components.wiz.const import DOMAIN
@@ -84,10 +84,23 @@ REAL_BULB_CONFIG = json.loads(
     "ewfHex":"ff00ffff000000",\
     "ping":0}}'
 )
+FAKE_DUAL_HEAD_RGBWW_BULB = BulbType(
+    bulb_type=BulbClass.RGB,
+    name="ESP01_DHRGB_03",
+    features=Features(
+        color=True, color_tmp=True, effect=True, brightness=True, dual_head=True
+    ),
+    kelvin_range=KelvinRange(2700, 6500),
+    fw_version="1.0.0",
+    white_channels=2,
+    white_to_color_ratio=80,
+)
 FAKE_RGBWW_BULB = BulbType(
     bulb_type=BulbClass.RGB,
     name="ESP01_SHRGB_03",
-    features=FEATURE_MAP[BulbClass.RGB],
+    features=Features(
+        color=True, color_tmp=True, effect=True, brightness=True, dual_head=False
+    ),
     kelvin_range=KelvinRange(2700, 6500),
     fw_version="1.0.0",
     white_channels=2,
@@ -96,7 +109,9 @@ FAKE_RGBWW_BULB = BulbType(
 FAKE_RGBW_BULB = BulbType(
     bulb_type=BulbClass.RGB,
     name="ESP01_SHRGB_03",
-    features=FEATURE_MAP[BulbClass.RGB],
+    features=Features(
+        color=True, color_tmp=True, effect=True, brightness=True, dual_head=False
+    ),
     kelvin_range=KelvinRange(2700, 6500),
     fw_version="1.0.0",
     white_channels=1,
@@ -105,7 +120,9 @@ FAKE_RGBW_BULB = BulbType(
 FAKE_DIMMABLE_BULB = BulbType(
     bulb_type=BulbClass.DW,
     name="ESP01_DW_03",
-    features=FEATURE_MAP[BulbClass.DW],
+    features=Features(
+        color=False, color_tmp=False, effect=True, brightness=True, dual_head=False
+    ),
     kelvin_range=KelvinRange(2700, 6500),
     fw_version="1.0.0",
     white_channels=1,
@@ -114,7 +131,9 @@ FAKE_DIMMABLE_BULB = BulbType(
 FAKE_TURNABLE_BULB = BulbType(
     bulb_type=BulbClass.TW,
     name="ESP01_TW_03",
-    features=FEATURE_MAP[BulbClass.TW],
+    features=Features(
+        color=False, color_tmp=True, effect=True, brightness=True, dual_head=False
+    ),
     kelvin_range=KelvinRange(2700, 6500),
     fw_version="1.0.0",
     white_channels=1,
@@ -123,7 +142,9 @@ FAKE_TURNABLE_BULB = BulbType(
 FAKE_SOCKET = BulbType(
     bulb_type=BulbClass.SOCKET,
     name="ESP01_SOCKET_03",
-    features=FEATURE_MAP[BulbClass.SOCKET],
+    features=Features(
+        color=False, color_tmp=False, effect=False, brightness=False, dual_head=False
+    ),
     kelvin_range=KelvinRange(2700, 6500),
     fw_version="1.0.0",
     white_channels=2,
@@ -171,6 +192,7 @@ def _mocked_wizlight(device, extended_white_range, bulb_type) -> wizlight:
     bulb.start_push = AsyncMock(side_effect=_save_setup_callback)
     bulb.async_close = AsyncMock()
     bulb.set_speed = AsyncMock()
+    bulb.set_ratio = AsyncMock()
     bulb.diagnostics = {
         "mocked": "mocked",
         "roomId": 123,

tests/components/wiz/test_number.py

@@ -6,12 +6,17 @@ from homeassistant.const import ATTR_ENTITY_ID, STATE_UNAVAILABLE
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers import entity_registry as er
 
-from . import FAKE_MAC, async_push_update, async_setup_integration
+from . import (
+    FAKE_DUAL_HEAD_RGBWW_BULB,
+    FAKE_MAC,
+    async_push_update,
+    async_setup_integration,
+)
 
 
 async def test_speed_operation(hass: HomeAssistant) -> None:
     """Test changing a speed."""
-    bulb, _ = await async_setup_integration(hass)
+    bulb, _ = await async_setup_integration(hass, bulb_type=FAKE_DUAL_HEAD_RGBWW_BULB)
     await async_push_update(hass, bulb, {"mac": FAKE_MAC})
     entity_id = "number.mock_title_effect_speed"
     entity_registry = er.async_get(hass)
@@ -19,7 +24,7 @@ async def test_speed_operation(hass: HomeAssistant) -> None:
     assert hass.states.get(entity_id).state == STATE_UNAVAILABLE
 
     await async_push_update(hass, bulb, {"mac": FAKE_MAC, "speed": 50})
-    assert hass.states.get(entity_id).state == "50"
+    assert hass.states.get(entity_id).state == "50.0"
 
     await hass.services.async_call(
         NUMBER_DOMAIN,
@@ -29,4 +34,29 @@ async def test_speed_operation(hass: HomeAssistant) -> None:
     )
     bulb.set_speed.assert_called_with(30)
     await async_push_update(hass, bulb, {"mac": FAKE_MAC, "speed": 30})
-    assert hass.states.get(entity_id).state == "30"
+    assert hass.states.get(entity_id).state == "30.0"
+
+
+async def test_ratio_operation(hass: HomeAssistant) -> None:
+    """Test changing a dual head ratio."""
+    bulb, _ = await async_setup_integration(hass, bulb_type=FAKE_DUAL_HEAD_RGBWW_BULB)
+    await async_push_update(hass, bulb, {"mac": FAKE_MAC})
+    entity_id = "number.mock_title_dual_head_ratio"
+    entity_registry = er.async_get(hass)
+    assert (
+        entity_registry.async_get(entity_id).unique_id == f"{FAKE_MAC}_dual_head_ratio"
+    )
+    assert hass.states.get(entity_id).state == STATE_UNAVAILABLE
+
+    await async_push_update(hass, bulb, {"mac": FAKE_MAC, "ratio": 50})
+    assert hass.states.get(entity_id).state == "50.0"
+
+    await hass.services.async_call(
+        NUMBER_DOMAIN,
+        SERVICE_SET_VALUE,
+        {ATTR_ENTITY_ID: entity_id, ATTR_VALUE: 30},
+        blocking=True,
+    )
+    bulb.set_ratio.assert_called_with(30)
+    await async_push_update(hass, bulb, {"mac": FAKE_MAC, "ratio": 30})
+    assert hass.states.get(entity_id).state == "30.0"