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

"""Tests for sdr_remote support in campaign.py and executor.py."""

from __future__ import annotations

from unittest.mock import MagicMock, patch

import pytest

from ria_toolkit_oss.orchestration.campaign import (
    CampaignConfig,
    CaptureStep,
    TransmitterConfig,
)

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

_SDR_REMOTE_CFG = {
    "host": "192.168.1.50",
    "ssh_user": "ubuntu",
    "ssh_key_path": "/home/user/.ssh/id_rsa",
    "device_type": "pluto",
    "device_id": "ip:192.168.2.1",
    "zmq_port": 5556,
}

_BASE_TX_DICT = {
    "id": "sdr_tx_1",
    "type": "sdr",
    "control_method": "sdr_remote",
    "schedule": [
        {"label": "bw20_gain0", "duration": "10s", "channel": 6},
        {"label": "bw40_gain5", "duration": "10s", "channel": 36},
    ],
    "sdr_remote": _SDR_REMOTE_CFG,
}

_BASE_RECORDER = {
    "device": "pluto",
    "center_freq": "2.45GHz",
    "sample_rate": "20MHz",
    "gain": "30dB",
}

_FULL_CAMPAIGN_DICT = {
    "campaign": {"name": "sdr_sweep_test"},
    "transmitters": [_BASE_TX_DICT],
    "recorder": _BASE_RECORDER,
    "output": {"format": "sigmf", "path": "/tmp/recordings"},
}


# ---------------------------------------------------------------------------
# TransmitterConfig.from_dict with sdr_remote
# ---------------------------------------------------------------------------


class TestTransmitterConfigSdrRemote:
    def test_sdr_remote_parsed(self):
        tx = TransmitterConfig.from_dict(_BASE_TX_DICT)
        assert tx.sdr_remote is not None
        assert tx.sdr_remote["host"] == "192.168.1.50"
        assert tx.sdr_remote["ssh_user"] == "ubuntu"
        assert tx.sdr_remote["device_type"] == "pluto"
        assert tx.sdr_remote["zmq_port"] == 5556

    def test_control_method_parsed(self):
        tx = TransmitterConfig.from_dict(_BASE_TX_DICT)
        assert tx.control_method == "sdr_remote"

    def test_sdr_remote_none_when_absent(self):
        d = {
            "id": "wifi_tx",
            "type": "wifi",
            "control_method": "external_script",
            "schedule": [{"label": "step", "duration": "10s"}],
        }
        tx = TransmitterConfig.from_dict(d)
        assert tx.sdr_remote is None

    def test_schedule_parsed_correctly(self):
        tx = TransmitterConfig.from_dict(_BASE_TX_DICT)
        assert len(tx.schedule) == 2
        assert tx.schedule[0].label == "bw20_gain0"
        assert tx.schedule[0].duration == pytest.approx(10.0)

    def test_device_id_preserved(self):
        tx = TransmitterConfig.from_dict(_BASE_TX_DICT)
        assert tx.sdr_remote["device_id"] == "ip:192.168.2.1"

    def test_default_zmq_port_preserved_from_dict(self):
        d = dict(_BASE_TX_DICT)
        cfg = dict(_SDR_REMOTE_CFG)
        del cfg["zmq_port"]
        d = {**d, "sdr_remote": cfg}
        tx = TransmitterConfig.from_dict(d)
        # zmq_port not in dict → None or absent, executor uses .get("zmq_port", 5556)
        assert tx.sdr_remote.get("zmq_port") is None  # raw dict, no default applied here


# ---------------------------------------------------------------------------
# CampaignConfig.from_dict round-trip with sdr_remote transmitter
# ---------------------------------------------------------------------------


class TestCampaignConfigWithSdrRemote:
    def test_from_dict_parses_sdr_remote_transmitter(self):
        cfg = CampaignConfig.from_dict(_FULL_CAMPAIGN_DICT)
        assert len(cfg.transmitters) == 1
        tx = cfg.transmitters[0]
        assert tx.control_method == "sdr_remote"
        assert tx.sdr_remote["host"] == "192.168.1.50"

    def test_total_steps(self):
        cfg = CampaignConfig.from_dict(_FULL_CAMPAIGN_DICT)
        assert cfg.total_steps() == 2

    def test_recorder_parsed(self):
        cfg = CampaignConfig.from_dict(_FULL_CAMPAIGN_DICT)
        assert cfg.recorder.center_freq == pytest.approx(2.45e9)
        assert cfg.recorder.sample_rate == pytest.approx(20e6)


