Add tests for push updates

pywizlight/tests/fake_bulb.py

@@ -1,9 +1,9 @@
 """Start up a fake bulb to test features without a real bulb."""
-import json
-from typing import cast, Any, Callable, Dict, Tuple
-from pywizlight.protocol import WizProtocol
 import asyncio
+import json
+from typing import Any, Callable, Dict, Tuple, cast
 
+from pywizlight.protocol import WizProtocol
 
 MODULE_CONFIGS = {
     ("ESP01_SHRGB_03", "1.25.0"): {
@@ -390,6 +390,7 @@ class BulbUDPRequestHandler:
     sys_config: Dict[str, Any]  # Will be set by constructor for the actual class
     model_config: Dict[str, Any]  # Will be set by constructor for the actual class
     user_config: Dict[str, Any]
+    registration: Dict[str, Any]
     transport: asyncio.DatagramTransport
 
     def handle(self, resp: bytes, addr: Tuple[str, int]) -> None:
@@ -415,6 +416,8 @@ def handle(self, resp: bytes, addr: Tuple[str, int]) -> None:
             self.transport.sendto(bytes(json.dumps(self.model_config), "utf-8"), addr)
         elif method == "getUserConfig":
             self.transport.sendto(bytes(json.dumps(self.user_config), "utf-8"), addr)
+        elif method == "registration":
+            self.transport.sendto(bytes(json.dumps(self.registration), "utf-8"), addr)
         else:
             raise RuntimeError(f"No handler for {method}")
 
@@ -434,6 +437,11 @@ async def make_udp_fake_bulb_server(
     handler.sys_config = get_initial_sys_config(module_name, firmware_version)
     handler.model_config = get_initial_model_config(module_name, firmware_version)
     handler.user_config = get_initial_user_config(module_name, firmware_version)
+    handler.registration = {
+        "method": "registration",
+        "env": "pro",
+        "result": {"mac": "a8bb5006033d", "success": True},
+    }
 
     transport_proto = await asyncio.get_event_loop().create_datagram_endpoint(
         lambda: WizProtocol(on_response=handler.handle),
@@ -443,6 +451,19 @@ async def make_udp_fake_bulb_server(
     return transport_proto
 
 
+async def make_udp_fake_bulb_push_server() -> Tuple[
+    asyncio.BaseTransport, asyncio.BaseProtocol
+]:
+    """Configure a fake push instance."""
+    handler = BulbUDPRequestHandler()
+    transport_proto = await asyncio.get_event_loop().create_datagram_endpoint(
+        lambda: WizProtocol(on_response=lambda resp, addr: None),
+        local_addr=("127.0.0.1", 0),
+    )
+    handler.transport = cast(asyncio.DatagramTransport, transport_proto[0])
+    return transport_proto
+
+
 async def startup_bulb(
     module_name: str = "ESP01_SHRGB_03", firmware_version: str = "1.25.0"
 ) -> Tuple[Callable[[], Any], int]:

pywizlight/tests/test_bulb.py

@@ -5,9 +5,9 @@
 import pytest
 
 from pywizlight import SCENES, PilotBuilder, wizlight
+from pywizlight.bulb import states_match
 from pywizlight.bulblibrary import BulbClass, BulbType, Features, KelvinRange
 from pywizlight.discovery import discover_lights
-from pywizlight.bulb import states_match
 from pywizlight.exceptions import WizLightTimeOutError
 from pywizlight.tests.fake_bulb import startup_bulb
 

pywizlight/tests/test_bulb_invalid_module_name.py

@@ -4,8 +4,8 @@
 import pytest
 
 from pywizlight import wizlight
-from pywizlight.tests.fake_bulb import startup_bulb
 from pywizlight.exceptions import WizLightNotKnownBulb
+from pywizlight.tests.fake_bulb import startup_bulb
 
 
 @pytest.fixture()

pywizlight/tests/test_bulb_missing_kelvin_range.py

@@ -4,8 +4,8 @@
 import pytest
 
 from pywizlight import wizlight
-from pywizlight.tests.fake_bulb import startup_bulb
 from pywizlight.exceptions import WizLightNotKnownBulb
+from pywizlight.tests.fake_bulb import startup_bulb
 
 
 @pytest.fixture()

pywizlight/tests/test_bulb_no_module_name.py

@@ -4,8 +4,8 @@
 import pytest
 
 from pywizlight import wizlight
-from pywizlight.tests.fake_bulb import startup_bulb
 from pywizlight.exceptions import WizLightNotKnownBulb
+from pywizlight.tests.fake_bulb import startup_bulb
 
 
 @pytest.fixture()

pywizlight/tests/test_push_manager.py

@@ -0,0 +1,90 @@
+"""Tests for the Bulb API with a socket."""
+import asyncio
+import logging
+from typing import AsyncGenerator, cast
+from unittest.mock import patch
+
+import pytest
+
+from pywizlight import wizlight
+from pywizlight.bulb import PilotParser
+from pywizlight.bulblibrary import BulbClass, BulbType, Features, KelvinRange
+from pywizlight.protocol import WizProtocol
+from pywizlight.push_manager import PushManager
+from pywizlight.tests.fake_bulb import startup_bulb
+from pywizlight.utils import to_wiz_json
+
+logging.getLogger("pywizlight").setLevel(logging.DEBUG)
+
+
+@pytest.fixture()
+async def socket_push() -> AsyncGenerator[wizlight, None]:
+    shutdown, port = await startup_bulb(
+        module_name="ESP10_SOCKET_06", firmware_version="1.25.0"
+    )
+    with patch("pywizlight.push_manager.RESPOND_PORT", port):
+        bulb = wizlight(ip="127.0.0.1", port=port)
+        yield bulb
+        await bulb.async_close()
+        shutdown()
+
+
+@pytest.mark.asyncio
+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),
+        name="ESP10_SOCKET_06",
+        kelvin_range=KelvinRange(max=2700, min=2700),
+        bulb_type=BulbClass.SOCKET,
+        fw_version="1.25.0",
+        white_channels=2,
+        white_to_color_ratio=20,
+    )
+    last_data = PilotParser({})
+    data_event = asyncio.Event()
+
+    def _on_push(data: PilotParser) -> None:
+        nonlocal last_data
+        last_data = data
+        data_event.set()
+
+    with patch("pywizlight.push_manager.LISTEN_PORT", 0):
+        await socket_push.start_push(_on_push)
+
+    push_manager = PushManager().get()
+    push_port = push_manager.push_transport.get_extra_info("sockname")[1]
+    push_in_transport_proto = await asyncio.get_event_loop().create_datagram_endpoint(
+        lambda: WizProtocol(on_response=lambda resp, addr: None),
+        remote_addr=("127.0.0.1", push_port),
+    )
+    push_transport = cast(asyncio.DatagramTransport, push_in_transport_proto[0])
+    params = {
+        "mac": "a8bb5006033d",
+        "rssi": -71,
+        "src": "hb",
+        "mqttCd": 255,
+        "ts": 1644593327,
+        "state": False,
+        "sceneId": 0,
+        "temp": 6500,
+        "dimming": 100,
+    }
+
+    push_transport.sendto(
+        to_wiz_json(
+            {
+                "method": "syncPilot",
+                "env": "pro",
+                "params": params,
+            }
+        ).encode(),
+        ("127.0.0.1", push_port),
+    )
+    await asyncio.wait_for(data_event.wait(), timeout=1)
+    assert last_data.pilotResult == params
+    update = await socket_push.updateState()
+    assert update is not None
+    assert update.pilotResult == params
+    push_transport.close()