ria-toolkit-oss/tests/remote_control/test_remote_transmitter.py

"""Tests for the server-side RemoteTransmitter ZMQ RPC dispatcher.

No real SDR hardware or ZMQ sockets are needed — we test run_function()
directly and mock the SDR drivers.
"""

from __future__ import annotations

from unittest.mock import MagicMock, patch

import pytest

from ria_toolkit_oss.remote_control.remote_transmitter import RemoteTransmitter

# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _make_mock_sdr():
    sdr = MagicMock()
    sdr.init_tx = MagicMock()
    sdr.tx_cw = MagicMock()
    sdr.close = MagicMock()
    return sdr


# ---------------------------------------------------------------------------
# set_radio dispatch
# ---------------------------------------------------------------------------


class TestSetRadio:
    def _pluto_module(self, mock_sdr):
        mod = MagicMock()
        mod.Pluto = MagicMock(return_value=mock_sdr)
        return mod

    def test_pluto_alias(self):
        tx = RemoteTransmitter()
        mock_sdr = _make_mock_sdr()
        with patch.dict("sys.modules", {"ria_toolkit_oss.sdr.pluto": self._pluto_module(mock_sdr)}):
            tx.set_radio("pluto", "ip:192.168.2.1")
        assert tx._sdr is mock_sdr

    def test_plutosdr_alias(self):
        tx = RemoteTransmitter()
        mock_sdr = _make_mock_sdr()
        with patch.dict("sys.modules", {"ria_toolkit_oss.sdr.pluto": self._pluto_module(mock_sdr)}):
            tx.set_radio("PlutoSDR", "ip:192.168.2.1")
        assert tx._sdr is mock_sdr

    def test_usrp_alias(self):
        tx = RemoteTransmitter()
        mock_sdr = _make_mock_sdr()
        mock_module = MagicMock()
        mock_module.USRP = MagicMock(return_value=mock_sdr)
        with patch.dict("sys.modules", {"ria_toolkit_oss.sdr.usrp": mock_module}):
            tx.set_radio("usrp", "usrp://addr=192.168.10.2")
        assert tx._sdr is mock_sdr

    def test_hackrf_alias(self):
        tx = RemoteTransmitter()
        mock_sdr = _make_mock_sdr()
        mock_module = MagicMock()
        mock_module.HackRF = MagicMock(return_value=mock_sdr)
        with patch.dict("sys.modules", {"ria_toolkit_oss.sdr.hackrf": mock_module}):
            tx.set_radio("hackrf", "")
        assert tx._sdr is mock_sdr

    def test_hackrf_one_alias(self):
        tx = RemoteTransmitter()
        mock_sdr = _make_mock_sdr()
        mock_module = MagicMock()
        mock_module.HackRF = MagicMock(return_value=mock_sdr)
        with patch.dict("sys.modules", {"ria_toolkit_oss.sdr.hackrf": mock_module}):
            tx.set_radio("hackrf_one", "")
        assert tx._sdr is mock_sdr

    def test_bladerf_alias(self):
        tx = RemoteTransmitter()
        mock_sdr = _make_mock_sdr()
        mock_module = MagicMock()
        mock_module.Blade = MagicMock(return_value=mock_sdr)
        with patch.dict("sys.modules", {"ria_toolkit_oss.sdr.blade": mock_module}):
            tx.set_radio("blade", "")
        assert tx._sdr is mock_sdr

    def test_bladerf_string_alias(self):
        """'bladerf' string (not 'blade') must also resolve to blade.Blade."""
        tx = RemoteTransmitter()
        mock_sdr = _make_mock_sdr()
        mock_module = MagicMock()
        mock_module.Blade = MagicMock(return_value=mock_sdr)
        with patch.dict("sys.modules", {"ria_toolkit_oss.sdr.blade": mock_module}):
            tx.set_radio("bladerf", "")
        assert tx._sdr is mock_sdr

    def test_case_insensitive(self):
        tx = RemoteTransmitter()
        mock_sdr = _make_mock_sdr()
        with patch.dict("sys.modules", {"ria_toolkit_oss.sdr.pluto": self._pluto_module(mock_sdr)}):
            tx.set_radio("PLUTO", "ip:192.168.2.1")
        assert tx._sdr is mock_sdr

    def test_unknown_radio_raises(self):
        tx = RemoteTransmitter()
        with pytest.raises(ValueError, match="Unknown SDR type"):
            tx.set_radio("nonexistent_radio")

    def test_import_error_raises_runtime(self):
        """ImportError during SDR driver load is re-raised as RuntimeError."""
        tx = RemoteTransmitter()
        # Inject a fake module whose Pluto class raises ImportError on import
        bad_module = MagicMock()
        bad_module.Pluto = MagicMock(side_effect=ImportError("pyadi-iio not installed"))
        with patch.dict("sys.modules", {"ria_toolkit_oss.sdr.pluto": bad_module}):
            with pytest.raises((RuntimeError, ImportError)):
                tx.set_radio("pluto")


# ---------------------------------------------------------------------------
# init_tx / transmit / stop guard
# ---------------------------------------------------------------------------


class TestInitTxGuards:
    def test_init_tx_without_set_radio_raises(self):
        tx = RemoteTransmitter()
        with pytest.raises(RuntimeError, match="set_radio"):
            tx.init_tx(center_frequency=2.4e9, sample_rate=20e6, gain=0)

    def test_transmit_without_set_radio_raises(self):
        tx = RemoteTransmitter()
        with pytest.raises(RuntimeError):
            tx.transmit(duration_s=0.1)

    def test_stop_without_set_radio_is_safe(self):
        tx = RemoteTransmitter()
        tx.stop()  # should not raise — nothing to close


