101 lines
3.7 KiB
Python
101 lines
3.7 KiB
Python
"""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
|