diff --git a/tests/search_tests/test_tavily_search.py b/tests/search_tests/test_tavily_search.py index c7e923d06b..92c99fc61e 100644 --- a/tests/search_tests/test_tavily_search.py +++ b/tests/search_tests/test_tavily_search.py @@ -4,22 +4,87 @@ Tests for Tavily Search API integration. import os import sys import pytest +from unittest.mock import AsyncMock, patch, MagicMock sys.path.insert( 0, os.path.abspath("../..") ) -from tests.search_tests.base_search_unit_tests import BaseSearchTest +import litellm -class TestTavilySearch(BaseSearchTest): +class TestTavilySearch: """ - Tests for Tavily Search functionality. + Tests for Tavily Search functionality with mocked network responses. """ - def get_search_provider(self) -> str: + @pytest.mark.asyncio + async def test_tavily_search_request_payload(self): """ - Return search_provider for Tavily Search. + Test that validates the Tavily search request payload structure without making real API calls. """ - return "tavily" + # Set environment variable for API key + os.environ["TAVILY_API_KEY"] = "test-api-key" + + # Create a mock response + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.json.return_value = { + "results": [ + { + "title": "Test Result 1", + "url": "https://example.com/1", + "content": "This is a test snippet for result 1" + }, + { + "title": "Test Result 2", + "url": "https://example.com/2", + "content": "This is a test snippet for result 2" + } + ] + } + + # Mock the httpx AsyncClient post method + with patch("litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", new_callable=AsyncMock) as mock_post: + mock_post.return_value = mock_response + + # Make the search call + response = await litellm.asearch( + query="latest developments in AI", + search_provider="tavily", + max_results=5 + ) + + # Verify the post method was called once + assert mock_post.call_count == 1 + + # Get the actual call arguments + call_args = mock_post.call_args + + # Verify URL + assert call_args.kwargs["url"] == "https://api.tavily.com/search" + + # Verify headers contain Authorization + headers = call_args.kwargs.get("headers", {}) + assert "Authorization" in headers + assert headers["Authorization"] == "Bearer test-api-key" + assert headers["Content-Type"] == "application/json" + + # Verify request payload + json_data = call_args.kwargs.get("json") + assert json_data is not None + assert json_data["query"] == "latest developments in AI" + assert json_data["max_results"] == 5 + + # Verify response structure + assert hasattr(response, "results") + assert hasattr(response, "object") + assert response.object == "search" + assert len(response.results) == 2 + + # Verify first result + first_result = response.results[0] + assert first_result.title == "Test Result 1" + assert first_result.url == "https://example.com/1" + assert first_result.snippet == "This is a test snippet for result 1"