Skip to content

Commit 6a753f2

Browse files
committed
add smart_scraper_openai_test
1 parent 4c8becc commit 6a753f2

File tree

2 files changed

+53
-0
lines changed

2 files changed

+53
-0
lines changed

tests/graphs/.env.example

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
OPENAI_API_KEY="YOUR OPENAI API KEY"
Lines changed: 52 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,52 @@
1+
"""
2+
Module for testing the smart scraper class
3+
"""
4+
5+
import os
6+
import pytest
7+
import pandas as pd
8+
from dotenv import load_dotenv
9+
from scrapegraphai.graphs import SmartScraperGraph
10+
from scrapegraphai.utils import prettify_exec_info
11+
12+
load_dotenv()
13+
14+
@pytest.fixture
15+
def graph_config():
16+
"""Configuration of the graph"""
17+
openai_key = os.getenv("OPENAI_APIKEY")
18+
return {
19+
"llm": {
20+
"api_key": openai_key,
21+
"model": "gpt-3.5-turbo",
22+
},
23+
"verbose": True,
24+
"headless": False,
25+
}
26+
27+
def test_scraping_pipeline(graph_config):
28+
"""Start of the scraping pipeline"""
29+
smart_scraper_graph = SmartScraperGraph(
30+
prompt="List me all the projects with their description.",
31+
source="https://perinim.github.io/projects/",
32+
config=graph_config,
33+
)
34+
35+
result = smart_scraper_graph.run()
36+
37+
assert result is not None
38+
assert isinstance(result, dict)
39+
40+
def test_get_execution_info(graph_config):
41+
"""Get the execution info"""
42+
smart_scraper_graph = SmartScraperGraph(
43+
prompt="List me all the projects with their description.",
44+
source="https://perinim.github.io/projects/",
45+
config=graph_config,
46+
)
47+
48+
smart_scraper_graph.run()
49+
50+
graph_exec_info = smart_scraper_graph.get_execution_info()
51+
52+
assert graph_exec_info is not None

0 commit comments

Comments
 (0)