ria-toolkit-oss/tests/agent/test_integration.py

"""End-to-end: local websockets server drives a Streamer with a MockSDR."""

from __future__ import annotations

import asyncio
import json

import websockets

from ria_toolkit_oss.agent.streamer import Streamer
from ria_toolkit_oss.agent.ws_client import WsClient
from ria_toolkit_oss.sdr.mock import MockSDR


def test_server_start_stream_stop_cycle_over_real_ws():
    async def scenario():
        control_frames: list[dict] = []
        binary_frames: list[bytes] = []
        ready = asyncio.Event()
        stopped = asyncio.Event()

        async def server_handler(ws):
            # Agent will open the connection; wait for heartbeat first.
            try:
                first = await asyncio.wait_for(ws.recv(), timeout=2.0)
                control_frames.append(json.loads(first))
                await ws.send(
                    json.dumps(
                        {
                            "type": "start",
                            "app_id": "app-1",
                            "radio_config": {
                                "device": "mock",
                                "buffer_size": 32,
                                "sample_rate": 1_000_000,
                                "center_frequency": 2.45e9,
                            },
                        }
                    )
                )
                while len(binary_frames) < 3:
                    msg = await asyncio.wait_for(ws.recv(), timeout=2.0)
                    if isinstance(msg, bytes):
                        binary_frames.append(msg)
                    else:
                        control_frames.append(json.loads(msg))
                ready.set()
                await ws.send(json.dumps({"type": "stop", "app_id": "app-1"}))
                # Drain final status frame.
                try:
                    while True:
                        msg = await asyncio.wait_for(ws.recv(), timeout=0.5)
                        if isinstance(msg, bytes):
                            binary_frames.append(msg)
                        else:
                            control_frames.append(json.loads(msg))
                except (asyncio.TimeoutError, Exception):
                    pass
                stopped.set()
            except Exception:
                stopped.set()

        server = await websockets.serve(server_handler, "127.0.0.1", 0)
        port = server.sockets[0].getsockname()[1]
        try:
            client = WsClient(
                f"ws://127.0.0.1:{port}",
                token="",
                heartbeat_interval=10.0,
                reconnect_pause=0.05,
            )
            streamer = Streamer(ws=client, sdr_factory=lambda d, i: MockSDR(buffer_size=32, seed=0))
            task = asyncio.create_task(client.run(on_message=streamer.on_message, heartbeat=streamer.build_heartbeat))
            await asyncio.wait_for(ready.wait(), timeout=3.0)
            await asyncio.wait_for(stopped.wait(), timeout=3.0)
            client.stop()
            task.cancel()
            try:
                await task
            except (asyncio.CancelledError, Exception):
                pass
        finally:
            server.close()
            await server.wait_closed()
        return control_frames, binary_frames

    controls, binaries = asyncio.run(scenario())

    # Heartbeat reached the server.
    assert any(f.get("type") == "heartbeat" for f in controls)
    # Status transitioned idle -> streaming -> idle.
    statuses = [f["status"] for f in controls if f.get("type") == "status"]
    assert "streaming" in statuses
    assert statuses[-1] == "idle"
    # Three binary IQ frames of 32 samples × 2 floats × 4 bytes.
    assert len(binaries) >= 3
    for b in binaries[:3]:
        assert len(b) == 32 * 2 * 4