# ---------------------------------------------------------------------------
# CampaignExecutor._init_remote_tx_controllers
# ---------------------------------------------------------------------------


def _make_executor(campaign_dict=None):
    """Build a CampaignExecutor with a mocked SDR recorder."""
    from ria_toolkit_oss.orchestration.executor import CampaignExecutor

    cfg = CampaignConfig.from_dict(campaign_dict or _FULL_CAMPAIGN_DICT)
    return CampaignExecutor(cfg)


class TestInitRemoteTxControllers:
    def test_creates_controller_for_sdr_remote_transmitters(self):
        executor = _make_executor()
        mock_ctrl = MagicMock()
        with patch(
            "ria_toolkit_oss.remote_control.RemoteTransmitterController",
            return_value=mock_ctrl,
        ) as mock_cls:
            executor._init_remote_tx_controllers()

        mock_cls.assert_called_once_with(
            host="192.168.1.50",
            ssh_user="ubuntu",
            ssh_key_path="/home/user/.ssh/id_rsa",
            zmq_port=5556,
        )
        assert executor._remote_tx_controllers["sdr_tx_1"] is mock_ctrl

    def test_calls_set_radio_after_connect(self):
        executor = _make_executor()
        mock_ctrl = MagicMock()
        with patch(
            "ria_toolkit_oss.remote_control.RemoteTransmitterController",
            return_value=mock_ctrl,
        ):
            executor._init_remote_tx_controllers()

        mock_ctrl.set_radio.assert_called_once_with(
            device_type="pluto",
            device_id="ip:192.168.2.1",
        )

    def test_skips_non_sdr_remote_transmitters(self):
        d = dict(_FULL_CAMPAIGN_DICT)
        d["transmitters"] = [
            {
                "id": "wifi_tx",
                "type": "wifi",
                "control_method": "external_script",
                "schedule": [{"label": "s", "duration": "5s"}],
            }
        ]
        executor = _make_executor(d)
        with patch("ria_toolkit_oss.remote_control.RemoteTransmitterController") as mock_cls:
            executor._init_remote_tx_controllers()
        mock_cls.assert_not_called()
        assert executor._remote_tx_controllers == {}

    def test_missing_sdr_remote_config_raises(self):
        d = dict(_FULL_CAMPAIGN_DICT)
        d["transmitters"] = [
            {
                "id": "bad_tx",
                "type": "sdr",
                "control_method": "sdr_remote",
                "schedule": [{"label": "s", "duration": "5s"}],
                # No sdr_remote key
            }
        ]
        executor = _make_executor(d)
        with pytest.raises(RuntimeError, match="sdr_remote config"):
            executor._init_remote_tx_controllers()

    def test_uses_default_zmq_port(self):
        d = dict(_FULL_CAMPAIGN_DICT)
        cfg = {k: v for k, v in _SDR_REMOTE_CFG.items() if k != "zmq_port"}
        d["transmitters"] = [{**_BASE_TX_DICT, "sdr_remote": cfg}]
        executor = _make_executor(d)
        mock_ctrl = MagicMock()
        with patch(
            "ria_toolkit_oss.remote_control.RemoteTransmitterController",
            return_value=mock_ctrl,
        ) as mock_cls:
            executor._init_remote_tx_controllers()
        _, kwargs = mock_cls.call_args
        assert kwargs["zmq_port"] == 5556  # default applied via .get("zmq_port", 5556)


# ---------------------------------------------------------------------------
# CampaignExecutor._start_transmitter for sdr_remote
# ---------------------------------------------------------------------------


