led-controller/tests/test_driver_delivery_wifi.py

"""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]