From 55ce399969b8150ceab55c31f6f6b4a7a59ef262 Mon Sep 17 00:00:00 2001
From: laundry <40748509+laundry2@users.noreply.github.com>
Date: Tue, 20 May 2025 14:25:35 +0800
Subject: [PATCH] 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
---
 src/llms/llm.py                 |   5 +-
 tests/integration/test_nodes.py | 130 ++++++++++++++++++++++++++++++++
 2 files changed, 132 insertions(+), 3 deletions(-)
 create mode 100644 tests/integration/test_nodes.py

diff --git a/src/llms/llm.py b/src/llms/llm.py
index 5ac7402..81fb748 100644
--- a/src/llms/llm.py
+++ b/src/llms/llm.py
@@ -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"))
diff --git a/tests/integration/test_nodes.py b/tests/integration/test_nodes.py
new file mode 100644
index 0000000..978acaa
--- /dev/null
+++ b/tests/integration/test_nodes.py
@@ -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