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")
|