Add support for Squire dual head lamp

pywizlight/bulb.py

@@ -106,6 +106,11 @@ def _validate_speed_or_raise(speed: int) -> None:
         raise ValueError("Value must be between 10 and 200")
 
 
+def _validate_ratio_or_raise(ratio: int) -> None:
+    if not 1 <= ratio <= 100:
+        raise ValueError("Value must be between 1 and 100")
+
+
 class PilotBuilder:
     """Get information from the bulb."""
 
@@ -122,11 +127,14 @@ def __init__(
         brightness: Optional[int] = None,
         colortemp: Optional[int] = None,
         state: bool = True,
+        ratio: Optional[int] = None,
     ) -> None:
         """Set the parameter."""
         self.pilot_params: Dict[str, Any] = {"state": state}
         if speed is not None:
             self._set_speed(speed)
+        if ratio is not None:
+            self._set_ratio(ratio)
         if scene is not None:
             self._set_scene(scene)
         if brightness is not None:
@@ -173,6 +181,11 @@ def _set_speed(self, value: int) -> None:
         _validate_speed_or_raise(value)
         self.pilot_params["speed"] = value
 
+    def _set_ratio(self, value: int) -> None:
+        """Set the ratio between the up and down light (1-100)."""
+        _validate_ratio_or_raise(value)
+        self.pilot_params["ratio"] = value
+
     def _set_scene(self, scene_id: int) -> None:
         """Set the scene by id."""
         if scene_id not in SCENES:
@@ -289,6 +302,10 @@ def get_speed(self) -> Optional[int]:
         """Get the color changing speed."""
         return _extract_int(self.pilotResult, "speed")
 
+    def get_ratio(self) -> Optional[int]:
+        """Get the ratio between the up and down light."""
+        return _extract_int(self.pilotResult, "ratio")
+
     def get_scene_id(self) -> Optional[int]:
         if "schdPsetId" in self.pilotResult:  # rhythm
             return 1000
@@ -601,6 +618,11 @@ async def set_speed(self, speed: int) -> None:
         _validate_speed_or_raise(speed)
         await self.send({"method": "setPilot", "params": {"speed": speed}})
 
