Merge pull request #911 from ScrapeGraphAI/codebeaver/pre/beta-904

codebeaver/pre/beta-904 - Unit Tests

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})

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."
+        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