From 0d9995b297c4bd19b6c915facc6c72199854aeb6 Mon Sep 17 00:00:00 2001
From: "codebeaver-ai[bot]"
 <192081515+codebeaver-ai[bot]@users.noreply.github.com>
Date: Wed, 29 Jan 2025 10:25:49 +0000
Subject: [PATCH] test: Add coverage improvement test for
 tests/test_depth_search_graph.py

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

diff --git a/tests/test_depth_search_graph.py b/tests/test_depth_search_graph.py
new file mode 100644
index 00000000..0197a6b8
--- /dev/null
+++ b/tests/test_depth_search_graph.py
@@ -0,0 +1,34 @@
+from unittest.mock import patch, MagicMock
+from scrapegraphai.graphs.depth_search_graph import DepthSearchGraph
+from scrapegraphai.graphs.abstract_graph import AbstractGraph
+import pytest
+
+
+class TestDepthSearchGraph:
+    """Test suite for DepthSearchGraph class"""
+
+    @pytest.mark.parametrize(
+        "source, expected_input_key",
+        [
+            ("https://example.com", "url"),
+            ("/path/to/local/directory", "local_dir"),
+        ],
+    )
+    def test_depth_search_graph_initialization(self, source, expected_input_key):
+        """
+        Test that DepthSearchGraph initializes correctly with different source types.
+        This test verifies that the input_key is set to 'url' for web sources and
+        'local_dir' for local directory sources.
+        """
+        prompt = "Test prompt"
+        config = {"llm": {"model": "mock_model"}}
+        
+        # Mock both BaseGraph and _create_llm method
+        with patch("scrapegraphai.graphs.depth_search_graph.BaseGraph"), \
+             patch.object(AbstractGraph, '_create_llm', return_value=MagicMock()):
+            graph = DepthSearchGraph(prompt, source, config)
+            
+            assert graph.prompt == prompt
+            assert graph.source == source
+            assert graph.config == config
+            assert graph.input_key == expected_input_key