From 0e03bc8b8b4208a57a2fc57a4dbe1fbc5e1ec838 Mon Sep 17 00:00:00 2001 From: khushishelat Date: Mon, 31 Aug 2026 16:13:52 -0700 Subject: [PATCH] test(parallel_ai): fake search at HTTP boundary --- .../parallel_ai/test_parallel_ai_search.py | 181 ++++++++---------- 1 file changed, 85 insertions(+), 96 deletions(-) diff --git a/tests/test_litellm/llms/parallel_ai/test_parallel_ai_search.py b/tests/test_litellm/llms/parallel_ai/test_parallel_ai_search.py index 88f542f64c4..e7ff6a5f5db 100644 --- a/tests/test_litellm/llms/parallel_ai/test_parallel_ai_search.py +++ b/tests/test_litellm/llms/parallel_ai/test_parallel_ai_search.py @@ -2,6 +2,7 @@ Tests for Parallel AI Search API integration (v1 endpoint). """ +import json from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -37,6 +38,17 @@ def _mock_response(payload=None): return mock_response +@pytest.fixture +def httpx_transport(monkeypatch): + monkeypatch.setattr( # test-quality-ok: respx needs HTTPX enabled to fake the provider HTTP boundary. + litellm, + "disable_aiohttp_transport", + True, + ) + litellm.in_memory_llm_clients_cache.flush_cache() + yield + litellm.in_memory_llm_clients_cache.flush_cache() + @pytest.fixture def bundled_cost_map(monkeypatch): @@ -152,9 +164,7 @@ class TestParallelAISearch: json_data = mock_post.call_args.kwargs.get("json") assert json_data["mode"] == "basic" - @pytest.mark.parametrize( - "processor,expected_mode", [("base", "basic"), ("pro", "advanced")] - ) + @pytest.mark.parametrize("processor,expected_mode", [("base", "basic"), ("pro", "advanced")]) @pytest.mark.asyncio async def test_legacy_processor_maps_to_mode(self, processor, expected_mode): with patch( @@ -239,9 +249,7 @@ class TestParallelAISearch: "arxiv.org", "nature.com", ] - assert advanced_settings["source_policy"]["exclude_domains"] == [ - "reddit.com" - ] + assert advanced_settings["source_policy"]["exclude_domains"] == ["reddit.com"] assert advanced_settings["excerpt_settings"]["max_chars_per_result"] == 1500 assert "max_results" not in json_data @@ -323,10 +331,7 @@ class TestParallelAISearch: ) call_args = mock_post.call_args - assert ( - call_args.kwargs["url"] - == "https://proxy.internal.example.com/v1/search" - ) + assert call_args.kwargs["url"] == "https://proxy.internal.example.com/v1/search" @pytest.mark.asyncio async def test_caller_api_base_without_key_is_refused(self, monkeypatch): @@ -357,62 +362,54 @@ class TestParallelAISearch: ) @pytest.mark.asyncio - async def test_flat_source_and_fetch_params_nest_under_advanced_settings(self): - with patch( - "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", - new_callable=AsyncMock, - ) as mock_post: - mock_post.return_value = _mock_response() + async def test_flat_source_and_fetch_params_nest_under_advanced_settings(self, respx_mock, httpx_transport): + route = respx_mock.post("https://api.parallel.ai/v1/search").respond(json=MOCK_V1_RESPONSE) - await litellm.asearch( - query="AI developments", - search_provider="parallel_ai", - objective="find peer-reviewed AI research", - include_domains=["arxiv.org"], - after_date="2026-01-01", - location="gb", - fetch_policy={"max_age_seconds": 600, "disable_cache_fallback": True}, - client_model="claude-fable-5", - ) + await litellm.asearch( + query="AI developments", + search_provider="parallel_ai", + objective="find peer-reviewed AI research", + include_domains=["arxiv.org"], + after_date="2026-01-01", + location="gb", + fetch_policy={"max_age_seconds": 600, "disable_cache_fallback": True}, + client_model="claude-fable-5", + ) - json_data = mock_post.call_args.kwargs.get("json") - assert json_data["objective"] == "find peer-reviewed AI research" - assert json_data["client_model"] == "claude-fable-5" + json_data = json.loads(route.calls[0].request.content) + assert json_data["objective"] == "find peer-reviewed AI research" + assert json_data["client_model"] == "claude-fable-5" - advanced_settings = json_data["advanced_settings"] - assert advanced_settings["location"] == "gb" - assert advanced_settings["fetch_policy"] == { - "max_age_seconds": 600, - "disable_cache_fallback": True, - } - assert advanced_settings["source_policy"]["include_domains"] == ["arxiv.org"] - assert advanced_settings["source_policy"]["after_date"] == "2026-01-01" + advanced_settings = json_data["advanced_settings"] + assert advanced_settings["location"] == "gb" + assert advanced_settings["fetch_policy"] == { + "max_age_seconds": 600, + "disable_cache_fallback": True, + } + assert advanced_settings["source_policy"]["include_domains"] == ["arxiv.org"] + assert advanced_settings["source_policy"]["after_date"] == "2026-01-01" - assert "include_domains" not in json_data - assert "after_date" not in json_data - assert "location" not in json_data - assert "fetch_policy" not in json_data + assert "include_domains" not in json_data + assert "after_date" not in json_data + assert "location" not in json_data + assert "fetch_policy" not in json_data @pytest.mark.asyncio - async def test_response_preserves_raw_parallel_fields(self): - with patch( - "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", - new_callable=AsyncMock, - ) as mock_post: - mock_post.return_value = _mock_response() + async def test_response_preserves_raw_parallel_fields(self, respx_mock, httpx_transport): + respx_mock.post("https://api.parallel.ai/v1/search").respond(json=MOCK_V1_RESPONSE) - response = await litellm.asearch( - query="AI developments", - search_provider="parallel_ai", - ) + response = await litellm.asearch( + query="AI developments", + search_provider="parallel_ai", + ) - dumped = response.model_dump() - assert dumped["search_id"] == "search_abc123" - assert dumped["session_id"] == "session_xyz" - assert dumped["parallel_usage"] == [{"name": "search_advanced", "count": 1}] + dumped = response.model_dump() + assert dumped["search_id"] == "search_abc123" + assert dumped["session_id"] == "session_xyz" + assert dumped["parallel_usage"] == [{"name": "search_advanced", "count": 1}] - first = response.results[0].model_dump() - assert first["excerpts"] == ["First excerpt.", "Second excerpt."] + first = response.results[0].model_dump() + assert first["excerpts"] == ["First excerpt.", "Second excerpt."] @pytest.mark.parametrize( "mode,usage,max_results,expected_cost", @@ -434,63 +431,55 @@ class TestParallelAISearch: ], ) @pytest.mark.asyncio - async def test_search_cost_uses_mode_and_provider_usage(self, mode, usage, max_results, expected_cost, bundled_cost_map): + async def test_search_cost_uses_mode_and_provider_usage( + self, mode, usage, max_results, expected_cost, bundled_cost_map, respx_mock, httpx_transport + ): response_payload = {**MOCK_V1_RESPONSE, "usage": usage} - with patch( - "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", - new_callable=AsyncMock, - ) as mock_post: - mock_post.return_value = _mock_response(response_payload) + respx_mock.post("https://api.parallel.ai/v1/search").respond(json=response_payload) - response = await litellm.asearch( - query="AI developments", - search_provider="parallel_ai", - mode=mode, - max_results=max_results, - ) + response = await litellm.asearch( + query="AI developments", + search_provider="parallel_ai", + mode=mode, + max_results=max_results, + ) - assert response._hidden_params["response_cost"] == pytest.approx(expected_cost) + assert response._hidden_params["response_cost"] == pytest.approx(expected_cost) @pytest.mark.asyncio - async def test_search_cost_treats_keyword_queries_as_one_request(self, bundled_cost_map): + async def test_search_cost_treats_keyword_queries_as_one_request( + self, bundled_cost_map, respx_mock, httpx_transport + ): response_payload = { **MOCK_V1_RESPONSE, "usage": [{"name": "sku_search", "count": 1}], } - with patch( - "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", - new_callable=AsyncMock, - ) as mock_post: - mock_post.return_value = _mock_response(response_payload) + respx_mock.post("https://api.parallel.ai/v1/search").respond(json=response_payload) - response = await litellm.asearch( - query=["AI developments", "machine learning trends"], - search_provider="parallel_ai", - mode="basic", - ) + response = await litellm.asearch( + query=["AI developments", "machine learning trends"], + search_provider="parallel_ai", + mode="basic", + ) - assert response._hidden_params["response_cost"] == pytest.approx(0.005) + assert response._hidden_params["response_cost"] == pytest.approx(0.005) @pytest.mark.asyncio - async def test_caller_cannot_supply_provider_usage(self, bundled_cost_map): + async def test_caller_cannot_supply_provider_usage(self, bundled_cost_map, respx_mock, httpx_transport): """`_parallel_ai_usage` prices the request, so a caller must not be able to set it. The provider reports no usage here, which is the case where a caller-supplied value would otherwise survive into the cost calculation. """ response_payload = {k: v for k, v in MOCK_V1_RESPONSE.items() if k != "usage"} - with patch( - "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", - new_callable=AsyncMock, - ) as mock_post: - mock_post.return_value = _mock_response(response_payload) + route = respx_mock.post("https://api.parallel.ai/v1/search").respond(json=response_payload) - response = await litellm.asearch( - query="AI developments", - search_provider="parallel_ai", - mode="basic", - _parallel_ai_usage=[{"name": "sku_search", "count": 0}], - ) + response = await litellm.asearch( + query="AI developments", + search_provider="parallel_ai", + mode="basic", + _parallel_ai_usage=[{"name": "sku_search", "count": 0}], + ) - assert response._hidden_params["response_cost"] == pytest.approx(0.005) - assert "_parallel_ai_usage" not in mock_post.call_args.kwargs["json"] + assert response._hidden_params["response_cost"] == pytest.approx(0.005) + assert "_parallel_ai_usage" not in json.loads(route.calls[0].request.content)