test_e2e_wire.py

"""End-to-end wire test: boot the real Flask+WS server (mock engine) and drive it.

Verifies the binary protocol and the full->patch->idle delta path, and measures
bytes shipped for a full frame vs a small interaction.
"""
import asyncio
import json
import os
import struct
import subprocess
import sys
import time
from io import BytesIO

import urllib.request

import websockets
from PIL import Image


def get_json(url, timeout=10):
    with urllib.request.urlopen(url, timeout=timeout) as r:
        return json.loads(r.read())

PORT = 7863
BASE = f"http://127.0.0.1:{PORT}"
WS = f"ws://127.0.0.1:{PORT}/kernel"
MAGIC = 0x4E


def unpack(buf):
    assert buf[0] == MAGIC, "bad magic"
    t = buf[1]
    hlen = struct.unpack(">I", buf[2:6])[0]
    header = json.loads(buf[6:6 + hlen])
    return t, header, buf[6 + hlen:]


async def run():
    full_bytes = patch_bytes = 0
    kinds = []

    async with websockets.connect(WS, max_size=None) as ws:
        deadline = time.time() + 20
        # 1) initial full frame
        while time.time() < deadline:
            t, h, payload = unpack(await asyncio.wait_for(ws.recv(), 20))
            if t == 2:
                continue
            assert t == 0, f"expected full frame, got type {t}"
            assert h["w"] == h["h"] == 1024
            img = Image.open(BytesIO(payload)); img.load()
            assert img.size == (1024, 1024), img.size
            full_bytes = len(payload)
            kinds.append("full")
            break

        # 2) click the start button -> menu opens -> patch(es)
        await ws.send(json.dumps({"type": "click", "x": 40, "y": 999}))
        got = []
        deadline = time.time() + 10
        while time.time() < deadline and not got:
            t, h, payload = unpack(await asyncio.wait_for(ws.recv(), 10))
            if t == 2:
                continue
            got.append(t)
            patch_bytes += len(payload)
        assert got, "no frame after start click"
        kinds.append("patch" if got[0] == 1 else "full")

        # 3) Text2OS over the wire: type "notepad" + Enter
        for ch in "notepad":
            await ws.send(json.dumps({"type": "key", "key": ch}))
        await ws.send(json.dumps({"type": "key", "key": "Enter"}))
        deadline = time.time() + 10
        saw = False
        while time.time() < deadline:
            t, h, payload = unpack(await asyncio.wait_for(ws.recv(), 10))
            if t == 2:
                continue
            saw = True
            if h.get("windows", 0) >= 1:
                break
        assert saw, "no frame after Text2OS"

    # 3.5) a late-joining client must receive the WHOLE desktop, not just the
    # tiles the cursor happens to dirty. Regression: "black tiles until
    # mouseover" -- the shared delta stream left new sockets almost blank.
    await asyncio.sleep(0.5)
    async with websockets.connect(WS, max_size=None) as late:
        deadline = time.time() + 10
        got_full = None
        while time.time() < deadline:
            t, h, payload = unpack(await asyncio.wait_for(late.recv(), 10))
            if t == 2:
                continue
            got_full = t
            break
        assert got_full == 0, f"late client got frame type {got_full}, want full(0)"
        img = Image.open(BytesIO(payload)); img.load()
        assert img.size == (1024, 1024), img.size
        kinds.append("late_full")

    # 4) metrics endpoint
    m = get_json(f"{BASE}/metrics")
    print("FULL bytes:", full_bytes)
    print("interaction bytes:", patch_bytes)
    print("kinds:", kinds)
    print("metrics:", json.dumps({k: m[k] for k in ("frames", "full", "patch", "idle", "tiles", "bytes")}))
    assert m["frames"] >= 1
    assert len(kinds) >= 2
    return full_bytes, patch_bytes, m


def main():
    env = dict(os.environ, NEURALOS_ENGINE="mock", PORT=str(PORT),
               NEURALOS_CACHE="/tmp/neuralos_e2e_cache")
    proc = subprocess.Popen([sys.executable, "app.py"], env=env,
                            stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
    try:
        for _ in range(80):
            try:
                get_json(f"{BASE}/healthz", timeout=2)
                break
            except Exception:
                time.sleep(0.25)
        else:
            raise SystemExit("server did not come up")
        full_bytes, patch_bytes, m = asyncio.run(run())
        assert full_bytes > 1000, full_bytes
        print("E2E OK")
    finally:
        proc.terminate()
        try:
            proc.wait(10)
        except Exception:
            proc.kill()


if __name__ == "__main__":
    main()