mirror of
https://github.com/VinciGit00/Scrapegraph-ai.git
synced 2026-07-12 21:01:56 +08:00
dedfa2eaf0
- Added pytest fixture to set up the RobotsNode with the initial state. - Implemented test_robots_node to test the execution of RobotsNode. - Used unittest.mock.patch to mock the execute method, ensuring faster and more reliable tests without actual network calls. - Added assertions to verify the correctness of the result and ensure the execute method is called once with the correct arguments.
56 lines
1.6 KiB
Python
56 lines
1.6 KiB
Python
import pytest
|
|
from scrapegraphai.models import Ollama
|
|
from scrapegraphai.nodes import RobotsNode
|
|
from unittest.mock import patch, MagicMock
|
|
|
|
@pytest.fixture
|
|
def setup():
|
|
"""
|
|
Setup the RobotsNode and initial state for testing.
|
|
"""
|
|
# Define the configuration for the graph
|
|
graph_config = {
|
|
"llm": {
|
|
"model_name": "ollama/llama3",
|
|
"temperature": 0,
|
|
"streaming": True
|
|
},
|
|
}
|
|
|
|
# Instantiate the LLM model with the configuration
|
|
llm_model = Ollama(graph_config["llm"])
|
|
|
|
# Define the RobotsNode with necessary configurations
|
|
robots_node = RobotsNode(
|
|
input="url",
|
|
output=["is_scrapable"],
|
|
node_config={
|
|
"llm_model": llm_model,
|
|
"headless": False
|
|
}
|
|
)
|
|
|
|
# Define the initial state for the node
|
|
initial_state = {
|
|
"url": "https://twitter.com/home"
|
|
}
|
|
|
|
return robots_node, initial_state
|
|
|
|
def test_robots_node(setup):
|
|
"""
|
|
Test the RobotsNode execution.
|
|
"""
|
|
robots_node, initial_state = setup
|
|
|
|
# Patch the execute method to avoid actual network calls and return a mock response
|
|
with patch.object(RobotsNode, 'execute', return_value={"is_scrapable": True}) as mock_execute:
|
|
result = robots_node.execute(initial_state)
|
|
|
|
# Check if the result is not None
|
|
assert result is not None
|
|
# Additional assertion to check the returned value
|
|
assert result["is_scrapable"] is True
|
|
# Ensure the execute method was called once
|
|
mock_execute.assert_called_once_with(initial_state)
|