Skip to content

Commit ba58568

Browse files
test: Update coverage improvement test for tests/test_search_graph.py
1 parent f5ebe8a commit ba58568

File tree

1 file changed

+24
-2
lines changed

1 file changed

+24
-2
lines changed

tests/test_search_graph.py

Lines changed: 24 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
import pytest
22

33
from scrapegraphai.graphs.search_graph import SearchGraph
4-
from unittest.mock import MagicMock, patch
4+
from unittest.mock import MagicMock, call, patch
55

66
class TestSearchGraph:
77
"""Test class for SearchGraph"""
@@ -57,4 +57,26 @@ def test_run_no_answer_found(self, mock_create_llm, mock_base_graph):
5757
result = search_graph.run()
5858

5959
# Assert
60-
assert result == "No answer found."
60+
assert result == "No answer found."
61+
62+
@patch('scrapegraphai.graphs.search_graph.SearchInternetNode')
63+
@patch('scrapegraphai.graphs.search_graph.GraphIteratorNode')
64+
@patch('scrapegraphai.graphs.search_graph.MergeAnswersNode')
65+
@patch('scrapegraphai.graphs.search_graph.BaseGraph')
66+
@patch('scrapegraphai.graphs.abstract_graph.AbstractGraph._create_llm')
67+
def test_max_results_config(self, mock_create_llm, mock_base_graph, mock_merge_answers, mock_graph_iterator, mock_search_internet):
68+
"""
69+
Test that the max_results parameter from the config is correctly passed to the SearchInternetNode.
70+
"""
71+
# Arrange
72+
prompt = "Test prompt"
73+
max_results = 5
74+
config = {"llm": {"model": "test-model"}, "max_results": max_results}
75+
76+
# Act
77+
search_graph = SearchGraph(prompt, config)
78+
79+
# Assert
80+
mock_search_internet.assert_called_once()
81+
call_args = mock_search_internet.call_args
82+
assert call_args.kwargs['node_config']['max_results'] == max_results

0 commit comments

Comments
 (0)