mirror of
https://github.com/VinciGit00/Scrapegraph-ai.git
synced 2026-07-09 21:19:20 +08:00
test: Add coverage improvement test for tests/test_search_graph.py
This commit is contained in:
parent
d5c7c9cd9d
commit
b300ca82bc
36
tests/test_search_graph.py
Normal file
36
tests/test_search_graph.py
Normal file
@ -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
|
||||
Loading…
Reference in New Issue
Block a user