test/test_web_search.py

"""
Tests for web search tool (Tavily and Exa)
Author: @mangubee
Date: 2026-01-02

Tests cover:
- Tavily search with mocked API
- Exa search with mocked API
- Retry logic simulation
- Fallback mechanism
- Error handling
"""

import pytest
from unittest.mock import Mock, patch, MagicMock
from src.tools.web_search import tavily_search, exa_search, search


# ============================================================================
# Test Fixtures
# ============================================================================

@pytest.fixture
def mock_tavily_response():
    """Mock Tavily API response"""
    return {
        "results": [
            {
                "title": "Test Result 1",
                "url": "https://example.com/1",
                "content": "This is test content 1"
            },
            {
                "title": "Test Result 2",
                "url": "https://example.com/2",
                "content": "This is test content 2"
            }
        ]
    }


@pytest.fixture
def mock_exa_response():
    """Mock Exa API response"""
    mock_result_1 = Mock()
    mock_result_1.title = "Exa Result 1"
    mock_result_1.url = "https://exa.com/1"
    mock_result_1.text = "This is exa content 1"

    mock_result_2 = Mock()
    mock_result_2.title = "Exa Result 2"
    mock_result_2.url = "https://exa.com/2"
    mock_result_2.text = "This is exa content 2"

    mock_response = Mock()
    mock_response.results = [mock_result_1, mock_result_2]
    return mock_response


@pytest.fixture
def mock_settings_tavily():
    """Mock Settings with Tavily API key"""
    with patch('src.tools.web_search.Settings') as mock:
        settings_instance = Mock()
        settings_instance.tavily_api_key = "test_tavily_key"
        settings_instance.exa_api_key = "test_exa_key"
        settings_instance.default_search_tool = "tavily"
        mock.return_value = settings_instance
        yield mock


@pytest.fixture
def mock_settings_exa():
    """Mock Settings with Exa as default"""
    with patch('src.tools.web_search.Settings') as mock:
        settings_instance = Mock()
        settings_instance.tavily_api_key = "test_tavily_key"
        settings_instance.exa_api_key = "test_exa_key"
        settings_instance.default_search_tool = "exa"
        mock.return_value = settings_instance
        yield mock


# ============================================================================
# Tavily Search Tests
# ============================================================================

def test_tavily_search_success(mock_settings_tavily, mock_tavily_response):
    """Test successful Tavily search"""
    with patch('tavily.TavilyClient') as mock_client_class:
        mock_client = Mock()
        mock_client.search.return_value = mock_tavily_response
        mock_client_class.return_value = mock_client

        result = tavily_search("test query", max_results=2)

        assert result["source"] == "tavily"
        assert result["query"] == "test query"
        assert result["count"] == 2
        assert len(result["results"]) == 2
        assert result["results"][0]["title"] == "Test Result 1"
        assert result["results"][0]["url"] == "https://example.com/1"
        assert result["results"][0]["snippet"] == "This is test content 1"


def test_tavily_search_missing_api_key():
    """Test Tavily search with missing API key"""
    with patch('src.tools.web_search.Settings') as mock_settings:
        settings_instance = Mock()
        settings_instance.tavily_api_key = None
        mock_settings.return_value = settings_instance

        with pytest.raises(ValueError, match="TAVILY_API_KEY not configured"):
            tavily_search("test query")


def test_tavily_search_connection_error(mock_settings_tavily):
    """Test Tavily search with connection error (triggers retry)"""
    with patch('tavily.TavilyClient') as mock_client_class:
        mock_client = Mock()
        mock_client.search.side_effect = ConnectionError("Network error")
        mock_client_class.return_value = mock_client

        with pytest.raises(ConnectionError):
            tavily_search("test query")

        # Verify retry happened (should be called MAX_RETRIES times)
        assert mock_client.search.call_count == 3


def test_tavily_search_empty_results(mock_settings_tavily):
    """Test Tavily search with empty results"""
    with patch('tavily.TavilyClient') as mock_client_class:
        mock_client = Mock()
        mock_client.search.return_value = {"results": []}
        mock_client_class.return_value = mock_client

        result = tavily_search("test query")

        assert result["count"] == 0
        assert result["results"] == []


# ============================================================================
# Exa Search Tests
# ============================================================================

def test_exa_search_success(mock_settings_exa, mock_exa_response):
    """Test successful Exa search"""
    with patch('exa_py.Exa') as mock_client_class:
        mock_client = Mock()
        mock_client.search.return_value = mock_exa_response
        mock_client_class.return_value = mock_client

        result = exa_search("test query", max_results=2)

        assert result["source"] == "exa"
        assert result["query"] == "test query"
        assert result["count"] == 2
        assert len(result["results"]) == 2
        assert result["results"][0]["title"] == "Exa Result 1"
        assert result["results"][0]["url"] == "https://exa.com/1"
        assert result["results"][0]["snippet"] == "This is exa content 1"


def test_exa_search_missing_api_key():
    """Test Exa search with missing API key"""
    with patch('src.tools.web_search.Settings') as mock_settings:
        settings_instance = Mock()
        settings_instance.exa_api_key = None
        mock_settings.return_value = settings_instance

        with pytest.raises(ValueError, match="EXA_API_KEY not configured"):
            exa_search("test query")


def test_exa_search_connection_error(mock_settings_exa):
    """Test Exa search with connection error (triggers retry)"""
    with patch('exa_py.Exa') as mock_client_class:
        mock_client = Mock()
        mock_client.search.side_effect = ConnectionError("Network error")
        mock_client_class.return_value = mock_client

        with pytest.raises(ConnectionError):
            exa_search("test query")

        # Verify retry happened
        assert mock_client.search.call_count == 3


# ============================================================================
# Unified Search with Fallback Tests
# ============================================================================

def test_search_tavily_success(mock_settings_tavily, mock_tavily_response):
    """Test unified search using Tavily successfully"""
    with patch('tavily.TavilyClient') as mock_client_class:
        mock_client = Mock()
        mock_client.search.return_value = mock_tavily_response
        mock_client_class.return_value = mock_client

        result = search("test query")

        assert result["source"] == "tavily"
        assert result["count"] == 2


def test_search_fallback_to_exa(mock_settings_tavily, mock_exa_response):
    """Test unified search falls back to Exa when Tavily fails"""
    with patch('tavily.TavilyClient') as mock_tavily_class:
        with patch('exa_py.Exa') as mock_exa_class:
            # Tavily fails
            mock_tavily_client = Mock()
            mock_tavily_client.search.side_effect = Exception("Tavily error")
            mock_tavily_class.return_value = mock_tavily_client

            # Exa succeeds
            mock_exa_client = Mock()
            mock_exa_client.search.return_value = mock_exa_response
            mock_exa_class.return_value = mock_exa_client

            result = search("test query")

            assert result["source"] == "exa"
            assert result["count"] == 2


def test_search_both_fail(mock_settings_tavily):
    """Test unified search when both Tavily and Exa fail"""
    with patch('tavily.TavilyClient') as mock_tavily_class:
        with patch('exa_py.Exa') as mock_exa_class:
            # Both fail
            mock_tavily_client = Mock()
            mock_tavily_client.search.side_effect = Exception("Tavily error")
            mock_tavily_class.return_value = mock_tavily_client

            mock_exa_client = Mock()
            mock_exa_client.search.side_effect = Exception("Exa error")
            mock_exa_class.return_value = mock_exa_client

            with pytest.raises(Exception, match="Search failed"):
                search("test query")