diff --git a/tests/test_json_scraper_graph.py b/tests/test_json_scraper_graph.py index 0f00691f..1572650e 100644 --- a/tests/test_json_scraper_graph.py +++ b/tests/test_json_scraper_graph.py @@ -95,4 +95,42 @@ class TestJSONScraperGraph: mock_execute.assert_called_once_with({"user_prompt": "Analyze the data from the JSON file", "json": "path/to/single/file.json"}) mock_fetch_node.assert_called_once() mock_generate_answer_node.assert_called_once() + mock_create_llm.assert_called_once_with({"model": "test-model", "temperature": 0}) + + @patch('scrapegraphai.graphs.json_scraper_graph.FetchNode') + @patch('scrapegraphai.graphs.json_scraper_graph.GenerateAnswerNode') + @patch.object(JSONScraperGraph, '_create_llm') + def test_json_scraper_graph_no_answer_found(self, mock_create_llm, mock_generate_answer_node, mock_fetch_node, mock_llm_model, mock_embedder_model): + """ + Test JSONScraperGraph when no answer is found. + This test checks if the graph correctly handles the scenario where no answer is generated, + ensuring it returns the default "No answer found." message. + """ + # Mock the _create_llm method to return a mock LLM model + mock_create_llm.return_value = mock_llm_model + + # Mock the execute method of BaseGraph to return an empty answer + with patch('scrapegraphai.graphs.json_scraper_graph.BaseGraph.execute') as mock_execute: + mock_execute.return_value = ({}, {}) # Empty state and execution info + + # Create a JSONScraperGraph instance + graph = JSONScraperGraph( + prompt="Query that produces no answer", + source="path/to/empty/file.json", + config={"llm": {"model": "test-model", "temperature": 0}}, + schema=BaseModel + ) + + # Set mocked embedder model + graph.embedder_model = mock_embedder_model + + # Run the graph + result = graph.run() + + # Assertions + assert result == "No answer found." + assert graph.input_key == "json" + mock_execute.assert_called_once_with({"user_prompt": "Query that produces no answer", "json": "path/to/empty/file.json"}) + mock_fetch_node.assert_called_once() + mock_generate_answer_node.assert_called_once() mock_create_llm.assert_called_once_with({"model": "test-model", "temperature": 0}) \ No newline at end of file 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