Skip to content

Commit

Permalink
test: Update coverage improvement test for tests/test_search_graph.py
Browse files Browse the repository at this point in the history
  • Loading branch information
codebeaver-ai[bot] authored Jan 29, 2025
1 parent c9d71af commit 80dd766
Showing 1 changed file with 26 additions and 1 deletion.
27 changes: 26 additions & 1 deletion tests/test_search_graph.py
Original file line number Diff line number Diff line change
Expand Up @@ -79,4 +79,29 @@ def test_max_results_config(self, mock_create_llm, mock_base_graph, mock_merge_a
# Assert
mock_search_internet.assert_called_once()
call_args = mock_search_internet.call_args
assert call_args.kwargs['node_config']['max_results'] == max_results
assert call_args.kwargs['node_config']['max_results'] == max_results

@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_custom_search_engine_config(self, mock_create_llm, mock_base_graph, mock_merge_answers, mock_graph_iterator, mock_search_internet):
"""
Test that the custom search_engine parameter from the config is correctly passed to the SearchInternetNode.
"""
# Arrange
prompt = "Test prompt"
custom_search_engine = "custom_engine"
config = {
"llm": {"model": "test-model"},
"search_engine": custom_search_engine
}

# 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']['search_engine'] == custom_search_engine

0 comments on commit 80dd766

Please sign in to comment.