mirror of
https://gitee.com/wanwujie/deer-flow
synced 2026-04-03 06:12:14 +08:00
test: add background node unit test (#198)
* test: add background node unit test Change-Id: Ia99f5a1687464387dcb01bbee04deaa371c6e490 * test: add background node unit test Change-Id: I9aabcf02ff04fda40c56f3ea22abe6b8f93bf9b6 * test: fix test error Change-Id: I3997dc53a2cfaa35501a1fbda5902ee15528124e * test: fix unit test error Change-Id: If4c4cd10673e76a30945674c7cda198aeabf28d0 * test: fix unit test error Change-Id: I3dd7a6179132e5497a30ada443d88de0c47af3d4
This commit is contained in:
@@ -44,13 +44,12 @@ def get_llm_by_type(
|
||||
return llm
|
||||
|
||||
|
||||
# Initialize LLMs for different purposes - now these will be cached
|
||||
basic_llm = get_llm_by_type("basic")
|
||||
|
||||
# In the future, we will use reasoning_llm and vl_llm for different purposes
|
||||
# reasoning_llm = get_llm_by_type("reasoning")
|
||||
# vl_llm = get_llm_by_type("vision")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Initialize LLMs for different purposes - now these will be cached
|
||||
basic_llm = get_llm_by_type("basic")
|
||||
print(basic_llm.invoke("Hello"))
|
||||
|
||||
130
tests/integration/test_nodes.py
Normal file
130
tests/integration/test_nodes.py
Normal file
@@ -0,0 +1,130 @@
|
||||
import json
|
||||
import pytest
|
||||
from unittest.mock import patch, MagicMock
|
||||
|
||||
# 在这里 mock 掉 get_llm_by_type,避免 ValueError
|
||||
with patch("src.llms.llm.get_llm_by_type", return_value=MagicMock()):
|
||||
from langgraph.types import Command
|
||||
from src.graph.nodes import background_investigation_node
|
||||
from src.config import SearchEngine
|
||||
from langchain_core.messages import HumanMessage
|
||||
|
||||
# Mock data
|
||||
MOCK_SEARCH_RESULTS = [
|
||||
{"title": "Test Title 1", "content": "Test Content 1"},
|
||||
{"title": "Test Title 2", "content": "Test Content 2"},
|
||||
]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_state():
|
||||
return {
|
||||
"messages": [HumanMessage(content="test query")],
|
||||
"background_investigation_results": None,
|
||||
}
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_configurable():
|
||||
mock = MagicMock()
|
||||
mock.max_search_results = 5
|
||||
return mock
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_config():
|
||||
# 你可以根据实际需要返回一个 MagicMock 或 dict
|
||||
return MagicMock()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def patch_config_from_runnable_config(mock_configurable):
|
||||
with patch(
|
||||
"src.graph.nodes.Configuration.from_runnable_config",
|
||||
return_value=mock_configurable,
|
||||
):
|
||||
yield
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_tavily_search():
|
||||
with patch("src.graph.nodes.LoggedTavilySearch") as mock:
|
||||
instance = mock.return_value
|
||||
instance.invoke.return_value = [
|
||||
{"title": "Test Title 1", "content": "Test Content 1"},
|
||||
{"title": "Test Title 2", "content": "Test Content 2"},
|
||||
]
|
||||
yield mock
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_web_search_tool():
|
||||
with patch("src.graph.nodes.get_web_search_tool") as mock:
|
||||
instance = mock.return_value
|
||||
instance.invoke.return_value = [
|
||||
{"title": "Test Title 1", "content": "Test Content 1"},
|
||||
{"title": "Test Title 2", "content": "Test Content 2"},
|
||||
]
|
||||
yield mock
|
||||
|
||||
|
||||
@pytest.mark.parametrize("search_engine", [SearchEngine.TAVILY, "other"])
|
||||
def test_background_investigation_node_tavily(
|
||||
mock_state,
|
||||
mock_tavily_search,
|
||||
mock_web_search_tool,
|
||||
search_engine,
|
||||
patch_config_from_runnable_config,
|
||||
mock_config,
|
||||
):
|
||||
"""Test background_investigation_node with Tavily search engine"""
|
||||
with patch("src.graph.nodes.SELECTED_SEARCH_ENGINE", search_engine):
|
||||
result = background_investigation_node(mock_state, mock_config)
|
||||
|
||||
# Verify the result structure
|
||||
assert isinstance(result, Command)
|
||||
assert result.goto == "planner"
|
||||
|
||||
# Verify the update contains background_investigation_results
|
||||
update = result.update
|
||||
assert "background_investigation_results" in update
|
||||
|
||||
# Parse and verify the JSON content
|
||||
results = json.loads(update["background_investigation_results"])
|
||||
assert isinstance(results, list)
|
||||
|
||||
if search_engine == SearchEngine.TAVILY:
|
||||
mock_tavily_search.return_value.invoke.assert_called_once_with(
|
||||
{"query": "test query"}
|
||||
)
|
||||
assert len(results) == 2
|
||||
assert results[0]["title"] == "Test Title 1"
|
||||
assert results[0]["content"] == "Test Content 1"
|
||||
else:
|
||||
mock_web_search_tool.return_value.invoke.assert_called_once_with(
|
||||
"test query"
|
||||
)
|
||||
assert len(results) == 2
|
||||
|
||||
|
||||
def test_background_investigation_node_malformed_response(
|
||||
mock_state, mock_tavily_search, patch_config_from_runnable_config, mock_config
|
||||
):
|
||||
"""Test background_investigation_node with malformed Tavily response"""
|
||||
with patch("src.graph.nodes.SELECTED_SEARCH_ENGINE", SearchEngine.TAVILY):
|
||||
# Mock a malformed response
|
||||
mock_tavily_search.return_value.invoke.return_value = "invalid response"
|
||||
|
||||
result = background_investigation_node(mock_state, mock_config)
|
||||
|
||||
# Verify the result structure
|
||||
assert isinstance(result, Command)
|
||||
assert result.goto == "planner"
|
||||
|
||||
# Verify the update contains background_investigation_results
|
||||
update = result.update
|
||||
assert "background_investigation_results" in update
|
||||
|
||||
# Parse and verify the JSON content
|
||||
results = json.loads(update["background_investigation_results"])
|
||||
assert results is None
|
||||
Reference in New Issue
Block a user