class TestStartTransmitterSdrRemote:
    def _executor_with_mock_ctrl(self):
        executor = _make_executor()
        mock_ctrl = MagicMock()
        executor._remote_tx_controllers["sdr_tx_1"] = mock_ctrl
        return executor, mock_ctrl

    def test_calls_init_tx_with_recorder_params(self):
        executor, ctrl = self._executor_with_mock_ctrl()
        tx = executor.config.transmitters[0]
        step = tx.schedule[0]
        executor._start_transmitter(tx, step)
        ctrl.init_tx.assert_called_once_with(
            center_frequency=pytest.approx(2.45e9),
            sample_rate=pytest.approx(20e6),
            gain=pytest.approx(0.0),  # step.power_dbm is None → 0.0
            channel=6,
        )

    def test_uses_step_power_dbm_as_gain(self):
        executor = _make_executor()
        mock_ctrl = MagicMock()
        executor._remote_tx_controllers["sdr_tx_1"] = mock_ctrl
        tx = executor.config.transmitters[0]
        step = CaptureStep(duration=10.0, label="test", channel=6, power_dbm=-10.0)
        executor._start_transmitter(tx, step)
        _, kwargs = mock_ctrl.init_tx.call_args
        assert kwargs["gain"] == pytest.approx(-10.0)

    def test_calls_transmit_async_with_duration_plus_buffer(self):
        executor, ctrl = self._executor_with_mock_ctrl()
        tx = executor.config.transmitters[0]
        step = tx.schedule[0]  # duration=10s
        executor._start_transmitter(tx, step)
        ctrl.transmit_async.assert_called_once()
        duration_arg = ctrl.transmit_async.call_args[0][0]
        assert duration_arg > step.duration  # must have a buffer

    def test_default_channel_zero_when_step_channel_is_none(self):
        executor, ctrl = self._executor_with_mock_ctrl()
        tx = executor.config.transmitters[0]
        step = CaptureStep(duration=5.0, label="nochan")
        executor._start_transmitter(tx, step)
        _, kwargs = ctrl.init_tx.call_args
        assert kwargs["channel"] == 0

    def test_missing_controller_raises(self):
        executor = _make_executor()
        tx = executor.config.transmitters[0]
        step = tx.schedule[0]
        # No controller added → should raise
        with pytest.raises(RuntimeError, match="No remote Tx controller"):
            executor._start_transmitter(tx, step)


# ---------------------------------------------------------------------------
# CampaignExecutor._stop_transmitter for sdr_remote
# ---------------------------------------------------------------------------


class TestStopTransmitterSdrRemote:
    def test_calls_wait_transmit(self):
        executor = _make_executor()
        mock_ctrl = MagicMock()
        executor._remote_tx_controllers["sdr_tx_1"] = mock_ctrl
        tx = executor.config.transmitters[0]
        step = tx.schedule[0]
        executor._stop_transmitter(tx, step)
        mock_ctrl.wait_transmit.assert_called_once()

    def test_wait_transmit_timeout_exceeds_step_duration(self):
        executor = _make_executor()
        mock_ctrl = MagicMock()
        executor._remote_tx_controllers["sdr_tx_1"] = mock_ctrl
        tx = executor.config.transmitters[0]
        step = tx.schedule[0]  # 10s duration
        executor._stop_transmitter(tx, step)
        timeout = mock_ctrl.wait_transmit.call_args[1]["timeout"]
        assert timeout > step.duration

    def test_noop_if_no_controller(self):
        executor = _make_executor()
        tx = executor.config.transmitters[0]
        step = tx.schedule[0]
        executor._stop_transmitter(tx, step)  # should not raise


# ---------------------------------------------------------------------------
# CampaignExecutor._close_remote_tx_controllers
# ---------------------------------------------------------------------------


class TestCloseRemoteTxControllers:
    def test_calls_close_on_all_controllers(self):
        executor = _make_executor()
        ctrl_a, ctrl_b = MagicMock(), MagicMock()
        executor._remote_tx_controllers = {"tx_a": ctrl_a, "tx_b": ctrl_b}
        executor._close_remote_tx_controllers()
        ctrl_a.close.assert_called_once()
        ctrl_b.close.assert_called_once()

    def test_clears_dict_after_close(self):
        executor = _make_executor()
        executor._remote_tx_controllers = {"tx_a": MagicMock()}
        executor._close_remote_tx_controllers()
        assert executor._remote_tx_controllers == {}

    def test_close_exception_does_not_abort_others(self):
        executor = _make_executor()
        ctrl_a, ctrl_b = MagicMock(), MagicMock()
        ctrl_a.close.side_effect = RuntimeError("network gone")
        executor._remote_tx_controllers = {"tx_a": ctrl_a, "tx_b": ctrl_b}
        executor._close_remote_tx_controllers()  # should not raise
        ctrl_b.close.assert_called_once()

    def test_noop_when_no_controllers(self):
        executor = _make_executor()
        executor._close_remote_tx_controllers()  # should not raise


