test: Update coverage improvement test for tests/test_json_scraper_graph.py

This commit is contained in:
codebeaver-ai[bot] 2025-01-28 11:02:15 +00:00 committed by GitHub
parent a406264a12
commit 9919e7c122
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -49,4 +49,50 @@ class TestJSONScraperGraph:
mock_execute.assert_called_once_with({"user_prompt": "Summarize the data from all JSON files", "json_dir": "path/to/json/directory"})
mock_fetch_node.assert_called_once()
mock_generate_answer_node.assert_called_once()
mock_create_llm.assert_called_once_with({"model": "test-model", "temperature": 0})
@pytest.fixture
def mock_llm_model(self):
return Mock()
@pytest.fixture
def mock_embedder_model(self):
return Mock()
@patch('scrapegraphai.graphs.json_scraper_graph.FetchNode')
@patch('scrapegraphai.graphs.json_scraper_graph.GenerateAnswerNode')
@patch.object(JSONScraperGraph, '_create_llm')
def test_json_scraper_graph_with_single_file(self, mock_create_llm, mock_generate_answer_node, mock_fetch_node, mock_llm_model, mock_embedder_model):
"""
Test JSONScraperGraph with a single JSON file.
This test checks if the graph correctly handles a single JSON file input
and processes it to generate an answer.
"""
# Mock the _create_llm method to return a mock LLM model
mock_create_llm.return_value = mock_llm_model
# Mock the execute method of BaseGraph
with patch('scrapegraphai.graphs.json_scraper_graph.BaseGraph.execute') as mock_execute:
mock_execute.return_value = ({"answer": "Mocked answer for single JSON file"}, {})
# Create a JSONScraperGraph instance with a single JSON file
graph = JSONScraperGraph(
prompt="Analyze the data from the JSON file",
source="path/to/single/file.json",
config={"llm": {"model": "test-model", "temperature": 0}},
schema=BaseModel
)
# Set mocked embedder model
graph.embedder_model = mock_embedder_model
# Run the graph
result = graph.run()
# Assertions
assert result == "Mocked answer for single JSON file"
assert graph.input_key == "json"
mock_execute.assert_called_once_with({"user_prompt": "Analyze the data from the JSON file", "json": "path/to/single/file.json"})
mock_fetch_node.assert_called_once()
mock_generate_answer_node.assert_called_once()
mock_create_llm.assert_called_once_with({"model": "test-model", "temperature": 0})