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

"""Agent-side TX interlocks: gain cap, freq ranges, duplicate sessions, disabled."""

from __future__ import annotations

import asyncio

from ria_toolkit_oss.agent.config import AgentConfig
from ria_toolkit_oss.agent.streamer import Streamer
from ria_toolkit_oss.sdr.mock import MockSDR


class FakeWs:
    def __init__(self):
        self.json_sent = []
        self.bytes_sent = []

    async def send_json(self, p):
        self.json_sent.append(p)

    async def send_bytes(self, b):
        self.bytes_sent.append(b)


def _last_tx_status(ws):
    frames = [m for m in ws.json_sent if m.get("type") == "tx_status"]
    return frames[-1] if frames else None


def _tx_start(app_id="a", **radio):
    rc = {
        "device": "mock",
        "buffer_size": 16,
        "tx_sample_rate": 1_000_000,
        "tx_center_frequency": 2.45e9,
        "tx_gain": -20,
        "underrun_policy": "zero",
    }
    rc.update(radio)
    return {"type": "tx_start", "app_id": app_id, "radio_config": rc}


def _make_streamer(cfg):
    built: list = []

    def factory(device, identifier):
        sdr = MockSDR(buffer_size=16)
        built.append(sdr)
        return sdr

    ws = FakeWs()
    s = Streamer(ws=ws, sdr_factory=factory, cfg=cfg)
    return s, ws, built


def test_rejects_when_tx_disabled():
    async def scenario():
        s, ws, built = _make_streamer(AgentConfig(tx_enabled=False))
        await s.on_message(_tx_start(tx_gain=-20, tx_center_frequency=2.45e9))
        return s, ws, built

    s, ws, built = asyncio.run(scenario())
    status = _last_tx_status(ws)
    assert status and status["state"] == "error"
    assert "disabled" in status["message"].lower()
    assert not built, "SDR should never have been constructed"
    assert s._tx is None


def test_rejects_when_tx_gain_exceeds_cap():
    async def scenario():
        s, ws, built = _make_streamer(AgentConfig(tx_enabled=True, tx_max_gain_db=-15.0))
        await s.on_message(_tx_start(tx_gain=-5, tx_center_frequency=2.45e9))
        return ws, built

    ws, built = asyncio.run(scenario())
    status = _last_tx_status(ws)
    assert status and status["state"] == "error"
    assert "exceeds cap" in status["message"]
    assert not built


def test_allows_gain_at_cap_boundary():
    async def scenario():
        s, ws, _ = _make_streamer(AgentConfig(tx_enabled=True, tx_max_gain_db=-10.0))
        await s.on_message(_tx_start(tx_gain=-10, tx_center_frequency=2.45e9))
        # Stop promptly to avoid keeping an executor thread around.
        await asyncio.sleep(0.02)
        await s.on_message({"type": "tx_stop", "app_id": "a"})
        return ws

    ws = asyncio.run(scenario())
    states = [m["state"] for m in ws.json_sent if m.get("type") == "tx_status"]
    assert "armed" in states
    assert states[-1] == "done"


def test_rejects_when_freq_outside_ranges():
    async def scenario():
        s, ws, built = _make_streamer(
            AgentConfig(
                tx_enabled=True,
                tx_allowed_freq_ranges=[[2.4e9, 2.5e9]],
            )
        )
        await s.on_message(_tx_start(tx_center_frequency=5.8e9, tx_gain=-20))
        return ws, built

    ws, built = asyncio.run(scenario())
    status = _last_tx_status(ws)
    assert status and status["state"] == "error"
    assert "outside allowed ranges" in status["message"]
    assert not built


def test_allows_freq_inside_a_range():
    async def scenario():
        s, ws, _ = _make_streamer(
            AgentConfig(
                tx_enabled=True,
                tx_allowed_freq_ranges=[[2.4e9, 2.5e9], [5.7e9, 5.8e9]],
            )
        )
        await s.on_message(_tx_start(tx_center_frequency=5.75e9, tx_gain=-20))
        await asyncio.sleep(0.02)
        await s.on_message({"type": "tx_stop", "app_id": "a"})
        return ws

    ws = asyncio.run(scenario())
    states = [m["state"] for m in ws.json_sent if m.get("type") == "tx_status"]
    assert "armed" in states
    assert states[-1] == "done"


def test_rejects_duplicate_tx_session():
    async def scenario():
        s, ws, _ = _make_streamer(AgentConfig(tx_enabled=True))
        await s.on_message(_tx_start(app_id="a", tx_gain=-20, tx_center_frequency=2.45e9))
        await asyncio.sleep(0.01)
        await s.on_message(_tx_start(app_id="b", tx_gain=-20, tx_center_frequency=2.45e9))
        # Let the second request process, then stop cleanly.
        await asyncio.sleep(0.01)
        await s.on_message({"type": "tx_stop", "app_id": "a"})
        return ws

    ws = asyncio.run(scenario())
    errors = [
        m for m in ws.json_sent
        if m.get("type") == "tx_status" and m.get("state") == "error"
    ]
    assert any("already active" in e.get("message", "") for e in errors)


def test_rejects_invalid_underrun_policy():
    async def scenario():
        s, ws, _ = _make_streamer(AgentConfig(tx_enabled=True))
        await s.on_message(
            {
                "type": "tx_start",
                "app_id": "a",
                "radio_config": {
                    "device": "mock",
                    "buffer_size": 8,
                    "tx_gain": -20,
                    "tx_center_frequency": 2.45e9,
                    "underrun_policy": "teleport",
                },
            }
        )
        return ws

    ws = asyncio.run(scenario())
    status = _last_tx_status(ws)
    assert status and status["state"] == "error"
    assert "underrun_policy" in status["message"]