From ba58568b8a7f3fba634069cf777474d2955475bc Mon Sep 17 00:00:00 2001
From: "codebeaver-ai[bot]"
 <192081515+codebeaver-ai[bot]@users.noreply.github.com>
Date: Wed, 29 Jan 2025 09:34:27 +0000
Subject: [PATCH] test: Update coverage improvement test for
 tests/test_search_graph.py

---
 tests/test_search_graph.py | 26 ++++++++++++++++++++++++--
 1 file changed, 24 insertions(+), 2 deletions(-)

diff --git a/tests/test_search_graph.py b/tests/test_search_graph.py
index f8e51c9c..0b8209c0 100644
--- a/tests/test_search_graph.py
+++ b/tests/test_search_graph.py
@@ -1,7 +1,7 @@
 import pytest
 
 from scrapegraphai.graphs.search_graph import SearchGraph
-from unittest.mock import MagicMock, patch
+from unittest.mock import MagicMock, call, patch
 
 class TestSearchGraph:
     """Test class for SearchGraph"""
@@ -57,4 +57,26 @@ class TestSearchGraph:
         result = search_graph.run()
 
         # Assert
-        assert result == "No answer found."
\ No newline at end of file
+        assert result == "No answer found."
+
+    @patch('scrapegraphai.graphs.search_graph.SearchInternetNode')
+    @patch('scrapegraphai.graphs.search_graph.GraphIteratorNode')
+    @patch('scrapegraphai.graphs.search_graph.MergeAnswersNode')
+    @patch('scrapegraphai.graphs.search_graph.BaseGraph')
+    @patch('scrapegraphai.graphs.abstract_graph.AbstractGraph._create_llm')
+    def test_max_results_config(self, mock_create_llm, mock_base_graph, mock_merge_answers, mock_graph_iterator, mock_search_internet):
+        """
+        Test that the max_results parameter from the config is correctly passed to the SearchInternetNode.
+        """
+        # Arrange
+        prompt = "Test prompt"
+        max_results = 5
+        config = {"llm": {"model": "test-model"}, "max_results": max_results}
+
+        # Act
+        search_graph = SearchGraph(prompt, config)
+
+        # Assert
+        mock_search_internet.assert_called_once()
+        call_args = mock_search_internet.call_args
+        assert call_args.kwargs['node_config']['max_results'] == max_results
\ No newline at end of file