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