Add tests for hero device and discovery

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,6 +428,20 @@
             "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,
 }
 

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_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