add tests

Author: VinciGit00

tests/graphs/code_generator_graph_openai_test.py

@@ -0,0 +1,79 @@
+"""
+code_generator_graph_openai_test module
+"""
+import os
+from typing import List
+import pytest
+from dotenv import load_dotenv
+from pydantic import BaseModel, Field
+from scrapegraphai.graphs import CodeGeneratorGraph
+
+load_dotenv()
+
+# ************************************************
+# Define the output schema for the graph
+# ************************************************
+
+class Project(BaseModel):
+    title: str = Field(description="The title of the project")
+    description: str = Field(description="The description of the project")
+
+class Projects(BaseModel):
+    projects: List[Project]
+
+@pytest.fixture
+def graph_config():
+    """
+    Configuration for the CodeGeneratorGraph
+    """
+    openai_key = os.getenv("OPENAI_APIKEY")
+    return {
+        "llm": {
+            "api_key": openai_key,
+            "model": "openai/gpt-4o-mini",
+        },
+        "verbose": True,
+        "headless": False,
+        "reduction": 2,
+        "max_iterations": {
+            "overall": 10,
+            "syntax": 3,
+            "execution": 3,
+            "validation": 3,
+            "semantic": 3
+        },
+        "output_file_name": "extracted_data.py"
+    }
+
+def test_code_generator_graph(graph_config: dict):
+    """
+    Test the CodeGeneratorGraph scraping pipeline
+    """
+    code_generator_graph = CodeGeneratorGraph(
+        prompt="List me all the projects with their description",
+        source="https://perinim.github.io/projects/",
+        schema=Projects,
+        config=graph_config
+    )
+
+    result = code_generator_graph.run()
+
+    assert result is not None
+
+
+def test_code_generator_execution_info(graph_config: dict):
+    """
+    Test getting the execution info of CodeGeneratorGraph
+    """
+    code_generator_graph = CodeGeneratorGraph(
+        prompt="List me all the projects with their description",
+        source="https://perinim.github.io/projects/",
+        schema=Projects,
+        config=graph_config
+    )
+
+    code_generator_graph.run()
+
+    graph_exec_info = code_generator_graph.get_execution_info()
+
+    assert graph_exec_info is not None

tests/graphs/depth_search_graph_openai_test.py

@@ -0,0 +1,57 @@
+"""
+depth_search_graph test
+"""
+import os
+import pytest
+from dotenv import load_dotenv
+from scrapegraphai.graphs import DepthSearchGraph
+
+load_dotenv()
+
+@pytest.fixture
+def graph_config():
+    """
+    Configuration for the DepthSearchGraph
+    """
+    openai_key = os.getenv("OPENAI_APIKEY")
+    return {
+        "llm": {
+            "api_key": openai_key,
+            "model": "openai/gpt-4o-mini",
+        },
+        "verbose": True,
+        "headless": False,
+        "depth": 2,
+        "only_inside_links": False,
+    }
+
+def test_depth_search_graph(graph_config: dict):
+    """
+    Test the DepthSearchGraph scraping pipeline
+    """
+    search_graph = DepthSearchGraph(
+        prompt="List me all the projects with their description",
+        source="https://perinim.github.io",
+        config=graph_config
+    )
+
+    result = search_graph.run()
+
+    assert result is not None
+
+
+def test_depth_search_execution_info(graph_config: dict):
+    """
+    Test getting the execution info of DepthSearchGraph
+    """
+    search_graph = DepthSearchGraph(
+        prompt="List me all the projects with their description",
+        source="https://perinim.github.io",
+        config=graph_config
+    )
+
+    search_graph.run()
+
+    graph_exec_info = search_graph.get_execution_info()
+
+    assert graph_exec_info is not None

tests/graphs/search_graph_openai_test.py

