test: fix unit test error

Change-Id: I3dd7a6179132e5497a30ada443d88de0c47af3d4

.gitignore

@@ -11,7 +11,6 @@ static/browser_history/*.gif
 
 # Virtual environments
 .venv
-venv/
 
 # Environment variables
 .env

server.py

@@ -7,8 +7,7 @@ Server script for running the DeerFlow API.
 
 import argparse
 import logging
-import signal
+
-import sys
 import uvicorn
 
 # Configure logging
@@ -19,17 +18,6 @@ logging.basicConfig(
 
 logger = logging.getLogger(__name__)
 
-
-def handle_shutdown(signum, frame):
-    """Handle graceful shutdown on SIGTERM/SIGINT"""
-    logger.info("Received shutdown signal. Starting graceful shutdown...")
-    sys.exit(0)
-
-
-# Register signal handlers
-signal.signal(signal.SIGTERM, handle_shutdown)
-signal.signal(signal.SIGINT, handle_shutdown)
-
 if __name__ == "__main__":
     # Parse command line arguments
     parser = argparse.ArgumentParser(description="Run the DeerFlow API server")
@@ -62,18 +50,16 @@ if __name__ == "__main__":
 
     # Determine reload setting
     reload = False
+
+    # Command line arguments override defaults
     if args.reload:
         reload = True
 
-    try:
+    logger.info("Starting DeerFlow API server")
-        logger.info(f"Starting DeerFlow API server on {args.host}:{args.port}")
+    uvicorn.run(
-        uvicorn.run(
+        "src.server:app",
-            "src.server:app",
+        host=args.host,
-            host=args.host,
+        port=args.port,
-            port=args.port,
+        reload=reload,
-            reload=reload,
+        log_level=args.log_level,
-            log_level=args.log_level,
+    )
-        )
-    except Exception as e:
-        logger.error(f"Failed to start server: {str(e)}")
-        sys.exit(1)

src/graph/nodes.py

@@ -50,10 +50,10 @@ def background_investigation_node(
     logger.info("background investigation node is running.")
     configurable = Configuration.from_runnable_config(config)
     query = state["messages"][-1].content
-    if SELECTED_SEARCH_ENGINE == SearchEngine.TAVILY.value:
+    if SELECTED_SEARCH_ENGINE == SearchEngine.TAVILY:
         searched_content = LoggedTavilySearch(
             max_results=configurable.max_search_results
-        ).invoke(query)
+        ).invoke({"query": query})
         background_investigation_results = None
         if isinstance(searched_content, list):
             background_investigation_results = [

tests/integration/test_nodes.py

@@ -68,7 +68,7 @@ def mock_web_search_tool():
         yield mock
 
 
-@pytest.mark.parametrize("search_engine", [SearchEngine.TAVILY.value, "other"])
+@pytest.mark.parametrize("search_engine", [SearchEngine.TAVILY, "other"])
 def test_background_investigation_node_tavily(
     mock_state,
     mock_tavily_search,
@@ -93,8 +93,10 @@ def test_background_investigation_node_tavily(
         results = json.loads(update["background_investigation_results"])
         assert isinstance(results, list)
 
-        if search_engine == SearchEngine.TAVILY.value:
+        if search_engine == SearchEngine.TAVILY:
-            mock_tavily_search.return_value.invoke.assert_called_once_with("test query")
+            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"
@@ -109,7 +111,7 @@ 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.value):
+    with patch("src.graph.nodes.SELECTED_SEARCH_ENGINE", SearchEngine.TAVILY):
         # Mock a malformed response
         mock_tavily_search.return_value.invoke.return_value = "invalid response"