"""Tests for dual-transport delivery (ESP-NOW bridge + Wi-Fi WebSocket) and Wi-Fi runtime.""" from __future__ import annotations import asyncio import json import socket import sys from pathlib import Path from typing import Any, Dict, List, Optional from unittest.mock import AsyncMock import pytest PROJECT_ROOT = Path(__file__).resolve().parents[1] SRC_PATH = PROJECT_ROOT / "src" for p in (str(PROJECT_ROOT), str(SRC_PATH)): if p in sys.path: sys.path.remove(p) sys.path.insert(0, p) _models = sys.modules.get("models") if _models is not None: _mf = (getattr(_models, "__file__", "") or "").replace("\\", "/") if "/tests/models" in _mf: for key in list(sys.modules): if key == "models" or key.startswith("models."): del sys.modules[key] from util.bridge_envelope import BROADCAST_MAC # noqa: E402 class FakeDevices: def __init__(self, docs: Dict[str, Dict[str, Any]]): self._docs = docs def read(self, mac: str) -> Optional[Dict[str, Any]]: return self._docs.get(mac) def items(self): return self._docs.items() class RecordingBridge: def __init__(self) -> None: self.envelopes: List[Dict[str, Any]] = [] async def send(self, data, addr=None): del addr if isinstance(data, dict): self.envelopes.append(data) elif isinstance(data, str): self.envelopes.append(json.loads(data)) return True def mac_keys(self) -> List[str]: keys: List[str] = [] for env in self.envelopes: devs = env.get("dv") or env.get("devices") or {} keys.extend(devs.keys()) return keys @pytest.fixture def bridge(): return RecordingBridge() @pytest.fixture def espnow_devices(): return FakeDevices( { "188b0e1560a8": { "id": "188b0e1560a8", "name": "esp-a", "transport": "espnow", "address": "188b0e1560a8", }, "e8f60a16ea10": { "id": "e8f60a16ea10", "name": "esp-b", "transport": "espnow", "address": "e8f60a16ea10", }, } ) @pytest.fixture def mixed_devices(): return FakeDevices( { "188b0e1560a8": { "id": "188b0e1560a8", "name": "esp-a", "transport": "espnow", "address": "188b0e1560a8", }, "102030405060": { "id": "102030405060", "name": "wifi-a", "transport": "wifi", "address": "192.168.50.10", }, } ) def test_wifi_message_for_device_narrows_select(): from util.driver_delivery import _wifi_message_for_device msg = json.dumps( {"v": "1", "select": {"wifi-a": 0, "esp-a": 1}}, separators=(",", ":"), ) narrowed = _wifi_message_for_device(msg, "wifi-a") body = json.loads(narrowed) assert body["select"] == {"wifi-a": 0} def test_combine_preset_chunks_for_wifi(): from util.driver_delivery import _combine_preset_chunks_for_wifi chunks = [ json.dumps({"v": "1", "presets": {"a": {"p": "on"}}}, separators=(",", ":")), json.dumps( {"v": "1", "presets": {"b": {"p": "blink"}}, "save": True, "default": "b"}, separators=(",", ":"), ), ] combined = json.loads(_combine_preset_chunks_for_wifi(chunks)) assert combined["presets"]["a"]["p"] == "on" assert combined["presets"]["b"]["p"] == "blink" assert combined["save"] is True assert combined["default"] == "b" def test_deliver_json_broadcast_espnow_only(bridge, espnow_devices, monkeypatch): from util import driver_delivery wifi_sends: list[tuple[str, str]] = [] async def fake_wifi(ip, msg): wifi_sends.append((ip, msg)) return True monkeypatch.setattr(driver_delivery, "send_json_line_to_ip", fake_wifi) async def _run(): return await driver_delivery.deliver_json_messages( bridge, [json.dumps({"v": "1", "select": ["off"]})], None, espnow_devices, delay_s=0, ) deliveries, n = asyncio.run(_run()) assert n == 1 assert deliveries >= 1 assert bridge.mac_keys() == [BROADCAST_MAC] assert wifi_sends == [] def test_deliver_json_broadcast_includes_wifi(bridge, mixed_devices, monkeypatch): from util import driver_delivery wifi_sends: list[tuple[str, str]] = [] async def fake_wifi(ip, msg): wifi_sends.append((ip, msg)) return True monkeypatch.setattr(driver_delivery, "send_json_line_to_ip", fake_wifi) async def _run(): return await driver_delivery.deliver_json_messages( bridge, [json.dumps({"v": "1", "select": ["off"]})], None, mixed_devices, delay_s=0, ) deliveries, _n = asyncio.run(_run()) assert deliveries >= 2 assert bridge.mac_keys() == [BROADCAST_MAC] assert len(wifi_sends) == 1 assert wifi_sends[0][0] == "192.168.50.10" def test_deliver_json_targeted_espnow_unicasts(bridge, espnow_devices, monkeypatch): from util import driver_delivery monkeypatch.setattr( driver_delivery, "send_json_line_to_ip", AsyncMock(return_value=True), ) async def _run(): return await driver_delivery.deliver_json_messages( bridge, [json.dumps({"v": "1", "select": ["2"]})], ["188b0e1560a8", "e8f60a16ea10"], espnow_devices, delay_s=0, ) asyncio.run(_run()) keys = bridge.mac_keys() assert "18:8b:0e:15:60:a8" in keys assert "e8:f6:0a:16:ea:10" in keys assert BROADCAST_MAC not in keys def test_deliver_json_targeted_wifi_uses_websocket(bridge, mixed_devices, monkeypatch): from util import driver_delivery wifi_sends: list[tuple[str, str]] = [] async def fake_wifi(ip, msg): wifi_sends.append((ip, msg)) return True monkeypatch.setattr(driver_delivery, "send_json_line_to_ip", fake_wifi) async def _run(): await driver_delivery.deliver_json_messages( bridge, [json.dumps({"v": "1", "select": {"wifi-a": 0}})], ["102030405060"], mixed_devices, delay_s=0, ) asyncio.run(_run()) assert bridge.mac_keys() == [] assert len(wifi_sends) == 1 assert wifi_sends[0][0] == "192.168.50.10" body = json.loads(wifi_sends[0][1]) assert body["select"] == {"wifi-a": 0} def test_deliver_json_unicast_flag_wifi(bridge, mixed_devices, monkeypatch): from util import driver_delivery wifi_sends: list[str] = [] async def fake_wifi(ip, msg): wifi_sends.append(msg) return True monkeypatch.setattr(driver_delivery, "send_json_line_to_ip", fake_wifi) async def _run(): await driver_delivery.deliver_json_messages( bridge, [json.dumps({"v": "1", "b": 128})], ["102030405060"], mixed_devices, delay_s=0, unicast=True, ) asyncio.run(_run()) assert len(wifi_sends) == 1 assert bridge.mac_keys() == [] def test_deliver_preset_broadcast_then_per_device_wifi( bridge, mixed_devices, monkeypatch ): from util import driver_delivery wifi_sends: list[str] = [] async def fake_wifi(ip, msg): wifi_sends.append(msg) return True monkeypatch.setattr(driver_delivery, "send_json_line_to_ip", fake_wifi) chunks = [ json.dumps( {"v": "1", "presets": {"p1": {"p": "on"}}, "save": True}, separators=(",", ":"), ) ] async def _run(): return await driver_delivery.deliver_preset_broadcast_then_per_device( bridge, chunks, None, mixed_devices, default_id=None, delay_s=0, ) count = asyncio.run(_run()) assert count >= 2 assert bridge.mac_keys() == [BROADCAST_MAC] assert len(wifi_sends) == 1 combined = json.loads(wifi_sends[0]) assert "p1" in combined["presets"] def test_deliver_json_requires_bridge(monkeypatch): from util import driver_delivery import models.transport as transport_mod monkeypatch.setattr(transport_mod, "get_current_bridge", lambda: None) async def _run(): with pytest.raises(RuntimeError, match="Transport not configured"): await driver_delivery.deliver_json_messages( None, ["{}"], None, FakeDevices({}), delay_s=0 ) asyncio.run(_run()) def test_device_status_broadcaster_send_text(): from util.device_status_broadcaster import ( _ws_send_text, broadcast_device_tcp_snapshot_to, broadcast_device_tcp_status, register_device_status_ws, unregister_device_status_ws, ) class StarletteLikeWS: def __init__(self): self.out: list[str] = [] async def send_text(self, msg: str): self.out.append(msg) class SendTextOnlyWS: def __init__(self): self.out: list[str] = [] async def send(self, msg: str): self.out.append(msg) async def _run(): starlette = StarletteLikeWS() legacy_ws = SendTextOnlyWS() await _ws_send_text(starlette, '{"ok":true}') await _ws_send_text(legacy_ws, '{"ok":true}') assert starlette.out == ['{"ok":true}'] assert legacy_ws.out == ['{"ok":true}'] await register_device_status_ws(starlette) await broadcast_device_tcp_status("192.168.1.5", True) assert len(starlette.out) == 2 status = json.loads(starlette.out[1]) assert status["type"] == "device_tcp" assert status["ip"] == "192.168.1.5" assert status["connected"] is True await broadcast_device_tcp_snapshot_to(starlette) snapshot = json.loads(starlette.out[2]) assert snapshot["type"] == "device_tcp_snapshot" assert "connected_ips" in snapshot await unregister_device_status_ws(starlette) asyncio.run(_run()) def test_process_udp_datagram_registers_and_connects(monkeypatch): from util import wifi_driver_runtime registered: list[tuple[str, str, str]] = [] connected: list[str] = [] def fake_register(device_name, peer_ip, mac, device_type=None): del device_type registered.append((device_name, peer_ip, str(mac))) monkeypatch.setattr( wifi_driver_runtime, "_register_udp_device_sync", fake_register, ) monkeypatch.setattr( wifi_driver_runtime.tcp_client_registry, "ensure_driver_connection", lambda ip: connected.append(ip), ) line = json.dumps( {"v": "1", "device_name": "strip-a", "mac": "aabbccddeeff", "type": "led"} ).encode() wifi_driver_runtime._process_udp_datagram(line, "192.168.1.42") assert registered == [("strip-a", "192.168.1.42", "aabbccddeeff")] assert connected == ["192.168.1.42"] def test_process_udp_datagram_ignores_invalid(): from util.wifi_driver_runtime import _process_udp_datagram _process_udp_datagram(b"not-json\n", "10.0.0.1") _process_udp_datagram(b'{"v":"1"}\n', "10.0.0.1") def test_discovery_protocol_uses_datagram_endpoint(monkeypatch): pytest.importorskip("uvloop") import uvloop from util.wifi_driver_runtime import _DiscoveryProtocol async def _run(): echoed: list[bytes] = [] asyncio.set_event_loop_policy(uvloop.EventLoopPolicy()) holder: dict = {"closing": False} loop = asyncio.get_running_loop() sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) sock.bind(("127.0.0.1", 0)) port = sock.getsockname()[1] transport, _protocol = await loop.create_datagram_endpoint( lambda: _DiscoveryProtocol(holder), sock=sock, ) class _EchoClient(asyncio.DatagramProtocol): def connection_made(self, t): self._transport = t def datagram_received(self, data, addr): del addr echoed.append(data) client_sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) client_sock.bind(("127.0.0.1", 0)) client_transport, _ = await loop.create_datagram_endpoint( _EchoClient, sock=client_sock, ) payload = b'{"v":"1","device_name":"x","mac":"112233445566"}\n' client_transport.sendto(payload, ("127.0.0.1", port)) await asyncio.sleep(0.05) holder["closing"] = True client_transport.close() transport.close() return echoed, payload monkeypatch.setattr( "util.wifi_driver_runtime._register_udp_device_sync", lambda *a, **k: None, ) monkeypatch.setattr( "util.wifi_driver_runtime.tcp_client_registry.ensure_driver_connection", lambda _ip: None, ) echoed, payload = asyncio.run(_run()) assert echoed == [payload]