Merge pull request #908 from ScrapeGraphAI/codebeaver/pre/beta

codebeaver/pre/beta - Unit Tests

codebeaver.yml

@@ -0,0 +1,2 @@
+from: pytest
+setup_commands: ['@merge', 'pip install -q selenium', 'pip install -q playwright', 'playwright install']

tests/test_json_scraper_graph.py

@@ -0,0 +1,52 @@
+import pytest
+
+from pydantic import BaseModel
+from scrapegraphai.graphs.json_scraper_graph import JSONScraperGraph
+from unittest.mock import Mock, patch
+
+class TestJSONScraperGraph:
+    @pytest.fixture
+    def mock_llm_model(self):
+        return Mock()
+
+    @pytest.fixture
+    def mock_embedder_model(self):
+        return Mock()
+
+    @patch('scrapegraphai.graphs.json_scraper_graph.FetchNode')
+    @patch('scrapegraphai.graphs.json_scraper_graph.GenerateAnswerNode')
+    @patch.object(JSONScraperGraph, '_create_llm')
+    def test_json_scraper_graph_with_directory(self, mock_create_llm, mock_generate_answer_node, mock_fetch_node, mock_llm_model, mock_embedder_model):
+        """
+        Test JSONScraperGraph with a directory of JSON files.
+        This test checks if the graph correctly handles multiple JSON files input
+        and processes them to generate an answer.
+        """
+        # 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
+        with patch('scrapegraphai.graphs.json_scraper_graph.BaseGraph.execute') as mock_execute:
+            mock_execute.return_value = ({"answer": "Mocked answer for multiple JSON files"}, {})
+
+            # Create a JSONScraperGraph instance
+            graph = JSONScraperGraph(
+                prompt="Summarize the data from all JSON files",
+                source="path/to/json/directory",
+                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 == "Mocked answer for multiple JSON files"
+            assert graph.input_key == "json_dir"
+            mock_execute.assert_called_once_with({"user_prompt": "Summarize the data from all JSON files", "json_dir": "path/to/json/directory"})
+            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

@@ -0,0 +1,36 @@
+import pytest
+
+from scrapegraphai.graphs.search_graph import SearchGraph
+from unittest.mock import MagicMock, patch
+
+class TestSearchGraph:
+    """Test class for SearchGraph"""
+
+    @pytest.mark.parametrize("urls", [
+        ["https://example.com", "https://test.com"],
+        [],
+        ["https://single-url.com"]
+    ])
+    @patch('scrapegraphai.graphs.search_graph.BaseGraph')
+    @patch('scrapegraphai.graphs.abstract_graph.AbstractGraph._create_llm')
+    def test_get_considered_urls(self, mock_create_llm, mock_base_graph, urls):
+        """
+        Test that get_considered_urls returns the correct list of URLs
+        considered during the search process.
+        """
+        # Arrange
+        prompt = "Test prompt"
+        config = {"llm": {"model": "test-model"}}
+
+        # Mock the _create_llm method to return a MagicMock
+        mock_create_llm.return_value = MagicMock()
+
+        # Mock the execute method to set the final_state
+        mock_base_graph.return_value.execute.return_value = ({"urls": urls}, {})
+
+        # Act
+        search_graph = SearchGraph(prompt, config)
+        search_graph.run()
+
+        # Assert
+        assert search_graph.get_considered_urls() == urls