Skip to content

Test ScriptCreatorGraph and print execution info #355

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Jun 9, 2024
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
72 changes: 35 additions & 37 deletions tests/graphs/script_generator_test.py
Original file line number Diff line number Diff line change
@@ -1,47 +1,45 @@
"""
"""
Module for making the tests for ScriptGeneratorGraph
"""
import pytest
from scrapegraphai.graphs import ScriptCreatorGraph
from scrapegraphai.utils import prettify_exec_info


@pytest.fixture
def graph_config():
"""
Configuration of the graph
"""
return {
"llm": {
"model": "ollama/mistral",
"temperature": 0,
"format": "json",
"base_url": "http://localhost:11434",
"library": "beautifoulsoup",
},
"embeddings": {
"model": "ollama/nomic-embed-text",
"temperature": 0,
"base_url": "http://localhost:11434",
},
"library": "beautifoulsoup"
}

"""
Configuration of the graph
"""
return {
"llm": {
"model": "ollama/mistral",
"temperature": 0,
"format": "json",
"base_url": "http://localhost:11434",
"library": "beautifulsoup",
},
"embeddings": {
"model": "ollama/nomic-embed-text",
"temperature": 0,
"base_url": "http://localhost:11434",
},
"library": "beautifulsoup"
}

def test_script_creator_graph(graph_config: dict):
"""
Start of the scraping pipeline
"""
smart_scraper_graph = ScriptCreatorGraph(
prompt="List me all the news with their description.",
source="https://perinim.github.io/projects",
config=graph_config
)

result = smart_scraper_graph.run()

assert result is not None

graph_exec_info = smart_scraper_graph.get_execution_info()

assert graph_exec_info is not None
"""
Test the ScriptCreatorGraph
"""
smart_scraper_graph = ScriptCreatorGraph(
prompt="List me all the news with their description.",
source="https://perinim.github.io/projects",
config=graph_config
)
result = smart_scraper_graph.run()
assert result is not None, "ScriptCreatorGraph execution failed to produce a result."
graph_exec_info = smart_scraper_graph.get_execution_info()
assert graph_exec_info is not None, "ScriptCreatorGraph execution info is None."
prettified_exec_info = prettify_exec_info(graph_exec_info)
print(prettified_exec_info)

# Perform additional assertions on the result or execution info as needed
Loading