File size: 8,641 Bytes
1041734
 
e7b4937
1041734
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
"""
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")