ledgrab/server/tests/test_ddp_client.py

"""Tests for DDP client packet construction (the LED hot path)."""

import struct
from unittest.mock import MagicMock

import numpy as np
import pytest

from ledgrab.core.devices.ddp_client import DDPClient


def _make_client(rgbw: bool = False) -> DDPClient:
    """Build a DDPClient with a mocked transport that captures sendto bytes."""
    client = DDPClient(host="127.0.0.1", rgbw=rgbw)
    transport = MagicMock()
    captured: list[bytes] = []

    def _sendto(data, addr=None):  # noqa: ARG001
        # asyncio's sendto accepts bytes-like; copy to bytes for assertion stability
        captured.append(bytes(data))

    transport.sendto.side_effect = _sendto
    client._transport = transport
    client._captured = captured  # type: ignore[attr-defined]
    return client


def _parse_header(packet: bytes) -> dict:
    """Parse a 10-byte DDP header back into named fields."""
    flags, seq, dtype, src, offset, dlen = struct.unpack("!BBB B I H", packet[:10])
    return {
        "flags": flags,
        "seq": seq,
        "type": dtype,
        "src": src,
        "offset": offset,
        "data_len": dlen,
        "payload": packet[10:],
    }


def test_send_pixels_numpy_single_packet_rgb():
    client = _make_client(rgbw=False)
    pixels = np.array(
        [[10, 20, 30], [40, 50, 60], [70, 80, 90]],
        dtype=np.uint8,
    )

    client.send_pixels_numpy(pixels)

    assert len(client._captured) == 1
    parsed = _parse_header(client._captured[0])
    # PUSH on the only packet, VER=1
    assert parsed["flags"] == 0x40 | 0x01
    assert parsed["type"] == 0x01
    assert parsed["offset"] == 0
    assert parsed["data_len"] == 9
    assert parsed["payload"] == bytes([10, 20, 30, 40, 50, 60, 70, 80, 90])


def test_send_pixels_numpy_increments_sequence():
    client = _make_client(rgbw=False)
    pixels = np.array([[1, 2, 3]], dtype=np.uint8)

    client.send_pixels_numpy(pixels)
    client.send_pixels_numpy(pixels)
    client.send_pixels_numpy(pixels)

    seqs = [_parse_header(p)["seq"] for p in client._captured]
    assert seqs == [1, 2, 3]


def test_send_pixels_numpy_chunks_large_payload():
    """A 1000-LED payload (3000 bytes) splits into multiple packets."""
    client = _make_client(rgbw=False)
    n = 1000
    pixels = np.arange(n * 3, dtype=np.uint8).reshape(n, 3)

    client.send_pixels_numpy(pixels, max_packet_size=1400)

    # max_payload = 1390, bytes_per_packet = 1389 (multiple of 3) → 3 packets
    assert len(client._captured) >= 2
    headers = [_parse_header(p) for p in client._captured]

    # Only the last packet has PUSH
    assert (headers[-1]["flags"] & 0x01) == 0x01
    for h in headers[:-1]:
        assert (h["flags"] & 0x01) == 0

    # Offsets are contiguous and lengths sum to total
    total = sum(h["data_len"] for h in headers)
    assert total == n * 3
    expected_offset = 0
    for h in headers:
        assert h["offset"] == expected_offset
        expected_offset += h["data_len"]


def test_send_pixels_numpy_rgbw_pads_alpha():
    client = _make_client(rgbw=True)
    pixels = np.array([[1, 2, 3], [4, 5, 6]], dtype=np.uint8)

    client.send_pixels_numpy(pixels)

    parsed = _parse_header(client._captured[0])
    # bpp=4 → payload is 8 bytes, alpha=0
    assert parsed["data_len"] == 8
    expected = bytes([1, 2, 3, 0, 4, 5, 6, 0])
    assert parsed["payload"] == expected


def test_send_pixels_numpy_reuses_send_buffer():
    """The internal send_buf must be allocated lazily and reused — no fresh
    bytearray per call."""
    client = _make_client(rgbw=False)
    pixels = np.array([[1, 2, 3]], dtype=np.uint8)

    assert client._send_buf is None
    client.send_pixels_numpy(pixels)
    first_buf_id = id(client._send_buf)
    assert first_buf_id is not None

    client.send_pixels_numpy(pixels)
    assert id(client._send_buf) == first_buf_id


def test_send_pixels_numpy_growing_payload_grows_buffer():
    client = _make_client(rgbw=False)
    small = np.array([[1, 2, 3]], dtype=np.uint8)
    big = np.zeros((600, 3), dtype=np.uint8)

    client.send_pixels_numpy(small)
    small_capacity = len(client._send_buf)

    client.send_pixels_numpy(big, max_packet_size=2048)
    big_capacity = len(client._send_buf)

    assert big_capacity >= small_capacity


def test_send_pixels_numpy_handles_non_contiguous():
    """Slicing a wider array yields a non-contiguous view; payload must
    still serialise to the same bytes as a contiguous copy."""
    client = _make_client(rgbw=False)
    wide = np.arange(60, dtype=np.uint8).reshape(5, 4, 3)
    # take the middle pixel column → (5, 3) but possibly non-contiguous
    pixels = wide[:, 1, :]

    client.send_pixels_numpy(pixels)
    payload = _parse_header(client._captured[0])["payload"]
    assert payload == np.ascontiguousarray(pixels).tobytes()


@pytest.mark.asyncio
async def test_send_pixels_async_delegates_to_numpy():
    client = _make_client(rgbw=False)
    pixels = [(10, 20, 30), (40, 50, 60)]

    result = await client.send_pixels(pixels)
    assert result is True

    payload = _parse_header(client._captured[0])["payload"]
    assert payload == bytes([10, 20, 30, 40, 50, 60])


def test_send_pixels_numpy_raises_when_disconnected():
    client = DDPClient(host="127.0.0.1")
    pixels = np.array([[1, 2, 3]], dtype=np.uint8)
    with pytest.raises(RuntimeError):
        client.send_pixels_numpy(pixels)