Add support for wall switches

pywizlight/bulblibrary.py

@@ -60,28 +60,24 @@ class BulbClass(Enum):
     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: {
         "brightness": True,
         "color": False,
-        "effect": True,
         "color_tmp": True,
     },
     # Dimmable white only supports brightness and some basic effects
     BulbClass.DW: {
         "brightness": True,
         "color": False,
-        "effect": True,
         "color_tmp": False,
     },
     # Socket supports only on/off
     BulbClass.SOCKET: {
         "brightness": False,
         "color": False,
-        "effect": False,
         "color_tmp": False,
     },
 }
@@ -125,12 +121,16 @@ def from_data(
                 )
             if "RGB" in _identifier:  # full RGB bulb
                 bulb_type = BulbClass.RGB
+                effect = True
             elif "TW" in _identifier:  # Non RGB but tunable white bulb
                 bulb_type = BulbClass.TW
+                effect = True
             elif "SOCKET" in _identifier:  # A smart socket
                 bulb_type = BulbClass.SOCKET
+                effect = False
             else:  # Plain brightness-only bulb
                 bulb_type = BulbClass.DW
+                effect = "DH" in _identifier or "SH" in _identifier
             dual_head = "DH" in _identifier
         elif type_id is not None:
             if type_id not in KNOWN_TYPE_IDS:
@@ -141,6 +141,7 @@ def from_data(
                 )
             bulb_type = KNOWN_TYPE_IDS.get(type_id, BulbClass.DW)
             dual_head = False
+            effect = True
         else:
             raise WizLightNotKnownBulb(
                 f"The bulb type could not be determined from the module name: {module_name} or type_id"
@@ -157,7 +158,9 @@ def from_data(
         else:
             kelvin_range = None
 
-        features = Features(**_BASE_FEATURE_MAP[bulb_type], dual_head=dual_head)
+        features = Features(
+            **_BASE_FEATURE_MAP[bulb_type], dual_head=dual_head, effect=effect
+        )
 
         return BulbType(
             bulb_type=bulb_type,

pywizlight/tests/fake_bulb.py

@@ -114,6 +114,43 @@
             "drvIface": 0,
         },
     },
+    ("ESP20_SHRGB_01BT", "1.23.70"): {
+        "method": "getModelConfig",
+        "env": "pro",
+        "result": {
+            "ps": 1,
+            "pwmFreq": 2000,
+            "pwmRange": [0, 100],
+            "wcr": 20,
+            "nowc": 2,
+            "cctRange": [2200, 2700, 6500, 6500],
+            "renderFactor": [200, 255, 255, 150, 0, 0, 40, 0, 0, 0],
+            "hasAdjMinDim": 0,
+            "hasTapSensor": 1,
+            "i2cTapSensor": [
+                {
+                    "chip": "SC7A22",
+                    "whoami": [24, 25, 15, 19],
+                    "freq": 100,
+                    "shockWin": 40,
+                },
+                {
+                    "chip": "IIS2DLPC",
+                    "whoami": [24, 25, 15, 68],
+                    "freq": 100,
+                    "shockWin": 40,
+                },
+                {
+                    "chip": "DA262",
+                    "whoami": [38, 39, 1, 19],
+                    "freq": 100,
+                    "shockWin": 50,
+                },
+            ],
+            "pm": 0,
+            "fanSpeed": 0,
+        },
+    },
 }
 
 SYSTEM_CONFIGS: Dict[Tuple[str, str], Any] = {  # AKA getSystemConfig
@@ -391,7 +428,37 @@
             "ping": 0,
         },
     },
+    ("ESP20_SHRGB_01BT", "1.23.70"): {
+        "method": "getSystemConfig",
+        "env": "pro",
+        "result": {
+            "mac": "d8a0119c42df",
+            "homeId": 5385975,
+            "roomId": 0,
+            "rgn": "eu",
+            "moduleName": "ESP20_SHRGB_01BT",
+            "fwVersion": "1.23.70",
+            "groupId": 0,
+            "ping": 0,
+        },
+    },
     ("BROKEN_JSON", "1.0.0"): json.JSONDecodeError,
