deepdata: core product code (backend, frontend, runtime, deepflow module)
This commit is contained in:
figmar
@@ -0,0 +1,133 @@
|
||||
import json
|
||||
from unittest.mock import Mock, patch
|
||||
|
||||
from src.tools.crawl import crawl_tool
|
||||
|
||||
|
||||
class TestCrawlTool:
|
||||
@patch("src.tools.crawl.Crawler")
|
||||
def test_crawl_tool_success(self, mock_crawler_class):
|
||||
# Arrange
|
||||
mock_crawler = Mock()
|
||||
mock_article = Mock()
|
||||
mock_article.to_markdown.return_value = (
|
||||
"# Test Article\nThis is test content." * 100
|
||||
)
|
||||
mock_crawler.crawl.return_value = mock_article
|
||||
mock_crawler_class.return_value = mock_crawler
|
||||
|
||||
url = "https://example.com"
|
||||
|
||||
# Act
|
||||
result = crawl_tool(url)
|
||||
|
||||
# Assert
|
||||
assert isinstance(result, str)
|
||||
result_dict = json.loads(result)
|
||||
assert result_dict["url"] == url
|
||||
assert "crawled_content" in result_dict
|
||||
assert len(result_dict["crawled_content"]) <= 1000
|
||||
mock_crawler_class.assert_called_once()
|
||||
mock_crawler.crawl.assert_called_once_with(url)
|
||||
mock_article.to_markdown.assert_called_once()
|
||||
|
||||
@patch("src.tools.crawl.Crawler")
|
||||
def test_crawl_tool_short_content(self, mock_crawler_class):
|
||||
# Arrange
|
||||
mock_crawler = Mock()
|
||||
mock_article = Mock()
|
||||
short_content = "Short content"
|
||||
mock_article.to_markdown.return_value = short_content
|
||||
mock_crawler.crawl.return_value = mock_article
|
||||
mock_crawler_class.return_value = mock_crawler
|
||||
|
||||
url = "https://example.com"
|
||||
|
||||
# Act
|
||||
result = crawl_tool(url)
|
||||
|
||||
# Assert
|
||||
result_dict = json.loads(result)
|
||||
assert result_dict["crawled_content"] == short_content
|
||||
|
||||
@patch("src.tools.crawl.Crawler")
|
||||
@patch("src.tools.crawl.logger")
|
||||
def test_crawl_tool_crawler_exception(self, mock_logger, mock_crawler_class):
|
||||
# Arrange
|
||||
mock_crawler = Mock()
|
||||
mock_crawler.crawl.side_effect = Exception("Network error")
|
||||
mock_crawler_class.return_value = mock_crawler
|
||||
|
||||
url = "https://example.com"
|
||||
|
||||
# Act
|
||||
result = crawl_tool(url)
|
||||
|
||||
# Assert
|
||||
assert isinstance(result, str)
|
||||
assert "Failed to crawl" in result
|
||||
assert "Network error" in result
|
||||
mock_logger.error.assert_called_once()
|
||||
|
||||
@patch("src.tools.crawl.Crawler")
|
||||
@patch("src.tools.crawl.logger")
|
||||
def test_crawl_tool_crawler_instantiation_exception(
|
||||
self, mock_logger, mock_crawler_class
|
||||
):
|
||||
# Arrange
|
||||
mock_crawler_class.side_effect = Exception("Crawler init error")
|
||||
|
||||
url = "https://example.com"
|
||||
|
||||
# Act
|
||||
result = crawl_tool(url)
|
||||
|
||||
# Assert
|
||||
assert isinstance(result, str)
|
||||
assert "Failed to crawl" in result
|
||||
assert "Crawler init error" in result
|
||||
mock_logger.error.assert_called_once()
|
||||
|
||||
@patch("src.tools.crawl.Crawler")
|
||||
@patch("src.tools.crawl.logger")
|
||||
def test_crawl_tool_markdown_conversion_exception(
|
||||
self, mock_logger, mock_crawler_class
|
||||
):
|
||||
# Arrange
|
||||
mock_crawler = Mock()
|
||||
mock_article = Mock()
|
||||
mock_article.to_markdown.side_effect = Exception("Markdown conversion error")
|
||||
mock_crawler.crawl.return_value = mock_article
|
||||
mock_crawler_class.return_value = mock_crawler
|
||||
|
||||
url = "https://example.com"
|
||||
|
||||
# Act
|
||||
result = crawl_tool(url)
|
||||
|
||||
# Assert
|
||||
assert isinstance(result, str)
|
||||
assert "Failed to crawl" in result
|
||||
assert "Markdown conversion error" in result
|
||||
mock_logger.error.assert_called_once()
|
||||
|
||||
@patch("src.tools.crawl.Crawler")
|
||||
def test_crawl_tool_with_none_content(self, mock_crawler_class):
|
||||
# Arrange
|
||||
mock_crawler = Mock()
|
||||
mock_article = Mock()
|
||||
mock_article.to_markdown.return_value = "# Article\n\n*No content available*\n"
|
||||
mock_crawler.crawl.return_value = mock_article
|
||||
mock_crawler_class.return_value = mock_crawler
|
||||
|
||||
url = "https://example.com"
|
||||
|
||||
# Act
|
||||
result = crawl_tool(url)
|
||||
|
||||
# Assert
|
||||
assert isinstance(result, str)
|
||||
result_dict = json.loads(result)
|
||||
assert result_dict["url"] == url
|
||||
assert "crawled_content" in result_dict
|
||||
assert "No content available" in result_dict["crawled_content"]
|
||||
@@ -0,0 +1,119 @@
|
||||
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
from unittest.mock import Mock, call, patch
|
||||
|
||||
from src.tools.decorators import create_logged_tool
|
||||
|
||||
|
||||
class MockBaseTool:
|
||||
"""Mock base tool class for testing."""
|
||||
|
||||
def _run(self, *args, **kwargs):
|
||||
return "base_result"
|
||||
|
||||
|
||||
class TestLoggedToolMixin:
|
||||
def test_run_calls_log_operation(self):
|
||||
"""Test that _run calls _log_operation with correct parameters."""
|
||||
# Create a logged tool instance
|
||||
LoggedTool = create_logged_tool(MockBaseTool)
|
||||
tool = LoggedTool()
|
||||
|
||||
# Mock the _log_operation method
|
||||
tool._log_operation = Mock()
|
||||
|
||||
# Call _run with test parameters
|
||||
args = ("arg1", "arg2")
|
||||
kwargs = {"key1": "value1", "key2": "value2"}
|
||||
tool._run(*args, **kwargs)
|
||||
|
||||
# Verify _log_operation was called with correct parameters
|
||||
tool._log_operation.assert_called_once_with("_run", *args, **kwargs)
|
||||
|
||||
def test_run_calls_super_run(self):
|
||||
"""Test that _run calls the parent class _run method."""
|
||||
# Create a logged tool instance
|
||||
LoggedTool = create_logged_tool(MockBaseTool)
|
||||
tool = LoggedTool()
|
||||
|
||||
# Mock the parent _run method
|
||||
with patch.object(
|
||||
MockBaseTool, "_run", return_value="mocked_result"
|
||||
) as mock_super_run:
|
||||
args = ("arg1", "arg2")
|
||||
kwargs = {"key1": "value1"}
|
||||
result = tool._run(*args, **kwargs)
|
||||
|
||||
# Verify super()._run was called with correct parameters
|
||||
mock_super_run.assert_called_once_with(*args, **kwargs)
|
||||
# Verify the result is returned
|
||||
assert result == "mocked_result"
|
||||
|
||||
def test_run_logs_result(self):
|
||||
"""Test that _run logs the result with debug level."""
|
||||
LoggedTool = create_logged_tool(MockBaseTool)
|
||||
tool = LoggedTool()
|
||||
|
||||
with patch("src.tools.decorators.logger.debug") as mock_debug:
|
||||
tool._run("test_arg")
|
||||
|
||||
# Verify debug log was called with correct message
|
||||
mock_debug.assert_has_calls(
|
||||
[
|
||||
call("Tool MockBaseTool._run called with parameters: test_arg"),
|
||||
call("Tool MockBaseTool returned: base_result"),
|
||||
]
|
||||
)
|
||||
|
||||
def test_run_returns_super_result(self):
|
||||
"""Test that _run returns the result from parent class."""
|
||||
LoggedTool = create_logged_tool(MockBaseTool)
|
||||
tool = LoggedTool()
|
||||
|
||||
result = tool._run()
|
||||
assert result == "base_result"
|
||||
|
||||
def test_run_with_no_args(self):
|
||||
"""Test _run method with no arguments."""
|
||||
LoggedTool = create_logged_tool(MockBaseTool)
|
||||
tool = LoggedTool()
|
||||
|
||||
with patch("src.tools.decorators.logger.debug") as mock_debug:
|
||||
tool._log_operation = Mock()
|
||||
|
||||
result = tool._run()
|
||||
|
||||
# Verify _log_operation called with no args
|
||||
tool._log_operation.assert_called_once_with("_run")
|
||||
# Verify result logging
|
||||
mock_debug.assert_called_once()
|
||||
assert result == "base_result"
|
||||
|
||||
def test_run_with_mixed_args_kwargs(self):
|
||||
"""Test _run method with both positional and keyword arguments."""
|
||||
LoggedTool = create_logged_tool(MockBaseTool)
|
||||
tool = LoggedTool()
|
||||
|
||||
tool._log_operation = Mock()
|
||||
|
||||
args = ("pos1", "pos2")
|
||||
kwargs = {"kw1": "val1", "kw2": "val2"}
|
||||
result = tool._run(*args, **kwargs)
|
||||
|
||||
# Verify all arguments passed correctly
|
||||
tool._log_operation.assert_called_once_with("_run", *args, **kwargs)
|
||||
assert result == "base_result"
|
||||
|
||||
def test_run_class_name_replacement(self):
|
||||
"""Test that class name 'Logged' prefix is correctly removed in logging."""
|
||||
LoggedTool = create_logged_tool(MockBaseTool)
|
||||
tool = LoggedTool()
|
||||
|
||||
with patch("src.tools.decorators.logger.debug") as mock_debug:
|
||||
tool._run()
|
||||
|
||||
# Verify the logged class name has 'Logged' prefix removed
|
||||
call_args = mock_debug.call_args[0][0]
|
||||
assert "Tool MockBaseTool returned:" in call_args
|
||||
assert "LoggedMockBaseTool" not in call_args
|
||||
@@ -0,0 +1,222 @@
|
||||
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
import os
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
|
||||
from src.tools.python_repl import python_repl_tool
|
||||
|
||||
|
||||
class TestPythonReplTool:
|
||||
@patch.dict(os.environ, {"ENABLE_PYTHON_REPL": "true"})
|
||||
@patch("src.tools.python_repl.repl")
|
||||
@patch("src.tools.python_repl.logger")
|
||||
def test_successful_code_execution(self, mock_logger, mock_repl):
|
||||
# Arrange
|
||||
code = "print('Hello, World!')"
|
||||
expected_output = "Hello, World!\n"
|
||||
mock_repl.run.return_value = expected_output
|
||||
|
||||
# Act
|
||||
result = python_repl_tool(code)
|
||||
|
||||
# Assert
|
||||
mock_repl.run.assert_called_once_with(code)
|
||||
mock_logger.info.assert_called_with("Code execution successful")
|
||||
assert "Successfully executed:" in result
|
||||
assert code in result
|
||||
assert expected_output in result
|
||||
|
||||
@patch.dict(os.environ, {"ENABLE_PYTHON_REPL": "true"})
|
||||
@patch("src.tools.python_repl.repl")
|
||||
@patch("src.tools.python_repl.logger")
|
||||
def test_invalid_input_type(self, mock_logger, mock_repl):
|
||||
# Arrange
|
||||
invalid_code = 123
|
||||
|
||||
# Act & Assert - expect ValidationError when passing invalid input
|
||||
with pytest.raises(Exception): # Could be ValidationError or similar
|
||||
python_repl_tool(invalid_code)
|
||||
|
||||
mock_repl.run.assert_not_called()
|
||||
|
||||
@patch.dict(os.environ, {"ENABLE_PYTHON_REPL": "true"})
|
||||
@patch("src.tools.python_repl.repl")
|
||||
@patch("src.tools.python_repl.logger")
|
||||
def test_code_execution_with_error_in_result(self, mock_logger, mock_repl):
|
||||
# Arrange
|
||||
code = "invalid_function()"
|
||||
error_result = "NameError: name 'invalid_function' is not defined"
|
||||
mock_repl.run.return_value = error_result
|
||||
|
||||
# Act
|
||||
result = python_repl_tool(code)
|
||||
|
||||
# Assert
|
||||
mock_repl.run.assert_called_once_with(code)
|
||||
mock_logger.error.assert_called_with(error_result)
|
||||
assert "Error executing code:" in result
|
||||
assert code in result
|
||||
assert error_result in result
|
||||
|
||||
@patch.dict(os.environ, {"ENABLE_PYTHON_REPL": "true"})
|
||||
@patch("src.tools.python_repl.repl")
|
||||
@patch("src.tools.python_repl.logger")
|
||||
def test_code_execution_with_exception_in_result(self, mock_logger, mock_repl):
|
||||
# Arrange
|
||||
code = "1/0"
|
||||
exception_result = "ZeroDivisionError: division by zero"
|
||||
mock_repl.run.return_value = exception_result
|
||||
|
||||
# Act
|
||||
result = python_repl_tool(code)
|
||||
|
||||
# Assert
|
||||
mock_repl.run.assert_called_once_with(code)
|
||||
mock_logger.error.assert_called_with(exception_result)
|
||||
assert "Error executing code:" in result
|
||||
assert code in result
|
||||
assert exception_result in result
|
||||
|
||||
@patch.dict(os.environ, {"ENABLE_PYTHON_REPL": "true"})
|
||||
@patch("src.tools.python_repl.repl")
|
||||
@patch("src.tools.python_repl.logger")
|
||||
def test_code_execution_raises_exception(self, mock_logger, mock_repl):
|
||||
# Arrange
|
||||
code = "print('test')"
|
||||
exception = RuntimeError("REPL failed")
|
||||
mock_repl.run.side_effect = exception
|
||||
|
||||
# Act
|
||||
result = python_repl_tool(code)
|
||||
|
||||
# Assert
|
||||
mock_repl.run.assert_called_once_with(code)
|
||||
mock_logger.error.assert_called_with(repr(exception))
|
||||
assert "Error executing code:" in result
|
||||
assert code in result
|
||||
assert repr(exception) in result
|
||||
|
||||
@patch.dict(os.environ, {"ENABLE_PYTHON_REPL": "true"})
|
||||
@patch("src.tools.python_repl.repl")
|
||||
@patch("src.tools.python_repl.logger")
|
||||
def test_successful_execution_with_calculation(self, mock_logger, mock_repl):
|
||||
# Arrange
|
||||
code = "result = 2 + 3\nprint(result)"
|
||||
expected_output = "5\n"
|
||||
mock_repl.run.return_value = expected_output
|
||||
|
||||
# Act
|
||||
result = python_repl_tool(code)
|
||||
|
||||
# Assert
|
||||
mock_repl.run.assert_called_once_with(code)
|
||||
mock_logger.info.assert_any_call("Executing Python code")
|
||||
mock_logger.info.assert_any_call("Code execution successful")
|
||||
assert "Successfully executed:" in result
|
||||
assert code in result
|
||||
assert expected_output in result
|
||||
|
||||
@patch.dict(os.environ, {"ENABLE_PYTHON_REPL": "true"})
|
||||
@patch("src.tools.python_repl.repl")
|
||||
@patch("src.tools.python_repl.logger")
|
||||
def test_empty_string_code(self, mock_logger, mock_repl):
|
||||
# Arrange
|
||||
code = ""
|
||||
mock_repl.run.return_value = ""
|
||||
|
||||
# Act
|
||||
result = python_repl_tool(code)
|
||||
|
||||
# Assert
|
||||
mock_repl.run.assert_called_once_with(code)
|
||||
mock_logger.info.assert_called_with("Code execution successful")
|
||||
assert "Successfully executed:" in result
|
||||
|
||||
@patch.dict(os.environ, {"ENABLE_PYTHON_REPL": "true"})
|
||||
@patch("src.tools.python_repl.repl")
|
||||
@patch("src.tools.python_repl.logger")
|
||||
def test_logging_calls(self, mock_logger, mock_repl):
|
||||
# Arrange
|
||||
code = "x = 1"
|
||||
mock_repl.run.return_value = ""
|
||||
|
||||
# Act
|
||||
python_repl_tool(code)
|
||||
|
||||
# Assert
|
||||
mock_logger.info.assert_any_call("Executing Python code")
|
||||
mock_logger.info.assert_any_call("Code execution successful")
|
||||
|
||||
# New tests for configuration behavior
|
||||
@patch.dict(os.environ, {"ENABLE_PYTHON_REPL": "false"})
|
||||
@patch("src.tools.python_repl.logger")
|
||||
def test_tool_disabled(self, mock_logger):
|
||||
# Arrange
|
||||
code = "print('test')"
|
||||
|
||||
# Act
|
||||
result = python_repl_tool(code)
|
||||
|
||||
# Assert
|
||||
mock_logger.warning.assert_called_with(
|
||||
"Python REPL tool is disabled. Please enable it in environment configuration."
|
||||
)
|
||||
assert "Tool disabled:" in result
|
||||
assert "Python REPL tool is disabled" in result
|
||||
|
||||
@patch.dict(os.environ, {}, clear=True)
|
||||
@patch("src.tools.python_repl.logger")
|
||||
def test_tool_disabled_by_default(self, mock_logger):
|
||||
# Arrange - remove any existing ENABLE_PYTHON_REPL variable
|
||||
if "ENABLE_PYTHON_REPL" in os.environ:
|
||||
del os.environ["ENABLE_PYTHON_REPL"]
|
||||
code = "print('test')"
|
||||
|
||||
# Act
|
||||
result = python_repl_tool(code)
|
||||
|
||||
# Assert
|
||||
mock_logger.warning.assert_called_with(
|
||||
"Python REPL tool is disabled. Please enable it in environment configuration."
|
||||
)
|
||||
assert "Tool disabled:" in result
|
||||
|
||||
@pytest.mark.parametrize("env_value", ["true", "True", "TRUE", "1", "yes", "on"])
|
||||
@patch("src.tools.python_repl.repl")
|
||||
@patch("src.tools.python_repl.logger")
|
||||
def test_tool_enabled_with_various_truthy_values(
|
||||
self, mock_logger, mock_repl, env_value
|
||||
):
|
||||
# Arrange
|
||||
with patch.dict(os.environ, {"ENABLE_PYTHON_REPL": env_value}):
|
||||
code = "print('enabled')"
|
||||
expected_output = "enabled\n"
|
||||
mock_repl.run.return_value = expected_output
|
||||
|
||||
# Act
|
||||
result = python_repl_tool(code)
|
||||
|
||||
# Assert
|
||||
mock_repl.run.assert_called_once_with(code)
|
||||
assert "Successfully executed:" in result
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"env_value", ["false", "False", "FALSE", "0", "no", "off", ""]
|
||||
)
|
||||
@patch("src.tools.python_repl.logger")
|
||||
def test_tool_disabled_with_various_falsy_values(self, mock_logger, env_value):
|
||||
# Arrange
|
||||
with patch.dict(os.environ, {"ENABLE_PYTHON_REPL": env_value}):
|
||||
code = "print('disabled')"
|
||||
|
||||
# Act
|
||||
result = python_repl_tool(code)
|
||||
|
||||
# Assert
|
||||
mock_logger.warning.assert_called_with(
|
||||
"Python REPL tool is disabled. Please enable it in environment configuration."
|
||||
)
|
||||
assert "Tool disabled:" in result
|
||||
@@ -0,0 +1,276 @@
|
||||
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
import os
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
|
||||
from src.config import SearchEngine
|
||||
from src.tools.search import get_web_search_tool
|
||||
|
||||
|
||||
class TestGetWebSearchTool:
|
||||
@patch("src.tools.search.SELECTED_SEARCH_ENGINE", SearchEngine.TAVILY.value)
|
||||
def test_get_web_search_tool_tavily(self):
|
||||
tool = get_web_search_tool(max_search_results=5)
|
||||
assert tool.name == "web_search"
|
||||
assert tool.max_results == 5
|
||||
assert tool.include_raw_content is True
|
||||
assert tool.include_images is True
|
||||
assert tool.include_image_descriptions is True
|
||||
assert tool.include_answer is False
|
||||
assert tool.search_depth == "advanced"
|
||||
|
||||
@patch("src.tools.search.SELECTED_SEARCH_ENGINE", SearchEngine.DUCKDUCKGO.value)
|
||||
def test_get_web_search_tool_duckduckgo(self):
|
||||
tool = get_web_search_tool(max_search_results=3)
|
||||
assert tool.name == "web_search"
|
||||
assert tool.max_results == 3
|
||||
|
||||
@patch("src.tools.search.SELECTED_SEARCH_ENGINE", SearchEngine.BRAVE_SEARCH.value)
|
||||
@patch.dict(os.environ, {"BRAVE_SEARCH_API_KEY": "test_api_key"})
|
||||
def test_get_web_search_tool_brave(self):
|
||||
tool = get_web_search_tool(max_search_results=4)
|
||||
assert tool.name == "web_search"
|
||||
assert tool.search_wrapper.api_key == "test_api_key"
|
||||
|
||||
@patch("src.tools.search.SELECTED_SEARCH_ENGINE", SearchEngine.ARXIV.value)
|
||||
def test_get_web_search_tool_arxiv(self):
|
||||
tool = get_web_search_tool(max_search_results=2)
|
||||
assert tool.name == "web_search"
|
||||
assert tool.api_wrapper.top_k_results == 2
|
||||
assert tool.api_wrapper.load_max_docs == 2
|
||||
assert tool.api_wrapper.load_all_available_meta is True
|
||||
|
||||
@patch("src.tools.search.SELECTED_SEARCH_ENGINE", "unsupported_engine")
|
||||
def test_get_web_search_tool_unsupported_engine(self):
|
||||
with pytest.raises(
|
||||
ValueError, match="Unsupported search engine: unsupported_engine"
|
||||
):
|
||||
get_web_search_tool(max_search_results=1)
|
||||
|
||||
@patch("src.tools.search.SELECTED_SEARCH_ENGINE", SearchEngine.BRAVE_SEARCH.value)
|
||||
@patch.dict(os.environ, {}, clear=True)
|
||||
def test_get_web_search_tool_brave_no_api_key(self):
|
||||
tool = get_web_search_tool(max_search_results=1)
|
||||
assert tool.search_wrapper.api_key == ""
|
||||
|
||||
@patch("src.tools.search.SELECTED_SEARCH_ENGINE", SearchEngine.TAVILY.value)
|
||||
@patch("src.tools.search.load_yaml_config")
|
||||
def test_get_web_search_tool_tavily_with_custom_config(self, mock_config):
|
||||
"""Test Tavily tool with custom configuration values."""
|
||||
mock_config.return_value = {
|
||||
"SEARCH_ENGINE": {
|
||||
"include_answer": True,
|
||||
"search_depth": "basic",
|
||||
"include_raw_content": False,
|
||||
"include_images": False,
|
||||
"include_image_descriptions": True,
|
||||
"include_domains": ["example.com"],
|
||||
"exclude_domains": ["spam.com"],
|
||||
}
|
||||
}
|
||||
tool = get_web_search_tool(max_search_results=5)
|
||||
assert tool.name == "web_search"
|
||||
assert tool.max_results == 5
|
||||
assert tool.include_answer is True
|
||||
assert tool.search_depth == "basic"
|
||||
assert tool.include_raw_content is False
|
||||
assert tool.include_images is False
|
||||
# include_image_descriptions should be False because include_images is False
|
||||
assert tool.include_image_descriptions is False
|
||||
assert tool.include_domains == ["example.com"]
|
||||
assert tool.exclude_domains == ["spam.com"]
|
||||
|
||||
@patch("src.tools.search.SELECTED_SEARCH_ENGINE", SearchEngine.TAVILY.value)
|
||||
@patch("src.tools.search.load_yaml_config")
|
||||
def test_get_web_search_tool_tavily_with_empty_config(self, mock_config):
|
||||
"""Test Tavily tool uses defaults when config is empty."""
|
||||
mock_config.return_value = {"SEARCH_ENGINE": {}}
|
||||
tool = get_web_search_tool(max_search_results=10)
|
||||
assert tool.name == "web_search"
|
||||
assert tool.max_results == 10
|
||||
assert tool.include_answer is False
|
||||
assert tool.search_depth == "advanced"
|
||||
assert tool.include_raw_content is True
|
||||
assert tool.include_images is True
|
||||
assert tool.include_image_descriptions is True
|
||||
assert tool.include_domains == []
|
||||
assert tool.exclude_domains == []
|
||||
|
||||
@patch("src.tools.search.SELECTED_SEARCH_ENGINE", SearchEngine.TAVILY.value)
|
||||
@patch("src.tools.search.load_yaml_config")
|
||||
def test_get_web_search_tool_tavily_image_descriptions_disabled_when_images_disabled(
|
||||
self, mock_config
|
||||
):
|
||||
"""Test that include_image_descriptions is False when include_images is False."""
|
||||
mock_config.return_value = {
|
||||
"SEARCH_ENGINE": {
|
||||
"include_images": False,
|
||||
"include_image_descriptions": True, # This should be ignored
|
||||
}
|
||||
}
|
||||
tool = get_web_search_tool(max_search_results=5)
|
||||
assert tool.include_images is False
|
||||
assert tool.include_image_descriptions is False
|
||||
|
||||
@patch("src.tools.search.SELECTED_SEARCH_ENGINE", SearchEngine.TAVILY.value)
|
||||
@patch("src.tools.search.load_yaml_config")
|
||||
def test_get_web_search_tool_tavily_partial_config(self, mock_config):
|
||||
"""Test Tavily tool with partial configuration."""
|
||||
mock_config.return_value = {
|
||||
"SEARCH_ENGINE": {
|
||||
"include_answer": True,
|
||||
"include_domains": ["trusted.com"],
|
||||
}
|
||||
}
|
||||
tool = get_web_search_tool(max_search_results=3)
|
||||
assert tool.include_answer is True
|
||||
assert tool.search_depth == "advanced" # default
|
||||
assert tool.include_raw_content is True # default
|
||||
assert tool.include_domains == ["trusted.com"]
|
||||
assert tool.exclude_domains == [] # default
|
||||
|
||||
@patch("src.tools.search.SELECTED_SEARCH_ENGINE", SearchEngine.TAVILY.value)
|
||||
@patch("src.tools.search.load_yaml_config")
|
||||
def test_get_web_search_tool_tavily_with_no_config_file(self, mock_config):
|
||||
"""Test Tavily tool when config file doesn't exist."""
|
||||
mock_config.return_value = {}
|
||||
tool = get_web_search_tool(max_search_results=5)
|
||||
assert tool.name == "web_search"
|
||||
assert tool.max_results == 5
|
||||
assert tool.include_answer is False
|
||||
assert tool.search_depth == "advanced"
|
||||
assert tool.include_raw_content is True
|
||||
assert tool.include_images is True
|
||||
|
||||
@patch("src.tools.search.SELECTED_SEARCH_ENGINE", SearchEngine.TAVILY.value)
|
||||
@patch("src.tools.search.load_yaml_config")
|
||||
def test_get_web_search_tool_tavily_multiple_domains(self, mock_config):
|
||||
"""Test Tavily tool with multiple domains in include/exclude lists."""
|
||||
mock_config.return_value = {
|
||||
"SEARCH_ENGINE": {
|
||||
"include_domains": ["example.com", "trusted.com", "gov.cn"],
|
||||
"exclude_domains": ["spam.com", "scam.org"],
|
||||
}
|
||||
}
|
||||
tool = get_web_search_tool(max_search_results=5)
|
||||
assert tool.include_domains == ["example.com", "trusted.com", "gov.cn"]
|
||||
assert tool.exclude_domains == ["spam.com", "scam.org"]
|
||||
|
||||
@patch("src.tools.search.SELECTED_SEARCH_ENGINE", SearchEngine.TAVILY.value)
|
||||
@patch("src.tools.search.load_yaml_config")
|
||||
def test_tavily_with_no_search_engine_section(self, mock_config):
|
||||
"""Test Tavily tool when SEARCH_ENGINE section doesn't exist in config."""
|
||||
mock_config.return_value = {"OTHER_CONFIG": {}}
|
||||
tool = get_web_search_tool(max_search_results=5)
|
||||
assert tool.name == "web_search"
|
||||
assert tool.max_results == 5
|
||||
assert tool.include_answer is False
|
||||
assert tool.search_depth == "advanced"
|
||||
assert tool.include_raw_content is True
|
||||
assert tool.include_images is True
|
||||
assert tool.include_domains == []
|
||||
assert tool.exclude_domains == []
|
||||
|
||||
@patch("src.tools.search.SELECTED_SEARCH_ENGINE", SearchEngine.TAVILY.value)
|
||||
@patch("src.tools.search.load_yaml_config")
|
||||
def test_tavily_with_completely_empty_config(self, mock_config):
|
||||
"""Test Tavily tool with completely empty config."""
|
||||
mock_config.return_value = {}
|
||||
tool = get_web_search_tool(max_search_results=5)
|
||||
assert tool.name == "web_search"
|
||||
assert tool.max_results == 5
|
||||
assert tool.include_answer is False
|
||||
assert tool.search_depth == "advanced"
|
||||
assert tool.include_raw_content is True
|
||||
assert tool.include_images is True
|
||||
|
||||
@patch("src.tools.search.SELECTED_SEARCH_ENGINE", SearchEngine.TAVILY.value)
|
||||
@patch("src.tools.search.load_yaml_config")
|
||||
def test_tavily_with_only_include_answer_param(self, mock_config):
|
||||
"""Test Tavily tool with only include_answer parameter specified."""
|
||||
mock_config.return_value = {"SEARCH_ENGINE": {"include_answer": True}}
|
||||
tool = get_web_search_tool(max_search_results=5)
|
||||
assert tool.include_answer is True
|
||||
assert tool.search_depth == "advanced"
|
||||
assert tool.include_raw_content is True
|
||||
assert tool.include_images is True
|
||||
|
||||
@patch("src.tools.search.SELECTED_SEARCH_ENGINE", SearchEngine.TAVILY.value)
|
||||
@patch("src.tools.search.load_yaml_config")
|
||||
def test_tavily_with_only_search_depth_param(self, mock_config):
|
||||
"""Test Tavily tool with only search_depth parameter specified."""
|
||||
mock_config.return_value = {"SEARCH_ENGINE": {"search_depth": "basic"}}
|
||||
tool = get_web_search_tool(max_search_results=5)
|
||||
assert tool.search_depth == "basic"
|
||||
assert tool.include_answer is False
|
||||
assert tool.include_raw_content is True
|
||||
assert tool.include_images is True
|
||||
|
||||
@patch("src.tools.search.SELECTED_SEARCH_ENGINE", SearchEngine.TAVILY.value)
|
||||
@patch("src.tools.search.load_yaml_config")
|
||||
def test_tavily_with_only_include_domains_param(self, mock_config):
|
||||
"""Test Tavily tool with only include_domains parameter specified."""
|
||||
mock_config.return_value = {
|
||||
"SEARCH_ENGINE": {"include_domains": ["example.com"]}
|
||||
}
|
||||
tool = get_web_search_tool(max_search_results=5)
|
||||
assert tool.include_domains == ["example.com"]
|
||||
assert tool.exclude_domains == []
|
||||
assert tool.include_answer is False
|
||||
assert tool.search_depth == "advanced"
|
||||
|
||||
@patch("src.tools.search.SELECTED_SEARCH_ENGINE", SearchEngine.TAVILY.value)
|
||||
@patch("src.tools.search.load_yaml_config")
|
||||
def test_tavily_with_explicit_false_boolean_values(self, mock_config):
|
||||
"""Test that explicitly False boolean values are respected (not treated as missing)."""
|
||||
mock_config.return_value = {
|
||||
"SEARCH_ENGINE": {
|
||||
"include_answer": False,
|
||||
"include_raw_content": False,
|
||||
"include_images": False,
|
||||
}
|
||||
}
|
||||
tool = get_web_search_tool(max_search_results=5)
|
||||
assert tool.include_answer is False
|
||||
assert tool.include_raw_content is False
|
||||
assert tool.include_images is False
|
||||
assert tool.include_image_descriptions is False
|
||||
|
||||
@patch("src.tools.search.SELECTED_SEARCH_ENGINE", SearchEngine.TAVILY.value)
|
||||
@patch("src.tools.search.load_yaml_config")
|
||||
def test_tavily_with_empty_domain_lists(self, mock_config):
|
||||
"""Test that empty domain lists are treated as optional."""
|
||||
mock_config.return_value = {
|
||||
"SEARCH_ENGINE": {
|
||||
"include_domains": [],
|
||||
"exclude_domains": [],
|
||||
}
|
||||
}
|
||||
tool = get_web_search_tool(max_search_results=5)
|
||||
assert tool.include_domains == []
|
||||
assert tool.exclude_domains == []
|
||||
|
||||
@patch("src.tools.search.SELECTED_SEARCH_ENGINE", SearchEngine.TAVILY.value)
|
||||
@patch("src.tools.search.load_yaml_config")
|
||||
def test_tavily_all_parameters_optional_mix(self, mock_config):
|
||||
"""Test that any combination of optional parameters works."""
|
||||
mock_config.return_value = {
|
||||
"SEARCH_ENGINE": {
|
||||
"include_answer": True,
|
||||
"include_images": False,
|
||||
# Deliberately omit search_depth, include_raw_content, domains
|
||||
}
|
||||
}
|
||||
tool = get_web_search_tool(max_search_results=5)
|
||||
assert tool.include_answer is True
|
||||
assert tool.include_images is False
|
||||
assert (
|
||||
tool.include_image_descriptions is False
|
||||
) # should be False since include_images is False
|
||||
assert tool.search_depth == "advanced" # default
|
||||
assert tool.include_raw_content is True # default
|
||||
assert tool.include_domains == [] # default
|
||||
assert tool.exclude_domains == [] # default
|
||||
@@ -0,0 +1,263 @@
|
||||
import pytest
|
||||
|
||||
from src.tools.search_postprocessor import SearchResultPostProcessor
|
||||
|
||||
|
||||
class TestSearchResultPostProcessor:
|
||||
"""Test cases for SearchResultPostProcessor"""
|
||||
|
||||
@pytest.fixture
|
||||
def post_processor(self):
|
||||
"""Create a SearchResultPostProcessor instance for testing"""
|
||||
return SearchResultPostProcessor(
|
||||
min_score_threshold=0.5, max_content_length_per_page=100
|
||||
)
|
||||
|
||||
def test_process_results_empty_input(self, post_processor):
|
||||
"""Test processing empty results"""
|
||||
results = []
|
||||
processed = post_processor.process_results(results)
|
||||
assert processed == []
|
||||
|
||||
def test_process_results_with_valid_page_results(self, post_processor):
|
||||
"""Test processing valid page results"""
|
||||
results = [
|
||||
{
|
||||
"type": "page",
|
||||
"title": "Test Page",
|
||||
"url": "https://example.com",
|
||||
"content": "Test content",
|
||||
"score": 0.8,
|
||||
}
|
||||
]
|
||||
processed = post_processor.process_results(results)
|
||||
assert len(processed) == 1
|
||||
assert processed[0]["title"] == "Test Page"
|
||||
assert processed[0]["url"] == "https://example.com"
|
||||
assert processed[0]["content"] == "Test content"
|
||||
assert processed[0]["score"] == 0.8
|
||||
|
||||
def test_process_results_filter_low_score(self, post_processor):
|
||||
"""Test filtering out low score results"""
|
||||
results = [
|
||||
{
|
||||
"type": "page",
|
||||
"title": "Low Score Page",
|
||||
"url": "https://example.com/low",
|
||||
"content": "Low score content",
|
||||
"score": 0.3, # Below threshold of 0.5
|
||||
},
|
||||
{
|
||||
"type": "page",
|
||||
"title": "High Score Page",
|
||||
"url": "https://example.com/high",
|
||||
"content": "High score content",
|
||||
"score": 0.9,
|
||||
},
|
||||
]
|
||||
processed = post_processor.process_results(results)
|
||||
assert len(processed) == 1
|
||||
assert processed[0]["title"] == "High Score Page"
|
||||
|
||||
def test_process_results_remove_duplicates(self, post_processor):
|
||||
"""Test removing duplicate URLs"""
|
||||
results = [
|
||||
{
|
||||
"type": "page",
|
||||
"title": "Page 1",
|
||||
"url": "https://example.com",
|
||||
"content": "Content 1",
|
||||
"score": 0.8,
|
||||
},
|
||||
{
|
||||
"type": "page",
|
||||
"title": "Page 2",
|
||||
"url": "https://example.com", # Duplicate URL
|
||||
"content": "Content 2",
|
||||
"score": 0.7,
|
||||
},
|
||||
]
|
||||
processed = post_processor.process_results(results)
|
||||
assert len(processed) == 1
|
||||
assert processed[0]["title"] == "Page 1" # First one should be kept
|
||||
|
||||
def test_process_results_sort_by_score(self, post_processor):
|
||||
"""Test sorting results by score in descending order"""
|
||||
results = [
|
||||
{
|
||||
"type": "page",
|
||||
"title": "Low Score",
|
||||
"url": "https://example.com/low",
|
||||
"content": "Low score content",
|
||||
"score": 0.3,
|
||||
},
|
||||
{
|
||||
"type": "page",
|
||||
"title": "High Score",
|
||||
"url": "https://example.com/high",
|
||||
"content": "High score content",
|
||||
"score": 0.9,
|
||||
},
|
||||
{
|
||||
"type": "page",
|
||||
"title": "Medium Score",
|
||||
"url": "https://example.com/medium",
|
||||
"content": "Medium score content",
|
||||
"score": 0.6,
|
||||
},
|
||||
]
|
||||
processed = post_processor.process_results(results)
|
||||
assert len(processed) == 2 # Low score filtered out
|
||||
# Should be sorted by score descending
|
||||
assert processed[0]["title"] == "High Score"
|
||||
assert processed[1]["title"] == "Medium Score"
|
||||
|
||||
def test_process_results_truncate_long_content(self, post_processor):
|
||||
"""Test truncating long content"""
|
||||
long_content = "A" * 150 # Longer than max_content_length of 100
|
||||
results = [
|
||||
{
|
||||
"type": "page",
|
||||
"title": "Long Content Page",
|
||||
"url": "https://example.com",
|
||||
"content": long_content,
|
||||
"score": 0.8,
|
||||
}
|
||||
]
|
||||
processed = post_processor.process_results(results)
|
||||
assert len(processed) == 1
|
||||
assert len(processed[0]["content"]) == 103 # 100 + "..."
|
||||
assert processed[0]["content"].endswith("...")
|
||||
|
||||
def test_process_results_remove_base64_images(self, post_processor):
|
||||
"""Test removing base64 images from content"""
|
||||
content_with_base64 = (
|
||||
"Content with image "
|
||||
+ "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8/5+hHgAHggJ/PchI7wAAAABJRU5ErkJggg=="
|
||||
)
|
||||
results = [
|
||||
{
|
||||
"type": "page",
|
||||
"title": "Page with Base64",
|
||||
"url": "https://example.com",
|
||||
"content": content_with_base64,
|
||||
"score": 0.8,
|
||||
}
|
||||
]
|
||||
processed = post_processor.process_results(results)
|
||||
assert len(processed) == 1
|
||||
assert processed[0]["content"] == "Content with image "
|
||||
|
||||
def test_process_results_with_image_type(self, post_processor):
|
||||
"""Test processing image type results"""
|
||||
results = [
|
||||
{
|
||||
"type": "image",
|
||||
"image_url": "https://example.com/image.jpg",
|
||||
"image_description": "Test image",
|
||||
}
|
||||
]
|
||||
processed = post_processor.process_results(results)
|
||||
assert len(processed) == 1
|
||||
assert processed[0]["type"] == "image"
|
||||
assert processed[0]["image_url"] == "https://example.com/image.jpg"
|
||||
assert processed[0]["image_description"] == "Test image"
|
||||
|
||||
def test_process_results_filter_base64_image_urls(self, post_processor):
|
||||
"""Test filtering out image results with base64 URLs"""
|
||||
results = [
|
||||
{
|
||||
"type": "image",
|
||||
"image_url": "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8/5+hHgAHggJ/PchI7wAAAABJRU5ErkJggg==",
|
||||
"image_description": "Base64 image",
|
||||
},
|
||||
{
|
||||
"type": "image",
|
||||
"image_url": "https://example.com/image.jpg",
|
||||
"image_description": "Regular image",
|
||||
},
|
||||
]
|
||||
processed = post_processor.process_results(results)
|
||||
assert len(processed) == 1
|
||||
assert processed[0]["image_url"] == "https://example.com/image.jpg"
|
||||
|
||||
def test_process_results_truncate_long_image_description(self, post_processor):
|
||||
"""Test truncating long image descriptions"""
|
||||
long_description = "A" * 150 # Longer than max_content_length of 100
|
||||
results = [
|
||||
{
|
||||
"type": "image",
|
||||
"image_url": "https://example.com/image.jpg",
|
||||
"image_description": long_description,
|
||||
}
|
||||
]
|
||||
processed = post_processor.process_results(results)
|
||||
assert len(processed) == 1
|
||||
assert len(processed[0]["image_description"]) == 103 # 100 + "..."
|
||||
assert processed[0]["image_description"].endswith("...")
|
||||
|
||||
def test_process_results_other_types_passthrough(self, post_processor):
|
||||
"""Test that other result types pass through unchanged"""
|
||||
results = [
|
||||
{
|
||||
"type": "video",
|
||||
"title": "Test Video",
|
||||
"url": "https://example.com/video.mp4",
|
||||
"score": 0.8,
|
||||
}
|
||||
]
|
||||
processed = post_processor.process_results(results)
|
||||
assert len(processed) == 1
|
||||
assert processed[0]["type"] == "video"
|
||||
assert processed[0]["title"] == "Test Video"
|
||||
|
||||
def test_process_results_truncate_long_content_with_no_config(self):
|
||||
"""Test truncating long content"""
|
||||
post_processor = SearchResultPostProcessor(None, None)
|
||||
long_content = "A" * 150 # Longer than max_content_length of 100
|
||||
results = [
|
||||
{
|
||||
"type": "page",
|
||||
"title": "Long Content Page",
|
||||
"url": "https://example.com",
|
||||
"content": long_content,
|
||||
"score": 0.8,
|
||||
}
|
||||
]
|
||||
processed = post_processor.process_results(results)
|
||||
assert len(processed) == 1
|
||||
assert len(processed[0]["content"]) == len("A" * 150)
|
||||
|
||||
def test_process_results_truncate_long_content_with_max_content_length_config(self):
|
||||
"""Test truncating long content"""
|
||||
post_processor = SearchResultPostProcessor(None, 100)
|
||||
long_content = "A" * 150 # Longer than max_content_length of 100
|
||||
results = [
|
||||
{
|
||||
"type": "page",
|
||||
"title": "Long Content Page",
|
||||
"url": "https://example.com",
|
||||
"content": long_content,
|
||||
"score": 0.8,
|
||||
}
|
||||
]
|
||||
processed = post_processor.process_results(results)
|
||||
assert len(processed) == 1
|
||||
assert len(processed[0]["content"]) == 103
|
||||
assert processed[0]["content"].endswith("...")
|
||||
|
||||
def test_process_results_truncate_long_content_with_min_score_config(self):
|
||||
"""Test truncating long content"""
|
||||
post_processor = SearchResultPostProcessor(0.8, None)
|
||||
long_content = "A" * 150 # Longer than max_content_length of 100
|
||||
results = [
|
||||
{
|
||||
"type": "page",
|
||||
"title": "Long Content Page",
|
||||
"url": "https://example.com",
|
||||
"content": long_content,
|
||||
"score": 0.3,
|
||||
}
|
||||
]
|
||||
processed = post_processor.process_results(results)
|
||||
assert len(processed) == 0
|
||||
@@ -0,0 +1,207 @@
|
||||
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
|
||||
# SPDX-License-Identifier: MIT
|
||||
import json
|
||||
from unittest.mock import AsyncMock, MagicMock, Mock, patch
|
||||
|
||||
import pytest
|
||||
import requests
|
||||
|
||||
from src.tools.tavily_search.tavily_search_api_wrapper import (
|
||||
EnhancedTavilySearchAPIWrapper,
|
||||
)
|
||||
|
||||
|
||||
class TestEnhancedTavilySearchAPIWrapper:
|
||||
@pytest.fixture
|
||||
def wrapper(self):
|
||||
with patch(
|
||||
"src.tools.tavily_search.tavily_search_api_wrapper.OriginalTavilySearchAPIWrapper"
|
||||
):
|
||||
wrapper = EnhancedTavilySearchAPIWrapper(tavily_api_key="dummy-key")
|
||||
# The parent class is mocked, so initialization won't fail
|
||||
return wrapper
|
||||
|
||||
@pytest.fixture
|
||||
def mock_response_data(self):
|
||||
return {
|
||||
"results": [
|
||||
{
|
||||
"title": "Test Title",
|
||||
"url": "https://example.com",
|
||||
"content": "Test content",
|
||||
"score": 0.9,
|
||||
"raw_content": "Raw test content",
|
||||
}
|
||||
],
|
||||
"images": [
|
||||
{
|
||||
"url": "https://example.com/image.jpg",
|
||||
"description": "Test image description",
|
||||
}
|
||||
],
|
||||
}
|
||||
|
||||
@patch("src.tools.tavily_search.tavily_search_api_wrapper.requests.post")
|
||||
def test_raw_results_success(self, mock_post, wrapper, mock_response_data):
|
||||
mock_response = Mock()
|
||||
mock_response.json.return_value = mock_response_data
|
||||
mock_response.raise_for_status.return_value = None
|
||||
mock_post.return_value = mock_response
|
||||
|
||||
result = wrapper.raw_results("test query", max_results=10)
|
||||
|
||||
assert result == mock_response_data
|
||||
mock_post.assert_called_once()
|
||||
call_args = mock_post.call_args
|
||||
assert "json" in call_args.kwargs
|
||||
assert call_args.kwargs["json"]["query"] == "test query"
|
||||
assert call_args.kwargs["json"]["max_results"] == 10
|
||||
|
||||
@patch("src.tools.tavily_search.tavily_search_api_wrapper.requests.post")
|
||||
def test_raw_results_with_all_parameters(
|
||||
self, mock_post, wrapper, mock_response_data
|
||||
):
|
||||
mock_response = Mock()
|
||||
mock_response.json.return_value = mock_response_data
|
||||
mock_response.raise_for_status.return_value = None
|
||||
mock_post.return_value = mock_response
|
||||
|
||||
result = wrapper.raw_results(
|
||||
"test query",
|
||||
max_results=3,
|
||||
search_depth="basic",
|
||||
include_domains=["example.com"],
|
||||
exclude_domains=["spam.com"],
|
||||
include_answer=True,
|
||||
include_raw_content=True,
|
||||
include_images=True,
|
||||
include_image_descriptions=True,
|
||||
)
|
||||
|
||||
assert result == mock_response_data
|
||||
call_args = mock_post.call_args
|
||||
params = call_args.kwargs["json"]
|
||||
assert params["include_domains"] == ["example.com"]
|
||||
assert params["exclude_domains"] == ["spam.com"]
|
||||
assert params["include_answer"] is True
|
||||
assert params["include_raw_content"] is True
|
||||
|
||||
@patch("src.tools.tavily_search.tavily_search_api_wrapper.requests.post")
|
||||
def test_raw_results_http_error(self, mock_post, wrapper):
|
||||
mock_response = Mock()
|
||||
mock_response.raise_for_status.side_effect = requests.HTTPError("API Error")
|
||||
mock_post.return_value = mock_response
|
||||
|
||||
with pytest.raises(requests.HTTPError):
|
||||
wrapper.raw_results("test query")
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_raw_results_async_success(self, wrapper, mock_response_data):
|
||||
# Create a mock that acts as both the response and its context manager
|
||||
mock_response_cm = AsyncMock()
|
||||
mock_response_cm.__aenter__ = AsyncMock(return_value=mock_response_cm)
|
||||
mock_response_cm.__aexit__ = AsyncMock(return_value=None)
|
||||
mock_response_cm.status = 200
|
||||
mock_response_cm.text = AsyncMock(return_value=json.dumps(mock_response_data))
|
||||
|
||||
# Create mock session that returns the context manager
|
||||
mock_session = AsyncMock()
|
||||
mock_session.post = MagicMock(
|
||||
return_value=mock_response_cm
|
||||
) # Use MagicMock, not AsyncMock
|
||||
|
||||
# Create mock session class
|
||||
mock_session_cm = AsyncMock()
|
||||
mock_session_cm.__aenter__ = AsyncMock(return_value=mock_session)
|
||||
mock_session_cm.__aexit__ = AsyncMock(return_value=None)
|
||||
|
||||
with patch(
|
||||
"src.tools.tavily_search.tavily_search_api_wrapper.aiohttp.ClientSession",
|
||||
return_value=mock_session_cm,
|
||||
):
|
||||
result = await wrapper.raw_results_async("test query")
|
||||
|
||||
assert result == mock_response_data
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_raw_results_async_error(self, wrapper):
|
||||
# Create a mock that acts as both the response and its context manager
|
||||
mock_response_cm = AsyncMock()
|
||||
mock_response_cm.__aenter__ = AsyncMock(return_value=mock_response_cm)
|
||||
mock_response_cm.__aexit__ = AsyncMock(return_value=None)
|
||||
mock_response_cm.status = 400
|
||||
mock_response_cm.reason = "Bad Request"
|
||||
|
||||
# Create mock session that returns the context manager
|
||||
mock_session = AsyncMock()
|
||||
mock_session.post = MagicMock(
|
||||
return_value=mock_response_cm
|
||||
) # Use MagicMock, not AsyncMock
|
||||
|
||||
# Create mock session class
|
||||
mock_session_cm = AsyncMock()
|
||||
mock_session_cm.__aenter__ = AsyncMock(return_value=mock_session)
|
||||
mock_session_cm.__aexit__ = AsyncMock(return_value=None)
|
||||
|
||||
with patch(
|
||||
"src.tools.tavily_search.tavily_search_api_wrapper.aiohttp.ClientSession",
|
||||
return_value=mock_session_cm,
|
||||
):
|
||||
with pytest.raises(Exception, match="Error 400: Bad Request"):
|
||||
await wrapper.raw_results_async("test query")
|
||||
|
||||
def test_clean_results_with_images(self, wrapper, mock_response_data):
|
||||
result = wrapper.clean_results_with_images(mock_response_data)
|
||||
|
||||
assert len(result) == 2
|
||||
|
||||
# Test page result
|
||||
page_result = result[0]
|
||||
assert page_result["type"] == "page"
|
||||
assert page_result["title"] == "Test Title"
|
||||
assert page_result["url"] == "https://example.com"
|
||||
assert page_result["content"] == "Test content"
|
||||
assert page_result["score"] == 0.9
|
||||
assert page_result["raw_content"] == "Raw test content"
|
||||
|
||||
# Test image result
|
||||
image_result = result[1]
|
||||
assert image_result["type"] == "image_url"
|
||||
assert image_result["image_url"] == {"url": "https://example.com/image.jpg"}
|
||||
assert image_result["image_description"] == "Test image description"
|
||||
|
||||
def test_clean_results_without_raw_content(self, wrapper):
|
||||
data = {
|
||||
"results": [
|
||||
{
|
||||
"title": "Test Title",
|
||||
"url": "https://example.com",
|
||||
"content": "Test content",
|
||||
"score": 0.9,
|
||||
}
|
||||
],
|
||||
"images": [],
|
||||
}
|
||||
|
||||
result = wrapper.clean_results_with_images(data)
|
||||
|
||||
assert len(result) == 1
|
||||
assert "raw_content" not in result[0]
|
||||
|
||||
def test_clean_results_empty_images(self, wrapper):
|
||||
data = {
|
||||
"results": [
|
||||
{
|
||||
"title": "Test Title",
|
||||
"url": "https://example.com",
|
||||
"content": "Test content",
|
||||
"score": 0.9,
|
||||
}
|
||||
],
|
||||
"images": [],
|
||||
}
|
||||
|
||||
result = wrapper.clean_results_with_images(data)
|
||||
|
||||
assert len(result) == 1
|
||||
assert result[0]["type"] == "page"
|
||||
@@ -0,0 +1,206 @@
|
||||
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
import json
|
||||
from unittest.mock import AsyncMock, Mock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from src.tools.tavily_search.tavily_search_api_wrapper import (
|
||||
EnhancedTavilySearchAPIWrapper,
|
||||
)
|
||||
from src.tools.tavily_search.tavily_search_results_with_images import (
|
||||
TavilySearchWithImages,
|
||||
)
|
||||
|
||||
|
||||
class TestTavilySearchWithImages:
|
||||
@pytest.fixture
|
||||
def mock_api_wrapper(self):
|
||||
"""Create a mock API wrapper."""
|
||||
wrapper = Mock(spec=EnhancedTavilySearchAPIWrapper)
|
||||
return wrapper
|
||||
|
||||
@pytest.fixture
|
||||
def search_tool(self, mock_api_wrapper):
|
||||
"""Create a TavilySearchWithImages instance with mocked dependencies."""
|
||||
tool = TavilySearchWithImages(
|
||||
max_results=5,
|
||||
include_answer=True,
|
||||
include_raw_content=True,
|
||||
include_images=True,
|
||||
include_image_descriptions=True,
|
||||
)
|
||||
tool.api_wrapper = mock_api_wrapper
|
||||
return tool
|
||||
|
||||
@pytest.fixture
|
||||
def sample_raw_results(self):
|
||||
"""Sample raw results from Tavily API."""
|
||||
return {
|
||||
"query": "test query",
|
||||
"answer": "Test answer",
|
||||
"images": ["https://example.com/image1.jpg"],
|
||||
"results": [
|
||||
{
|
||||
"title": "Test Title",
|
||||
"url": "https://example.com",
|
||||
"content": "Test content",
|
||||
"score": 0.95,
|
||||
"raw_content": "Raw test content",
|
||||
}
|
||||
],
|
||||
"response_time": 1.5,
|
||||
}
|
||||
|
||||
@pytest.fixture
|
||||
def sample_cleaned_results(self):
|
||||
"""Sample cleaned results."""
|
||||
return [
|
||||
{
|
||||
"title": "Test Title",
|
||||
"url": "https://example.com",
|
||||
"content": "Test content",
|
||||
}
|
||||
]
|
||||
|
||||
def test_init_default_values(self):
|
||||
"""Test initialization with default values."""
|
||||
tool = TavilySearchWithImages()
|
||||
assert tool.include_image_descriptions is False
|
||||
assert isinstance(tool.api_wrapper, EnhancedTavilySearchAPIWrapper)
|
||||
|
||||
def test_init_custom_values(self):
|
||||
"""Test initialization with custom values."""
|
||||
tool = TavilySearchWithImages(max_results=10, include_image_descriptions=True)
|
||||
assert tool.max_results == 10
|
||||
assert tool.include_image_descriptions is True
|
||||
|
||||
def test_run_success(
|
||||
self,
|
||||
search_tool,
|
||||
mock_api_wrapper,
|
||||
sample_raw_results,
|
||||
sample_cleaned_results,
|
||||
):
|
||||
"""Test successful synchronous run."""
|
||||
mock_api_wrapper.raw_results.return_value = sample_raw_results
|
||||
mock_api_wrapper.clean_results_with_images.return_value = sample_cleaned_results
|
||||
|
||||
result, raw = search_tool._run("test query")
|
||||
|
||||
assert result == json.dumps(sample_cleaned_results, ensure_ascii=False)
|
||||
assert raw == sample_raw_results
|
||||
|
||||
mock_api_wrapper.raw_results.assert_called_once_with(
|
||||
"test query",
|
||||
search_tool.max_results,
|
||||
search_tool.search_depth,
|
||||
search_tool.include_domains,
|
||||
search_tool.exclude_domains,
|
||||
search_tool.include_answer,
|
||||
search_tool.include_raw_content,
|
||||
search_tool.include_images,
|
||||
search_tool.include_image_descriptions,
|
||||
)
|
||||
|
||||
mock_api_wrapper.clean_results_with_images.assert_called_once_with(
|
||||
sample_raw_results
|
||||
)
|
||||
|
||||
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")
|
||||
|
||||
result, raw = search_tool._run("test query")
|
||||
|
||||
result_dict = json.loads(result)
|
||||
assert "error" in result_dict
|
||||
assert "API Error" in result_dict["error"]
|
||||
assert raw == {}
|
||||
mock_api_wrapper.clean_results_with_images.assert_not_called()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_arun_success(
|
||||
self,
|
||||
search_tool,
|
||||
mock_api_wrapper,
|
||||
sample_raw_results,
|
||||
sample_cleaned_results,
|
||||
):
|
||||
"""Test successful asynchronous run."""
|
||||
mock_api_wrapper.raw_results_async = AsyncMock(return_value=sample_raw_results)
|
||||
mock_api_wrapper.clean_results_with_images.return_value = sample_cleaned_results
|
||||
|
||||
result, raw = await search_tool._arun("test query")
|
||||
|
||||
assert result == json.dumps(sample_cleaned_results, ensure_ascii=False)
|
||||
assert raw == sample_raw_results
|
||||
|
||||
mock_api_wrapper.raw_results_async.assert_called_once_with(
|
||||
"test query",
|
||||
search_tool.max_results,
|
||||
search_tool.search_depth,
|
||||
search_tool.include_domains,
|
||||
search_tool.exclude_domains,
|
||||
search_tool.include_answer,
|
||||
search_tool.include_raw_content,
|
||||
search_tool.include_images,
|
||||
search_tool.include_image_descriptions,
|
||||
)
|
||||
|
||||
mock_api_wrapper.clean_results_with_images.assert_called_once_with(
|
||||
sample_raw_results
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
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")
|
||||
)
|
||||
|
||||
result, raw = await search_tool._arun("test query")
|
||||
|
||||
result_dict = json.loads(result)
|
||||
assert "error" in result_dict
|
||||
assert "Async API Error" in result_dict["error"]
|
||||
assert raw == {}
|
||||
mock_api_wrapper.clean_results_with_images.assert_not_called()
|
||||
|
||||
def test_run_with_run_manager(
|
||||
self,
|
||||
search_tool,
|
||||
mock_api_wrapper,
|
||||
sample_raw_results,
|
||||
sample_cleaned_results,
|
||||
):
|
||||
"""Test run with callback manager."""
|
||||
mock_run_manager = Mock()
|
||||
mock_api_wrapper.raw_results.return_value = sample_raw_results
|
||||
mock_api_wrapper.clean_results_with_images.return_value = sample_cleaned_results
|
||||
|
||||
result, raw = search_tool._run("test query", run_manager=mock_run_manager)
|
||||
|
||||
assert result == json.dumps(sample_cleaned_results, ensure_ascii=False)
|
||||
assert raw == sample_raw_results
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_arun_with_run_manager(
|
||||
self,
|
||||
search_tool,
|
||||
mock_api_wrapper,
|
||||
sample_raw_results,
|
||||
sample_cleaned_results,
|
||||
):
|
||||
"""Test async run with callback manager."""
|
||||
mock_run_manager = Mock()
|
||||
mock_api_wrapper.raw_results_async = AsyncMock(return_value=sample_raw_results)
|
||||
mock_api_wrapper.clean_results_with_images.return_value = sample_cleaned_results
|
||||
|
||||
result, raw = await search_tool._arun(
|
||||
"test query", run_manager=mock_run_manager
|
||||
)
|
||||
|
||||
assert result == json.dumps(sample_cleaned_results, ensure_ascii=False)
|
||||
assert raw == sample_raw_results
|
||||
@@ -0,0 +1,124 @@
|
||||
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
from unittest.mock import Mock, patch
|
||||
|
||||
import pytest
|
||||
from langchain_core.callbacks import (
|
||||
AsyncCallbackManagerForToolRun,
|
||||
CallbackManagerForToolRun,
|
||||
)
|
||||
|
||||
from src.rag import Chunk, Document, Resource, Retriever
|
||||
from src.tools.retriever import RetrieverInput, RetrieverTool, get_retriever_tool
|
||||
|
||||
|
||||
def test_retriever_input_model():
|
||||
input_data = RetrieverInput(keywords="test keywords")
|
||||
assert input_data.keywords == "test keywords"
|
||||
|
||||
|
||||
def test_retriever_tool_init():
|
||||
mock_retriever = Mock(spec=Retriever)
|
||||
resources = [Resource(uri="test://uri", title="Test")]
|
||||
tool = RetrieverTool(retriever=mock_retriever, resources=resources)
|
||||
|
||||
assert tool.name == "local_search_tool"
|
||||
assert "retrieving information" in tool.description
|
||||
assert tool.args_schema == RetrieverInput
|
||||
assert tool.retriever == mock_retriever
|
||||
assert tool.resources == resources
|
||||
|
||||
|
||||
def test_retriever_tool_run_with_results():
|
||||
mock_retriever = Mock(spec=Retriever)
|
||||
chunk = Chunk(content="test content", similarity=0.9)
|
||||
doc = Document(id="doc1", chunks=[chunk])
|
||||
mock_retriever.query_relevant_documents.return_value = [doc]
|
||||
|
||||
resources = [Resource(uri="test://uri", title="Test")]
|
||||
tool = RetrieverTool(retriever=mock_retriever, resources=resources)
|
||||
|
||||
result = tool._run("test keywords")
|
||||
|
||||
mock_retriever.query_relevant_documents.assert_called_once_with(
|
||||
"test keywords", resources
|
||||
)
|
||||
assert isinstance(result, list)
|
||||
assert len(result) == 1
|
||||
assert result[0] == doc.to_dict()
|
||||
|
||||
|
||||
def test_retriever_tool_run_no_results():
|
||||
mock_retriever = Mock(spec=Retriever)
|
||||
mock_retriever.query_relevant_documents.return_value = []
|
||||
|
||||
resources = [Resource(uri="test://uri", title="Test")]
|
||||
tool = RetrieverTool(retriever=mock_retriever, resources=resources)
|
||||
|
||||
result = tool._run("test keywords")
|
||||
|
||||
assert result == "No results found from the local knowledge base."
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_retriever_tool_arun():
|
||||
mock_retriever = Mock(spec=Retriever)
|
||||
chunk = Chunk(content="async content", similarity=0.8)
|
||||
doc = Document(id="doc2", chunks=[chunk])
|
||||
mock_retriever.query_relevant_documents.return_value = [doc]
|
||||
|
||||
resources = [Resource(uri="test://uri", title="Test")]
|
||||
tool = RetrieverTool(retriever=mock_retriever, resources=resources)
|
||||
|
||||
mock_run_manager = Mock(spec=AsyncCallbackManagerForToolRun)
|
||||
mock_sync_manager = Mock(spec=CallbackManagerForToolRun)
|
||||
mock_run_manager.get_sync.return_value = mock_sync_manager
|
||||
|
||||
result = await tool._arun("async keywords", mock_run_manager)
|
||||
|
||||
mock_run_manager.get_sync.assert_called_once()
|
||||
assert isinstance(result, list)
|
||||
assert len(result) == 1
|
||||
assert result[0] == doc.to_dict()
|
||||
|
||||
|
||||
@patch("src.tools.retriever.build_retriever")
|
||||
def test_get_retriever_tool_success(mock_build_retriever):
|
||||
mock_retriever = Mock(spec=Retriever)
|
||||
mock_build_retriever.return_value = mock_retriever
|
||||
|
||||
resources = [Resource(uri="test://uri", title="Test")]
|
||||
tool = get_retriever_tool(resources)
|
||||
|
||||
assert isinstance(tool, RetrieverTool)
|
||||
assert tool.retriever == mock_retriever
|
||||
assert tool.resources == resources
|
||||
|
||||
|
||||
def test_get_retriever_tool_empty_resources():
|
||||
result = get_retriever_tool([])
|
||||
assert result is None
|
||||
|
||||
|
||||
@patch("src.tools.retriever.build_retriever")
|
||||
def test_get_retriever_tool_no_retriever(mock_build_retriever):
|
||||
mock_build_retriever.return_value = None
|
||||
|
||||
resources = [Resource(uri="test://uri", title="Test")]
|
||||
result = get_retriever_tool(resources)
|
||||
|
||||
assert result is None
|
||||
|
||||
|
||||
def test_retriever_tool_run_with_callback_manager():
|
||||
mock_retriever = Mock(spec=Retriever)
|
||||
mock_retriever.query_relevant_documents.return_value = []
|
||||
|
||||
resources = [Resource(uri="test://uri", title="Test")]
|
||||
tool = RetrieverTool(retriever=mock_retriever, resources=resources)
|
||||
|
||||
mock_callback_manager = Mock(spec=CallbackManagerForToolRun)
|
||||
result = tool._run("test keywords", mock_callback_manager)
|
||||
|
||||
assert result == "No results found from the local knowledge base."
|
||||
Reference in New Issue
Block a user