add smart_scraper_openai_test

Aliang-CN · Jun 17, 2024 · 6a753f2 · 6a753f2
1 parent 4c8becc
commit 6a753f2
Show file tree

Hide file tree

Showing 2 changed files with 53 additions and 0 deletions.
diff --git a/tests/graphs/.env.example b/tests/graphs/.env.example
@@ -0,0 +1 @@
+OPENAI_API_KEY="YOUR OPENAI API KEY"
diff --git a/tests/graphs/smart_scraper_openai_test.py b/tests/graphs/smart_scraper_openai_test.py
@@ -0,0 +1,52 @@
+"""
+Module for testing the smart scraper class
+"""
+
+import os
+import pytest
+import pandas as pd
+from dotenv import load_dotenv
+from scrapegraphai.graphs import SmartScraperGraph
+from scrapegraphai.utils import prettify_exec_info
+
+load_dotenv()
+
+@pytest.fixture
+def graph_config():
+    """Configuration of the graph"""
+    openai_key = os.getenv("OPENAI_APIKEY")
+    return {
+        "llm": {
+            "api_key": openai_key,
+            "model": "gpt-3.5-turbo",
+        },
+        "verbose": True,
+        "headless": False,
+    }
+
+def test_scraping_pipeline(graph_config):
+    """Start of the scraping pipeline"""
+    smart_scraper_graph = SmartScraperGraph(
+        prompt="List me all the projects with their description.",
+        source="https://perinim.github.io/projects/",
+        config=graph_config,
+    )
+
+    result = smart_scraper_graph.run()
+
+    assert result is not None
+    assert isinstance(result, dict) 
+
+def test_get_execution_info(graph_config):
+    """Get the execution info"""
+    smart_scraper_graph = SmartScraperGraph(
+        prompt="List me all the projects with their description.",
+        source="https://perinim.github.io/projects/",
+        config=graph_config,
+    )
+
+    smart_scraper_graph.run()
+
+    graph_exec_info = smart_scraper_graph.get_execution_info()
+
+    assert graph_exec_info is not None