+    async def set_ratio(self, ratio: int) -> None:
+        """Set the ratio between the up and down light."""
+        _validate_ratio_or_raise(ratio)
+        await self.send({"method": "setPilot", "params": {"ratio": ratio}})
+
     async def turn_on(self, pilot_builder: PilotBuilder = PilotBuilder()) -> None:
         """Turn the light on with defined message.
 

pywizlight/bulblibrary.py

@@ -3,6 +3,7 @@
 Bulb Type detection:
 ESP01_SHDW1C_31
 ESP01 -- defines the module family (WiFi only bulb in this case)
+DH -- Dual Head light
 SH -- Single Head light (most bulbs are single heads) / LED Strip
 TW -- Tunable White - can only control CCT and dimming; no color
 DW -- Dimmable White (most filament bulbs)
@@ -25,6 +26,7 @@ class Features:
     color_tmp: bool
     effect: bool
     brightness: bool
+    dual_head: bool
 
 
 @dataclasses.dataclass(frozen=True)
@@ -48,17 +50,35 @@ class BulbClass(Enum):
     """Smart socket with only on/off."""
 
 
-FEATURE_MAP = {
+_BASE_FEATURE_MAP = {
     # RGB supports effects and tuneable white
-    BulbClass.RGB: Features(brightness=True, color=True, effect=True, color_tmp=True),
+    BulbClass.RGB: {
+        "brightness": True,
+        "color": True,
+        "effect": True,
+        "color_tmp": True,
+    },
     # TODO: TW supports effects but only "some"; improve the mapping to supported effects
-    BulbClass.TW: Features(brightness=True, color=False, effect=True, color_tmp=True),
+    BulbClass.TW: {
+        "brightness": True,
+        "color": False,
+        "effect": True,
+        "color_tmp": True,
+    },
     # Dimmable white only supports brightness and some basic effects
-    BulbClass.DW: Features(brightness=True, color=False, effect=True, color_tmp=False),
+    BulbClass.DW: {
+        "brightness": True,
+        "color": False,
+        "effect": True,
+        "color_tmp": False,
+    },
     # Socket supports only on/off
-    BulbClass.SOCKET: Features(
-        brightness=False, color=False, effect=False, color_tmp=False
-    ),
+    BulbClass.SOCKET: {
+        "brightness": False,
+        "color": False,
+        "effect": False,
+        "color_tmp": False,
+    },
 }
 
 
@@ -117,7 +137,9 @@ def from_data(
         else:
             kelvin_range = None
 
-        features = FEATURE_MAP[bulb_type]
+        features = Features(
+            **_BASE_FEATURE_MAP[bulb_type], dual_head="DH" in _identifier
+        )
 
         return BulbType(
             bulb_type=bulb_type,

pywizlight/tests/fake_bulb.py

@@ -331,6 +331,22 @@
             "drvConf": [20, 1],
         },
     },
+    ("ESP20_DHRGB_01B", "1.21.40"): {
+        "method": "getSystemConfig",
+        "env": "pro",
+        "result": {
+            "mac": "d8a01199cf31",
+            "homeId": 5385975,
+            "roomId": 8201410,
+            "moduleName": "ESP20_DHRGB_01B",
+            "fwVersion": "1.21.40",
+            "groupId": 0,
+            "drvConf": [20, 2],
+            "ewf": [200, 255, 150, 255, 0, 0, 40],
+            "ewfHex": "c8ff96ff000028",
+            "ping": 0,
+        },
+    },
 }
 
 USER_CONFIGS = {  # AKA getUserConfig
@@ -419,6 +435,20 @@
             "extRange": [2700, 6500],
         },
     },
+    ("ESP20_DHRGB_01B", "1.21.40"): {
+        "method": "getUserConfig",
+        "env": "pro",
+        "result": {
+            "fadeIn": 500,
+            "fadeOut": 500,
+            "dftDim": 100,
+            "pwmRange": [0, 100],
+            "whiteRange": [2700, 6500],
+            "extRange": [2200, 6500],
+            "opMode": 0,
+            "po": True,
+        },
+    },
 }
 
 

pywizlight/tests/test_bulb.py

@@ -246,7 +246,9 @@ async def test_fw_version(correct_bulb: wizlight) -> None:
     """Test fetching the firmware version."""
     bulb_type = await correct_bulb.get_bulbtype()
     assert bulb_type == BulbType(
-        features=Features(color=True, color_tmp=True, effect=True, brightness=True),
+        features=Features(
+            color=True, color_tmp=True, effect=True, brightness=True, dual_head=False
+        ),
         name="ESP01_SHRGB_03",
         kelvin_range=KelvinRange(max=6500, min=2200),
         bulb_type=BulbClass.RGB,

pywizlight/tests/test_bulb_dimmable_white.py

@@ -24,7 +24,9 @@ async def test_model_description_dimmable_bulb(dimmable_bulb: wizlight) -> None:
     """Test fetching the model description dimmable bulb."""
     bulb_type = await dimmable_bulb.get_bulbtype()
     assert bulb_type == BulbType(
-        features=Features(color=False, color_tmp=False, effect=True, brightness=True),
+        features=Features(
+            color=False, color_tmp=False, effect=True, brightness=True, dual_head=False
+        ),
         name="ESP05_SHDW_21",
         kelvin_range=KelvinRange(max=2700, min=2700),
         bulb_type=BulbClass.DW,

pywizlight/tests/test_bulb_dimmable_white_1_11_7.py

@@ -29,7 +29,9 @@ async def test_model_description_dimmable_bulb(dimmable_bulb: wizlight) -> None:
     """Test fetching the model description dimmable bulb."""
     bulb_type = await dimmable_bulb.get_bulbtype()
     assert bulb_type == BulbType(
-        features=Features(color=False, color_tmp=False, effect=True, brightness=True),
+        features=Features(
+            color=False, color_tmp=False, effect=True, brightness=True, dual_head=False
+        ),
         name="ESP06_SHDW9_01",
         kelvin_range=KelvinRange(max=6500, min=2700),
         bulb_type=BulbClass.DW,

pywizlight/tests/test_bulb_light_strip_1_16_64.py

@@ -32,7 +32,9 @@ async def test_model_description_light_strip(light_strip: wizlight) -> None:
     """Test fetching the model description for a light strip."""
     bulb_type = await light_strip.get_bulbtype()
     assert bulb_type == BulbType(
-        features=Features(color=True, color_tmp=True, effect=True, brightness=True),
+        features=Features(
+            color=True, color_tmp=True, effect=True, brightness=True, dual_head=False
+        ),
         name="ESP03_SHRGB3_01ABI",
         kelvin_range=KelvinRange(max=6500, min=2700),
         bulb_type=BulbClass.RGB,

pywizlight/tests/test_bulb_light_strip_1_21_4.py

@@ -32,7 +32,9 @@ async def test_model_description_light_strip(light_strip: wizlight) -> None:
     """Test fetching the model description for a light strip."""
     bulb_type = await light_strip.get_bulbtype()
     assert bulb_type == BulbType(
-        features=Features(color=True, color_tmp=True, effect=True, brightness=True),
+        features=Features(
+            color=True, color_tmp=True, effect=True, brightness=True, dual_head=False
+        ),
         name="ESP20_SHRGB_01ABI",
         kelvin_range=KelvinRange(max=6500, min=2700),
         bulb_type=BulbClass.RGB,

pywizlight/tests/test_bulb_light_strip_1_25_0.py

@@ -72,7 +72,9 @@ async def test_model_description_light_strip(light_strip: wizlight) -> None:
     """Test fetching the model description for a light strip."""
     bulb_type = await light_strip.get_bulbtype()
     assert bulb_type == BulbType(
-        features=Features(color=True, color_tmp=True, effect=True, brightness=True),
+        features=Features(
+            color=True, color_tmp=True, effect=True, brightness=True, dual_head=False
+        ),
         name="ESP20_SHRGB_01ABI",
         kelvin_range=KelvinRange(max=6500, min=2700),
         bulb_type=BulbClass.RGB,

pywizlight/tests/test_bulb_rgbw_1_21_4.py

@@ -24,7 +24,9 @@ async def test_model_description_rgbw_bulb(rgbw_bulb: wizlight) -> None:
     """Test fetching the model description rgbw bulb."""
     bulb_type = await rgbw_bulb.get_bulbtype()
     assert bulb_type == BulbType(
-        features=Features(color=True, color_tmp=True, effect=True, brightness=True),
+        features=Features(
+            color=True, color_tmp=True, effect=True, brightness=True, dual_head=False
+        ),
         name="ESP20_SHRGBC_01",
         kelvin_range=KelvinRange(max=6500, min=2200),
         bulb_type=BulbClass.RGB,

pywizlight/tests/test_bulb_rgbww_1_17_1.py

@@ -24,7 +24,9 @@ async def test_model_description_rgbww_bulb(rgbww_bulb: wizlight) -> None:
     """Test fetching the model description rgbww bulb."""
     bulb_type = await rgbww_bulb.get_bulbtype()
     assert bulb_type == BulbType(
-        features=Features(color=True, color_tmp=True, effect=True, brightness=True),
+        features=Features(
+            color=True, color_tmp=True, effect=True, brightness=True, dual_head=False
+        ),
         name="ESP01_SHRGB1C_31",
         kelvin_range=KelvinRange(max=6500, min=2700),
         bulb_type=BulbClass.RGB,

pywizlight/tests/test_bulb_socket.py

@@ -24,7 +24,13 @@ async def test_model_description_socket(socket: wizlight) -> None:
     """Test fetching the model description of a socket is None."""
     bulb_type = await socket.get_bulbtype()
     assert bulb_type == BulbType(
-        features=Features(color=False, color_tmp=False, effect=False, brightness=False),
+        features=Features(
+            color=False,
+            color_tmp=False,
+            effect=False,
+            brightness=False,
+            dual_head=False,
+        ),
         name="ESP10_SOCKET_06",
         kelvin_range=KelvinRange(max=2700, min=2700),
         bulb_type=BulbClass.SOCKET,

pywizlight/tests/test_bulb_squire_1_21_40.py

@@ -0,0 +1,49 @@
+"""Tests for the Bulb API with a light strip."""
+from typing import AsyncGenerator
+
+import pytest
+
+from pywizlight import PilotBuilder, wizlight
+from pywizlight.bulblibrary import BulbClass, BulbType, Features, KelvinRange
+from pywizlight.tests.fake_bulb import startup_bulb
+
+
+@pytest.fixture()
+async def squire() -> AsyncGenerator[wizlight, None]:
+    shutdown, port = await startup_bulb(
+        module_name="ESP20_DHRGB_01B", firmware_version="1.21.40"
+    )
+    bulb = wizlight(ip="127.0.0.1", port=port)
+    yield bulb
+    await bulb.async_close()
+    shutdown()
+
+
+@pytest.mark.asyncio
+async def test_setting_ratio(squire: wizlight) -> None:
+    """Test setting ratio."""
+    await squire.set_ratio(50)
+    state = await squire.updateState()
+    assert state and state.get_ratio() == 50
+    await squire.turn_on(PilotBuilder(ratio=20))
+    state = await squire.updateState()
+    assert state and state.get_ratio() == 20
+    with pytest.raises(ValueError):
+        await squire.set_ratio(500)
+
+
+@pytest.mark.asyncio
+async def test_model_description_squire(squire: wizlight) -> None:
+    """Test fetching the model description for a squire."""
+    bulb_type = await squire.get_bulbtype()
+    assert bulb_type == BulbType(
+        features=Features(
+            color=True, color_tmp=True, effect=True, brightness=True, dual_head=True
+        ),
+        name="ESP20_DHRGB_01B",
+        kelvin_range=KelvinRange(max=6500, min=2200),
+        bulb_type=BulbClass.RGB,
+        fw_version="1.21.40",
+        white_channels=2,
+        white_to_color_ratio=20,
+    )

pywizlight/tests/test_bulb_turnable_white.py

@@ -24,7 +24,9 @@ async def test_model_description_dimmable_bulb(turnable_bulb: wizlight) -> None:
     """Test fetching the model description dimmable bulb."""
     bulb_type = await turnable_bulb.get_bulbtype()
     assert bulb_type == BulbType(
-        features=Features(color=False, color_tmp=True, effect=True, brightness=True),
+        features=Features(
+            color=False, color_tmp=True, effect=True, brightness=True, dual_head=False
+        ),
         name="ESP21_SHTW_01",
         kelvin_range=KelvinRange(max=5000, min=2700),
         bulb_type=BulbClass.TW,

pywizlight/tests/test_bulb_turnable_white_1_18_0.py

@@ -24,7 +24,9 @@ async def test_model_description_dimmable_bulb(turnable_bulb: wizlight) -> None:
     """Test fetching the model description dimmable bulb."""
     bulb_type = await turnable_bulb.get_bulbtype()
     assert bulb_type == BulbType(
-        features=Features(color=False, color_tmp=True, effect=True, brightness=True),
+        features=Features(
+            color=False, color_tmp=True, effect=True, brightness=True, dual_head=False
+        ),
         name="ESP14_SHTW1C_01",
         kelvin_range=KelvinRange(max=6500, min=2700),
         bulb_type=BulbClass.TW,

pywizlight/tests/test_push_manager.py

@@ -65,7 +65,13 @@ async def test_push_updates(socket_push: wizlight) -> None:
     """Test push updates."""
     bulb_type = await socket_push.get_bulbtype()
     assert bulb_type == BulbType(
-        features=Features(color=False, color_tmp=False, effect=False, brightness=False),
+        features=Features(
+            color=False,
+            color_tmp=False,
+            effect=False,
+            brightness=False,
+            dual_head=False,
+        ),
         name="ESP10_SOCKET_06",
         kelvin_range=KelvinRange(max=2700, min=2700),
         bulb_type=BulbClass.SOCKET,
@@ -140,7 +146,13 @@ async def test_discovery_by_firstbeat(
     """Test discovery from first beat."""
     bulb_type = await socket_push.get_bulbtype()
     assert bulb_type == BulbType(
-        features=Features(color=False, color_tmp=False, effect=False, brightness=False),
+        features=Features(
+            color=False,
+            color_tmp=False,
+            effect=False,
+            brightness=False,
+            dual_head=False,
+        ),
         name="ESP10_SOCKET_06",
         kelvin_range=KelvinRange(max=2700, min=2700),
         bulb_type=BulbClass.SOCKET,