ria-toolkit-oss/tests/ria_toolkit_oss_cli/test_capture.py

# flake8: noqa
"""Tests for capture command."""

import os
import tempfile
from unittest.mock import MagicMock, patch

import numpy as np
import pytest
import yaml
from click.testing import CliRunner

from ria_toolkit_oss.ria_toolkit_oss_cli.ria_toolkit_oss.capture import (
    auto_select_device,
    capture,
    get_sdr_device,
    save_visualization,
)


class TestGetSdrDevice:
    """Tests for get_sdr_device function."""

    def test_get_pluto_device(self):
        """Test getting PlutoSDR device."""
        mock_sdr_class = MagicMock()
        mock_sdr_instance = MagicMock()
        mock_sdr_class.return_value = mock_sdr_instance

        with patch.dict("sys.modules", {"src.ria_toolkit_oss.sdr.pluto": MagicMock(Pluto=mock_sdr_class)}):
            device = get_sdr_device("pluto")
            assert device is mock_sdr_instance

    def test_get_hackrf_device(self):
        """Test getting HackRF device."""
        mock_sdr_class = MagicMock()
        mock_sdr_instance = MagicMock()
        mock_sdr_class.return_value = mock_sdr_instance

        with patch.dict("sys.modules", {"src.ria_toolkit_oss.sdr.hackrf": MagicMock(HackRF=mock_sdr_class)}):
            device = get_sdr_device("hackrf")
            assert device is mock_sdr_instance

    def test_get_unknown_device(self):
        """Test getting unknown device type."""
        from click.exceptions import ClickException

        with pytest.raises(ClickException) as exc_info:
            get_sdr_device("unknown_device")

        assert "Unknown device type" in str(exc_info.value)


class TestAutoSelectDevice:
    """Tests for auto_select_device function."""

    def test_auto_select_no_devices(self):
        """Test auto-select with no devices found."""
        from click.exceptions import ClickException

        with patch("ria_toolkit_oss.ria_toolkit_oss_cli.ria_toolkit_oss.capture.list_all_devices") as mock_discover:
            mock_discover.return_value = []

            with pytest.raises(ClickException) as exc_info:
                auto_select_device()

            assert "No SDR devices found" in str(exc_info.value)

    def test_auto_select_single_device(self):
        """Test auto-select with single device."""
        with patch("ria_toolkit_oss.ria_toolkit_oss_cli.ria_toolkit_oss.capture.list_all_devices") as mock_discover:
            mock_discover.return_value = [{"type": "HackRF", "serial": "123456"}]

            device_type = auto_select_device(quiet=True)
            assert device_type == "hackrf"

    def test_auto_select_single_device_with_warning(self):
        """Test auto-select shows warning when not quiet."""
        with (
            patch("ria_toolkit_oss.ria_toolkit_oss_cli.ria_toolkit_oss.capture.list_all_devices") as mock_discover,
            patch("ria_toolkit_oss.ria_toolkit_oss_cli.ria_toolkit_oss.capture.click.echo") as mock_echo,
        ):

            mock_discover.return_value = [{"type": "PlutoSDR", "uri": "ip:pluto.local"}]

            device_type = auto_select_device(quiet=False)

            assert device_type == "pluto"
            # Should have called echo twice (warning + hint)
            assert mock_echo.call_count == 2

    def test_auto_select_multiple_devices(self):
        """Test auto-select with multiple devices raises error."""
        from click.exceptions import ClickException

        with patch("ria_toolkit_oss.ria_toolkit_oss_cli.ria_toolkit_oss.capture.list_all_devices") as mock_discover:
            mock_discover.return_value = [
                {"type": "HackRF", "serial": "123456"},
                {"type": "PlutoSDR", "uri": "ip:pluto.local"},
            ]

            with pytest.raises(ClickException) as exc_info:
                auto_select_device()

            assert "Multiple devices found" in str(exc_info.value)

    def test_auto_select_device_name_mapping(self):
        """Test device name mapping."""
        with patch("ria_toolkit_oss.ria_toolkit_oss_cli.ria_toolkit_oss.capture.list_all_devices") as mock_discover:
            # Test various device name formats
            test_cases = [
                ("PlutoSDR", "pluto"),
                ("HackRF", "hackrf"),
                ("BladeRF", "bladerf"),
                ("RTL-SDR", "rtlsdr"),
            ]

            for device_name, expected_type in test_cases:
                mock_discover.return_value = [{"type": device_name}]
                device_type = auto_select_device(quiet=True)
                assert device_type == expected_type


class TestSaveVisualization:
    """Tests for save_visualization function."""

    def test_save_visualization_success(self):
        """Test successful visualization save."""
        mock_recording = MagicMock()

        with patch("ria_toolkit_oss.ria_toolkit_oss_cli.ria_toolkit_oss.capture.view_simple_sig") as mock_view:
            save_visualization(mock_recording, "test.png", quiet=True)

            mock_view.assert_called_once_with(
                mock_recording, output_path="test.png", saveplot=True, fast_mode=False, labels_mode=True
            )

    def test_save_visualization_import_error(self):
        """Test visualization save with import error."""
        mock_recording = MagicMock()

        with (
            patch(
                "ria_toolkit_oss.ria_toolkit_oss_cli.ria_toolkit_oss.capture.view_simple_sig",
                side_effect=ImportError("Module not found"),
            ),
            patch("ria_toolkit_oss.ria_toolkit_oss_cli.ria_toolkit_oss.capture.click.echo") as mock_echo,
        ):

            save_visualization(mock_recording, "test.png", quiet=True)

            # Should catch error and echo warning
            mock_echo.assert_called_once()
            assert "Warning" in str(mock_echo.call_args)

    def test_save_visualization_general_error(self):
        """Test visualization save with general error."""
        mock_recording = MagicMock()

        with (
            patch(
                "ria_toolkit_oss.ria_toolkit_oss_cli.ria_toolkit_oss.capture.view_simple_sig",
                side_effect=Exception("Failed to plot"),
            ),
            patch("ria_toolkit_oss.ria_toolkit_oss_cli.ria_toolkit_oss.capture.click.echo") as mock_echo,
        ):

            save_visualization(mock_recording, "test.png", quiet=True)

            mock_echo.assert_called_once()
            assert "Failed to save visualization" in str(mock_echo.call_args)