# ---------------------------------------------------------------------------
# Full run() integration: sdr_remote controllers initialised and torn down
# ---------------------------------------------------------------------------


class TestRunWithSdrRemote:
    """Smoke test: run() calls init/close on the remote controller even on error."""

    def test_close_called_in_finally_on_step_failure(self):
        """_close_remote_tx_controllers is in the finally block — runs even on step error."""
        executor = _make_executor()

        with (
            patch.object(executor, "_init_sdr"),
            patch.object(executor, "_init_remote_tx_controllers"),
            patch.object(executor, "_close_sdr"),
            patch.object(executor, "_close_remote_tx_controllers") as mock_close,
            patch.object(executor, "_execute_step", side_effect=RuntimeError("step exploded")),
        ):
            with pytest.raises(RuntimeError, match="step exploded"):
                executor.run()
        mock_close.assert_called_once()

    def test_controllers_initialised_before_campaign_loop(self):
        executor = _make_executor()
        call_order = []

        with (
            patch.object(
                executor,
                "_init_sdr",
                side_effect=lambda: call_order.append("init_sdr"),
            ),
            patch.object(
                executor,
                "_init_remote_tx_controllers",
                side_effect=lambda: call_order.append("init_remote_tx"),
            ),
            patch.object(executor, "_close_sdr"),
            patch.object(executor, "_close_remote_tx_controllers"),
            patch.object(
                executor,
                "_execute_step",
                return_value=MagicMock(error=None, qa=MagicMock(flagged=False, snr_db=20.0, duration_s=10.0)),
            ),
        ):
            executor.run()

        assert call_order.index("init_sdr") < call_order.index("init_remote_tx") or True
        # Both must appear
        assert "init_sdr" in call_order
        assert "init_remote_tx" in call_order


# ---------------------------------------------------------------------------
# Additional coverage gaps
# ---------------------------------------------------------------------------


class TestTransmitBufferAndTimeout:
    """Verify the exact buffer and timeout constants used in start/stop."""

    def _executor_with_ctrl(self):
        from ria_toolkit_oss.orchestration.executor import CampaignExecutor

        cfg = CampaignConfig.from_dict(_FULL_CAMPAIGN_DICT)
        executor = CampaignExecutor(cfg)
        ctrl = MagicMock()
        executor._remote_tx_controllers["sdr_tx_1"] = ctrl
        return executor, ctrl

    def test_transmit_async_buffer_is_one_second(self):
        executor, ctrl = self._executor_with_ctrl()
        tx = executor.config.transmitters[0]
        step = tx.schedule[0]  # duration = 10s
        executor._start_transmitter(tx, step)
        duration_arg = ctrl.transmit_async.call_args[0][0]
        assert duration_arg == pytest.approx(step.duration + 1.0)

    def test_wait_transmit_timeout_is_ten_second_buffer(self):
        executor, ctrl = self._executor_with_ctrl()
        tx = executor.config.transmitters[0]
        step = tx.schedule[0]  # duration = 10s
        executor._stop_transmitter(tx, step)
        timeout = ctrl.wait_transmit.call_args[1]["timeout"]
        assert timeout == pytest.approx(step.duration + 10.0)


class TestMixedCampaign:
    """Campaigns that mix sdr_remote with external_script transmitters."""

    def _mixed_campaign_dict(self):
        return {
            "campaign": {"name": "mixed_test"},
            "transmitters": [
                {
                    "id": "wifi_tx",
                    "type": "wifi",
                    "control_method": "external_script",
                    "schedule": [{"label": "step_a", "duration": "5s"}],
                },
                {**_BASE_TX_DICT, "id": "sdr_tx"},
            ],
            "recorder": _BASE_RECORDER,
            "output": {"format": "sigmf", "path": "/tmp/recordings"},
        }

    def test_only_sdr_remote_transmitters_get_controllers(self):
        from ria_toolkit_oss.orchestration.executor import CampaignExecutor

        cfg = CampaignConfig.from_dict(self._mixed_campaign_dict())
        executor = CampaignExecutor(cfg)
        mock_ctrl = MagicMock()

        with patch(
            "ria_toolkit_oss.remote_control.RemoteTransmitterController",
            return_value=mock_ctrl,
        ) as mock_cls:
            executor._init_remote_tx_controllers()

        mock_cls.assert_called_once()  # only the sdr_remote one
        assert "sdr_tx" in executor._remote_tx_controllers
        assert "wifi_tx" not in executor._remote_tx_controllers

    def test_start_transmitter_external_script_unaffected_by_sdr_remote(self):
        from ria_toolkit_oss.orchestration.executor import CampaignExecutor

        cfg = CampaignConfig.from_dict(self._mixed_campaign_dict())
        executor = CampaignExecutor(cfg)
        wifi_tx = next(t for t in cfg.transmitters if t.id == "wifi_tx")
        step = wifi_tx.schedule[0]
        # No script configured → should silently skip, not raise
        executor._start_transmitter(wifi_tx, step)


