From b300ca82bc9b4f42552f9f91e0eadc9ea59ef877 Mon Sep 17 00:00:00 2001
From: "codebeaver-ai[bot]"
 <192081515+codebeaver-ai[bot]@users.noreply.github.com>
Date: Mon, 27 Jan 2025 10:57:28 +0000
Subject: [PATCH] test: Add coverage improvement test for
 tests/test_search_graph.py

---
 tests/test_search_graph.py | 36 ++++++++++++++++++++++++++++++++++++
 1 file changed, 36 insertions(+)
 create mode 100644 tests/test_search_graph.py

diff --git a/tests/test_search_graph.py b/tests/test_search_graph.py
new file mode 100644
index 00000000..1384988d
--- /dev/null
+++ b/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
\ No newline at end of file