|
| 1 | +from unittest.mock import patch, MagicMock |
| 2 | +from scrapegraphai.graphs.depth_search_graph import DepthSearchGraph |
| 3 | +from scrapegraphai.graphs.abstract_graph import AbstractGraph |
| 4 | +import pytest |
| 5 | + |
| 6 | + |
| 7 | +class TestDepthSearchGraph: |
| 8 | + """Test suite for DepthSearchGraph class""" |
| 9 | + |
| 10 | + @pytest.mark.parametrize( |
| 11 | + "source, expected_input_key", |
| 12 | + [ |
| 13 | + ("https://example.com", "url"), |
| 14 | + ("/path/to/local/directory", "local_dir"), |
| 15 | + ], |
| 16 | + ) |
| 17 | + def test_depth_search_graph_initialization(self, source, expected_input_key): |
| 18 | + """ |
| 19 | + Test that DepthSearchGraph initializes correctly with different source types. |
| 20 | + This test verifies that the input_key is set to 'url' for web sources and |
| 21 | + 'local_dir' for local directory sources. |
| 22 | + """ |
| 23 | + prompt = "Test prompt" |
| 24 | + config = {"llm": {"model": "mock_model"}} |
| 25 | + |
| 26 | + # Mock both BaseGraph and _create_llm method |
| 27 | + with patch("scrapegraphai.graphs.depth_search_graph.BaseGraph"), \ |
| 28 | + patch.object(AbstractGraph, '_create_llm', return_value=MagicMock()): |
| 29 | + graph = DepthSearchGraph(prompt, source, config) |
| 30 | + |
| 31 | + assert graph.prompt == prompt |
| 32 | + assert graph.source == source |
| 33 | + assert graph.config == config |
| 34 | + assert graph.input_key == expected_input_key |
0 commit comments