"""
Tests for Tavily Search API integration.
"""

import os
import pytest
from unittest.mock import AsyncMock, patch, MagicMock


import litellm


class TestTavilySearch:
    """
    Tests for Tavily Search functionality with mocked network responses.
    """

    @pytest.mark.asyncio
    async def test_tavily_search_request_payload(self):
        """
        Test that validates the Tavily search request payload structure without making real API calls.
        """
        # 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"