mirror of
https://gitee.com/wanwujie/deer-flow
synced 2026-04-18 03:54:46 +08:00
refine the research prompt (#460)
This commit is contained in:
@@ -1,9 +1,8 @@
|
||||
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
import json
|
||||
import pytest
|
||||
from unittest.mock import Mock, patch, AsyncMock
|
||||
from unittest.mock import Mock, AsyncMock
|
||||
from src.tools.tavily_search.tavily_search_results_with_images import (
|
||||
TavilySearchResultsWithImages,
|
||||
)
|
||||
@@ -77,10 +76,8 @@ class TestTavilySearchResultsWithImages:
|
||||
assert tool.max_results == 10
|
||||
assert tool.include_image_descriptions is True
|
||||
|
||||
@patch("builtins.print")
|
||||
def test_run_success(
|
||||
self,
|
||||
mock_print,
|
||||
search_tool,
|
||||
mock_api_wrapper,
|
||||
sample_raw_results,
|
||||
@@ -110,10 +107,8 @@ class TestTavilySearchResultsWithImages:
|
||||
mock_api_wrapper.clean_results_with_images.assert_called_once_with(
|
||||
sample_raw_results
|
||||
)
|
||||
mock_print.assert_called_once()
|
||||
|
||||
@patch("builtins.print")
|
||||
def test_run_exception(self, mock_print, search_tool, mock_api_wrapper):
|
||||
def test_run_exception(self, search_tool, mock_api_wrapper):
|
||||
"""Test synchronous run with exception."""
|
||||
mock_api_wrapper.raw_results.side_effect = Exception("API Error")
|
||||
|
||||
@@ -124,10 +119,8 @@ class TestTavilySearchResultsWithImages:
|
||||
mock_api_wrapper.clean_results_with_images.assert_not_called()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch("builtins.print")
|
||||
async def test_arun_success(
|
||||
self,
|
||||
mock_print,
|
||||
search_tool,
|
||||
mock_api_wrapper,
|
||||
sample_raw_results,
|
||||
@@ -157,11 +150,9 @@ class TestTavilySearchResultsWithImages:
|
||||
mock_api_wrapper.clean_results_with_images.assert_called_once_with(
|
||||
sample_raw_results
|
||||
)
|
||||
mock_print.assert_called_once()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch("builtins.print")
|
||||
async def test_arun_exception(self, mock_print, search_tool, mock_api_wrapper):
|
||||
async def test_arun_exception(self, search_tool, mock_api_wrapper):
|
||||
"""Test asynchronous run with exception."""
|
||||
mock_api_wrapper.raw_results_async = AsyncMock(
|
||||
side_effect=Exception("Async API Error")
|
||||
@@ -173,10 +164,8 @@ class TestTavilySearchResultsWithImages:
|
||||
assert raw == {}
|
||||
mock_api_wrapper.clean_results_with_images.assert_not_called()
|
||||
|
||||
@patch("builtins.print")
|
||||
def test_run_with_run_manager(
|
||||
self,
|
||||
mock_print,
|
||||
search_tool,
|
||||
mock_api_wrapper,
|
||||
sample_raw_results,
|
||||
@@ -193,10 +182,8 @@ class TestTavilySearchResultsWithImages:
|
||||
assert raw == sample_raw_results
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch("builtins.print")
|
||||
async def test_arun_with_run_manager(
|
||||
self,
|
||||
mock_print,
|
||||
search_tool,
|
||||
mock_api_wrapper,
|
||||
sample_raw_results,
|
||||
@@ -213,52 +200,3 @@ class TestTavilySearchResultsWithImages:
|
||||
|
||||
assert result == sample_cleaned_results
|
||||
assert raw == sample_raw_results
|
||||
|
||||
@patch("builtins.print")
|
||||
def test_print_output_format(
|
||||
self,
|
||||
mock_print,
|
||||
search_tool,
|
||||
mock_api_wrapper,
|
||||
sample_raw_results,
|
||||
sample_cleaned_results,
|
||||
):
|
||||
"""Test that print outputs correctly formatted JSON."""
|
||||
mock_api_wrapper.raw_results.return_value = sample_raw_results
|
||||
mock_api_wrapper.clean_results_with_images.return_value = sample_cleaned_results
|
||||
|
||||
search_tool._run("test query")
|
||||
|
||||
# Verify print was called with expected format
|
||||
call_args = mock_print.call_args[0]
|
||||
assert call_args[0] == "sync"
|
||||
assert isinstance(call_args[1], str) # Should be JSON string
|
||||
|
||||
# Verify it's valid JSON
|
||||
json_data = json.loads(call_args[1])
|
||||
assert json_data == sample_cleaned_results
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch("builtins.print")
|
||||
async def test_async_print_output_format(
|
||||
self,
|
||||
mock_print,
|
||||
search_tool,
|
||||
mock_api_wrapper,
|
||||
sample_raw_results,
|
||||
sample_cleaned_results,
|
||||
):
|
||||
"""Test that async print outputs correctly formatted JSON."""
|
||||
mock_api_wrapper.raw_results_async = AsyncMock(return_value=sample_raw_results)
|
||||
mock_api_wrapper.clean_results_with_images.return_value = sample_cleaned_results
|
||||
|
||||
await search_tool._arun("test query")
|
||||
|
||||
# Verify print was called with expected format
|
||||
call_args = mock_print.call_args[0]
|
||||
assert call_args[0] == "async"
|
||||
assert isinstance(call_args[1], str) # Should be JSON string
|
||||
|
||||
# Verify it's valid JSON
|
||||
json_data = json.loads(call_args[1])
|
||||
assert json_data == sample_cleaned_results
|
||||
|
||||
Reference in New Issue
Block a user