+    ("ESP01_DIMTRIACS_01", "1.16.68"): {
+        "method": "getSystemConfig",
+        "env": "pro",
+        "result": {
+            "mac": "a8bb50f30985",
+            "homeId": 5385975,
+            "roomId": 8201410,
+            "homeLock": False,
+            "pairingLock": False,
+            "typeId": 0,
+            "moduleName": "ESP01_DIMTRIACS_01",
+            "fwVersion": "1.16.68",
+            "groupId": 0,
+            "drvConf": [20, 1],
+        },
+    },
 }
 
 USER_CONFIGS: Dict[Tuple[str, str], Any] = {  # AKA getUserConfig
@@ -533,6 +600,20 @@
             "po": False,
         },
     },
+    ("ESP01_DIMTRIACS_01", "1.16.68"): {
+        "method": "getUserConfig",
+        "env": "pro",
+        "result": {
+            "fadeIn": 500,
+            "fadeOut": 500,
+            "fadeNight": False,
+            "dftDim": 100,
+            "pwmRange": [10, 100],
+            "whiteRange": [2700, 2700],
+            "extRange": [2700, 2700],
+            "po": True,
+        },
+    },
 }
 
 MODEL_CONFIG_NOT_FOUND = {

pywizlight/tests/test_bulb_hero_1_23_70.py

@@ -0,0 +1,36 @@
+"""Tests for the Bulb API with a Hero."""
+from typing import AsyncGenerator
+
+import pytest
+
+from pywizlight import wizlight
+from pywizlight.bulblibrary import BulbClass, BulbType, Features, KelvinRange
+from pywizlight.tests.fake_bulb import startup_bulb
+
+
+@pytest.fixture()
+async def hero() -> AsyncGenerator[wizlight, None]:
+    shutdown, port = await startup_bulb(
+        module_name="ESP20_SHRGB_01BT", firmware_version="1.23.70"
+    )
+    bulb = wizlight(ip="127.0.0.1", port=port)
+    yield bulb
+    await bulb.async_close()
+    shutdown()
+
+
+@pytest.mark.asyncio
+async def test_model_description_squire(hero: wizlight) -> None:
+    """Test fetching the model description for a hero."""
+    bulb_type = await hero.get_bulbtype()
+    assert bulb_type == BulbType(
+        features=Features(
+            color=True, color_tmp=True, effect=True, brightness=True, dual_head=False
+        ),
+        name="ESP20_SHRGB_01BT",
+        kelvin_range=KelvinRange(max=6500, min=2200),
+        bulb_type=BulbClass.RGB,
+        fw_version="1.23.70",
+        white_channels=2,
+        white_to_color_ratio=20,
+    )

pywizlight/tests/test_bulb_squire_1_21_40.py

@@ -1,4 +1,4 @@
-"""Tests for the Bulb API with a light strip."""
+"""Tests for the Bulb API with a Squire."""
 from typing import AsyncGenerator
 
 import pytest

pywizlight/tests/test_bulb_unknown_1_8_0.py

@@ -1,4 +1,4 @@
-"""Tests for the Bulb API with a light strip."""
+"""Tests for the Bulb API with an unknown bulb."""
 from typing import AsyncGenerator
 
 import pytest

pywizlight/tests/test_bulb_wall_switch_1_16_68.py

@@ -0,0 +1,36 @@
+"""Tests for the Bulb API with a wall switch dimmer."""
+from typing import AsyncGenerator
+
+import pytest
+
+from pywizlight import wizlight
+from pywizlight.bulblibrary import BulbClass, BulbType, Features, KelvinRange
+from pywizlight.tests.fake_bulb import startup_bulb
+
+
+@pytest.fixture()
+async def wall_switch() -> AsyncGenerator[wizlight, None]:
+    shutdown, port = await startup_bulb(
+        module_name="ESP01_DIMTRIACS_01", firmware_version="1.16.68"
+    )
+    bulb = wizlight(ip="127.0.0.1", port=port)
+    yield bulb
+    await bulb.async_close()
+    shutdown()
+
+
+@pytest.mark.asyncio
+async def test_model_description_wall_switch(wall_switch: wizlight) -> None:
+    """Test fetching the model description wall switch."""
+    bulb_type = await wall_switch.get_bulbtype()
+    assert bulb_type == BulbType(
+        features=Features(
+            color=False, color_tmp=False, effect=False, brightness=True, dual_head=False
+        ),
+        name="ESP01_DIMTRIACS_01",
+        kelvin_range=KelvinRange(max=2700, min=2700),
+        bulb_type=BulbClass.DW,
+        fw_version="1.16.68",
+        white_channels=1,
+        white_to_color_ratio=20,
+    )

pywizlight/tests/test_discovery.py

@@ -0,0 +1,94 @@
+"""Tests for discovery."""
+import asyncio
+import contextlib
+import logging
+from typing import AsyncGenerator, Tuple
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from pywizlight.discovery import (
+    BroadcastProtocol,
+    DiscoveredBulb,
+    discover_lights,
+    find_wizlights,
+)
+
+logging.getLogger("pywizlight").setLevel(logging.DEBUG)
+
+
+@pytest.fixture
+async def mock_discovery_aio_protocol() -> AsyncGenerator:
+    """Fixture to mock an asyncio connection."""
+    loop = asyncio.get_running_loop()
+    future: asyncio.Future[
+        Tuple[asyncio.DatagramProtocol, BroadcastProtocol]
+    ] = asyncio.Future()
+
+    async def _wait_for_connection():
+        transport, protocol = await future
+        await asyncio.sleep(0)
+        await asyncio.sleep(0)
+        return transport, protocol
+
+    async def _mock_create_datagram_endpoint(func, sock=None):
+        protocol: BroadcastProtocol = func()
+        transport = MagicMock()
+        protocol.connection_made(transport)
+        with contextlib.suppress(asyncio.InvalidStateError):
+            future.set_result((transport, protocol))
+        return transport, protocol
+
+    with patch.object(loop, "create_datagram_endpoint", _mock_create_datagram_endpoint):
+        yield _wait_for_connection
+
+
+@pytest.mark.asyncio
+async def test_find_wizlights(mock_discovery_aio_protocol):
+    """Test find_wizlights."""
+    task = asyncio.create_task(
+        find_wizlights(wait_time=0.02, broadcast_address="192.168.213.252")
+    )
+    transport_protocol = await mock_discovery_aio_protocol()
+    protocol: BroadcastProtocol = transport_protocol[1]
+    protocol.datagram_received(
+        b'{"method":"registration","env":"pro","result":{"mac":"d8a01199cf31","success":true}}',
+        ("1.3.4.2", 1234),
+    )
+    protocol.datagram_received(
+        b"garbage",
+        ("1.3.4.2", 1234),
+    )
+    protocol.connection_lost(None)
+    bulbs = await task
+    assert bulbs == [DiscoveredBulb(ip_address="1.3.4.2", mac_address="d8a01199cf31")]
+
+
+@pytest.mark.asyncio
+async def test_discover_lights_fails(mock_discovery_aio_protocol):
+    """Test discover_lights."""
+    task = asyncio.create_task(discover_lights(wait_time=0.02))
+    transport_protocol = await mock_discovery_aio_protocol()
+    protocol: BroadcastProtocol = transport_protocol[1]
+    protocol.connection_lost(OSError)
+    with pytest.raises(OSError):
+        await task
+
+
+@pytest.mark.asyncio
+async def test_discover_lights(mock_discovery_aio_protocol):
+    """Test discover_lights success."""
+    task = asyncio.create_task(discover_lights(wait_time=0.02))
+    transport_protocol = await mock_discovery_aio_protocol()
+    protocol: BroadcastProtocol = transport_protocol[1]
+    protocol.datagram_received(
+        b'{"method":"registration","env":"pro","result":{"mac":"d8a01199cf31","success":true}}',
+        ("1.3.4.2", 1234),
+    )
+    protocol.datagram_received(
+        b'{"method":"registration","env":"pro","result":{"mac":"d8a01199cf32","success":true}}',
+        ("1.3.4.3", 1234),
+    )
+    protocol.connection_lost(None)
+    bulbs = await task
+    assert len(bulbs) == 2