@@ -0,0 +1,62 @@
+"""
+search_graph_openai_test.py module
+"""
+import os
+import pytest
+from dotenv import load_dotenv
+from scrapegraphai.graphs import SearchGraph
+
+load_dotenv()
+
+# ************************************************
+# Define the test fixtures and helpers
+# ************************************************
+
+@pytest.fixture
+def graph_config():
+    """
+    Configuration for the SearchGraph
+    """
+    openai_key = os.getenv("OPENAI_APIKEY")
+    return {
+        "llm": {
+            "api_key": openai_key,
+            "model": "openai/gpt-4o",
+        },
+        "max_results": 2,
+        "verbose": True,
+    }
+
+# ************************************************
+# Define the test cases
+# ************************************************
+
+def test_search_graph(graph_config: dict):
+    """
+    Test the SearchGraph functionality
+    """
+    search_graph = SearchGraph(
+        prompt="List me Chioggia's famous dishes",
+        config=graph_config
+    )
+
+    result = search_graph.run()
+
+    assert result is not None
+    assert len(result) > 0
+
+
+def test_search_graph_execution_info(graph_config: dict):
+    """
+    Test getting the execution info of SearchGraph
+    """
+    search_graph = SearchGraph(
+        prompt="List me Chioggia's famous dishes",
+        config=graph_config
+    )
+
+    search_graph.run()
+
+    graph_exec_info = search_graph.get_execution_info()
+
+    assert graph_exec_info is not None

tests/graphs/smart_scraper_ernie_test.py

@@ -49,4 +49,4 @@ def test_get_execution_info(graph_config: dict):
 
     graph_exec_info = smart_scraper_graph.get_execution_info()
 
-    assert graph_exec_info is not None
+    assert graph_exec_info is not None

tests/graphs/xml_scraper_openai_test.py

@@ -0,0 +1,94 @@
+"""
+xml_scraper_test
+"""
+import os
+import pytest
+from dotenv import load_dotenv
+from scrapegraphai.graphs import XMLScraperGraph
+from scrapegraphai.utils import convert_to_csv, convert_to_json, prettify_exec_info
+
+load_dotenv()
+
+# ************************************************
+# Define the test fixtures and helpers
+# ************************************************
+
+@pytest.fixture
+def graph_config():
+    """
+    Configuration for the XMLScraperGraph
+    """
+    openai_key = os.getenv("OPENAI_APIKEY")
+    return {
+        "llm": {
+            "api_key": openai_key,
+            "model": "openai/gpt-4o",
+        },
+        "verbose": False,
+    }
+
+@pytest.fixture
+def xml_content():
+    """
+    Fixture to read the XML file content
+    """
+    FILE_NAME = "inputs/books.xml"
+    curr_dir = os.path.dirname(os.path.realpath(__file__))
+    file_path = os.path.join(curr_dir, FILE_NAME)
+
+    with open(file_path, 'r', encoding="utf-8") as file:
+        return file.read()
+
+# ************************************************
+# Define the test cases
+# ************************************************
+
+def test_xml_scraper_graph(graph_config: dict, xml_content: str):
+    """
+    Test the XMLScraperGraph scraping pipeline
+    """
+    xml_scraper_graph = XMLScraperGraph(
+        prompt="List me all the authors, title and genres of the books",
+        source=xml_content,  # Pass the XML content
+        config=graph_config
+    )
+
+    result = xml_scraper_graph.run()
+
+    assert result is not None
+
+def test_xml_scraper_execution_info(graph_config: dict, xml_content: str):
+    """
+    Test getting the execution info of XMLScraperGraph
+    """
+    xml_scraper_graph = XMLScraperGraph(
+        prompt="List me all the authors, title and genres of the books",
+        source=xml_content,  # Pass the XML content
+        config=graph_config
+    )
+
+    xml_scraper_graph.run()
+
+    graph_exec_info = xml_scraper_graph.get_execution_info()
+
+    assert graph_exec_info is not None
+    print(prettify_exec_info(graph_exec_info))
+
+def test_xml_scraper_save_results(graph_config: dict, xml_content: str):
+    """
+    Test saving the results of XMLScraperGraph to CSV and JSON
+    """
+    xml_scraper_graph = XMLScraperGraph(
+        prompt="List me all the authors, title and genres of the books",
+        source=xml_content,  # Pass the XML content
+        config=graph_config
+    )
+
+    result = xml_scraper_graph.run()
+
+    # Save to csv and json
+    convert_to_csv(result, "result")
+    convert_to_json(result, "result")
+
+    assert os.path.exists("result.csv")
+    assert os.path.exists("result.json")