class TestInitTx:
    def _tx_with_mock_sdr(self):
        tx = RemoteTransmitter()
        tx._sdr = _make_mock_sdr()
        return tx

    def test_delegates_to_sdr(self):
        tx = self._tx_with_mock_sdr()
        tx.init_tx(center_frequency=2.4e9, sample_rate=20e6, gain=30, channel=1)
        tx._sdr.init_tx.assert_called_once_with(
            center_frequency=2.4e9,
            sample_rate=20e6,
            gain=30,
            channel=1,
        )

    def test_default_channel_zero(self):
        tx = self._tx_with_mock_sdr()
        tx.init_tx(center_frequency=2.4e9, sample_rate=20e6, gain=30)
        _, kwargs = tx._sdr.init_tx.call_args
        assert kwargs["channel"] == 0


class TestTransmit:
    def test_calls_tx_cw_until_duration(self):
        tx = RemoteTransmitter()
        tx._sdr = _make_mock_sdr()
        tx.init_tx(center_frequency=2.4e9, sample_rate=20e6, gain=0)
        tx.transmit(duration_s=0.05)
        assert tx._sdr.tx_cw.called

    def test_zero_duration_does_not_call_tx_cw(self):
        tx = RemoteTransmitter()
        tx._sdr = _make_mock_sdr()
        tx.init_tx(center_frequency=2.4e9, sample_rate=20e6, gain=0)
        tx.transmit(duration_s=0.0)
        tx._sdr.tx_cw.assert_not_called()

    def test_missing_tx_cw_method_handled(self):
        """AttributeError on tx_cw should not crash transmit()."""
        tx = RemoteTransmitter()
        sdr = MagicMock(spec=[])  # no tx_cw attribute
        sdr.init_tx = MagicMock()
        tx._sdr = sdr
        # Should not raise — AttributeError is caught and slept through
        tx.transmit(duration_s=0.01)


class TestStop:
    def test_calls_close_and_clears_sdr(self):
        tx = RemoteTransmitter()
        mock_sdr = _make_mock_sdr()
        tx._sdr = mock_sdr
        tx.stop()
        mock_sdr.close.assert_called_once()
        assert tx._sdr is None

    def test_close_exception_is_swallowed(self):
        tx = RemoteTransmitter()
        sdr = _make_mock_sdr()
        sdr.close.side_effect = RuntimeError("hardware error")
        tx._sdr = sdr
        tx.stop()  # should not raise
        assert tx._sdr is None

    def test_stop_idempotent(self):
        tx = RemoteTransmitter()
        tx.stop()
        tx.stop()  # second call is safe


# ---------------------------------------------------------------------------
# run_function dispatcher
# ---------------------------------------------------------------------------


class TestRunFunction:
    def _tx_with_mock_sdr(self):
        tx = RemoteTransmitter()
        tx._sdr = _make_mock_sdr()
        return tx

    def test_unknown_function_returns_failure(self):
        tx = RemoteTransmitter()
        resp = tx.run_function({"function_name": "explode"})
        assert resp["status"] is False
        assert "explode" in resp["error_message"]

    def test_set_radio_success(self):
        tx = RemoteTransmitter()
        mock_sdr = _make_mock_sdr()
        mod = MagicMock()
        mod.Pluto = MagicMock(return_value=mock_sdr)
        with patch.dict("sys.modules", {"ria_toolkit_oss.sdr.pluto": mod}):
            resp = tx.run_function({"function_name": "set_radio", "radio_str": "pluto", "identifier": "ip:1.2.3.4"})
        assert resp["status"] is True

    def test_set_radio_bad_type_returns_failure(self):
        tx = RemoteTransmitter()
        resp = tx.run_function({"function_name": "set_radio", "radio_str": "alien_device"})
        assert resp["status"] is False

    def test_init_tx_without_radio_returns_failure(self):
        tx = RemoteTransmitter()
        resp = tx.run_function(
            {
                "function_name": "init_tx",
                "center_frequency": 2.4e9,
                "sample_rate": 20e6,
                "gain": 0,
            }
        )
        assert resp["status"] is False
        assert resp["error_message"]

    def test_init_tx_with_radio_success(self):
        tx = self._tx_with_mock_sdr()
        resp = tx.run_function(
            {
                "function_name": "init_tx",
                "center_frequency": 2.4e9,
                "sample_rate": 20e6,
                "gain": 30,
            }
        )
        assert resp["status"] is True

    def test_transmit_runs_for_short_duration(self):
        tx = self._tx_with_mock_sdr()
        tx._sdr.init_tx = MagicMock()
        resp = tx.run_function(
            {
                "function_name": "init_tx",
                "center_frequency": 2.4e9,
                "sample_rate": 20e6,
                "gain": 0,
            }
        )
        resp = tx.run_function({"function_name": "transmit", "duration_s": 0.02})
        assert resp["status"] is True

    def test_stop_via_run_function(self):
        tx = self._tx_with_mock_sdr()
        resp = tx.run_function({"function_name": "stop"})
        assert resp["status"] is True
        assert tx._sdr is None

    def test_response_always_has_required_keys(self):
        tx = RemoteTransmitter()
        for fn in ("set_radio", "init_tx", "transmit", "stop", "bogus"):
            resp = tx.run_function({"function_name": fn})
            assert "status" in resp
            assert "message" in resp
            assert "error_message" in resp