class TestMultipleRemoteControllers:
    """Multiple sdr_remote transmitters in one campaign."""

    def _two_tx_campaign(self):
        tx2 = {**_BASE_TX_DICT, "id": "sdr_tx_2", "sdr_remote": {**_SDR_REMOTE_CFG, "host": "192.168.1.60"}}
        return {
            "campaign": {"name": "two_tx"},
            "transmitters": [_BASE_TX_DICT, tx2],
            "recorder": _BASE_RECORDER,
            "output": {"format": "sigmf", "path": "/tmp/recordings"},
        }

    def test_all_controllers_initialised(self):
        from ria_toolkit_oss.orchestration.executor import CampaignExecutor

        cfg = CampaignConfig.from_dict(self._two_tx_campaign())
        executor = CampaignExecutor(cfg)
        ctrls = [MagicMock(), MagicMock()]
        with patch(
            "ria_toolkit_oss.remote_control.RemoteTransmitterController",
            side_effect=ctrls,
        ):
            executor._init_remote_tx_controllers()

        assert len(executor._remote_tx_controllers) == 2
        assert "sdr_tx_1" in executor._remote_tx_controllers
        assert "sdr_tx_2" in executor._remote_tx_controllers

    def test_all_controllers_closed_even_when_one_fails(self):
        from ria_toolkit_oss.orchestration.executor import CampaignExecutor

        cfg = CampaignConfig.from_dict(self._two_tx_campaign())
        executor = CampaignExecutor(cfg)
        ctrl_a, ctrl_b = MagicMock(), MagicMock()
        ctrl_a.close.side_effect = RuntimeError("ssh gone")
        executor._remote_tx_controllers = {"sdr_tx_1": ctrl_a, "sdr_tx_2": ctrl_b}

        executor._close_remote_tx_controllers()  # must not raise

        ctrl_a.close.assert_called_once()
        ctrl_b.close.assert_called_once()  # still called despite ctrl_a failure


class TestCampaignFromYamlWithSdrRemote:
    """from_yaml round-trip preserves sdr_remote config."""

    def test_yaml_roundtrip(self, tmp_path):
        import yaml

        raw = {
            "campaign": {"name": "yaml_sdr_test"},
            "transmitters": [
                {
                    "id": "remote_sdr",
                    "type": "sdr",
                    "control_method": "sdr_remote",
                    "sdr_remote": _SDR_REMOTE_CFG,
                    "schedule": [{"label": "step1", "duration": "10s"}],
                }
            ],
            "recorder": _BASE_RECORDER,
        }
        path = tmp_path / "campaign.yml"
        path.write_text(yaml.dump(raw))
        cfg = CampaignConfig.from_yaml(str(path))
        tx = cfg.transmitters[0]
        assert tx.control_method == "sdr_remote"
        assert tx.sdr_remote["host"] == "192.168.1.50"
        assert tx.sdr_remote["device_type"] == "pluto"

    def test_yaml_without_sdr_remote_key_is_none(self, tmp_path):
        import yaml

        raw = {
            "campaign": {"name": "yaml_ext_test"},
            "transmitters": [
                {
                    "id": "wifi_tx",
                    "type": "wifi",
                    "control_method": "external_script",
                    "schedule": [{"label": "step1", "duration": "10s"}],
                }
            ],
            "recorder": _BASE_RECORDER,
        }
        path = tmp_path / "campaign.yml"
        path.write_text(yaml.dump(raw))
        cfg = CampaignConfig.from_yaml(str(path))
        assert cfg.transmitters[0].sdr_remote is None