import pytest from pydantic import BaseModel from scrapegraphai.graphs.script_creator_graph import ScriptCreatorGraph from scrapegraphai.graphs.script_creator_multi_graph import ( BaseGraph, ScriptCreatorMultiGraph, ) @pytest.fixture(autouse=True) def set_api_key_env(monkeypatch): monkeypatch.setenv("OPENAI_API_KEY", "dummy") # Dummy classes to simulate behavior for testing class DummyGraph: def __init__(self, final_state, execution_info): self.final_state = final_state self.execution_info = execution_info def execute(self, inputs): return self.final_state, self.execution_info class DummySchema(BaseModel): field: str = "dummy" class TestScriptCreatorMultiGraph: """Tests for ScriptCreatorMultiGraph.""" def test_run_success(self): """Test run() returns the merged script when execution is successful.""" prompt = "Test prompt" source = ["http://example.com"] config = {"llm": {"model": "openai/test-model"}} schema = DummySchema instance = ScriptCreatorMultiGraph(prompt, source, config, schema) # Set necessary attributes that are expected by _create_graph() and the run() method. instance.llm_model = {"model": "openai/test-model"} instance.schema = {"type": "dummy"} # Replace the graph with a dummy graph that simulates successful execution. dummy_final_state = {"merged_script": "print('Hello World')"} dummy_execution_info = {"info": "dummy"} instance.graph = DummyGraph(dummy_final_state, dummy_execution_info) result = instance.run() assert result == "print('Hello World')" def test_run_failure(self): """Test run() returns failure message when merged_script is missing.""" prompt = "Test prompt" source = ["http://example.com"] config = {"llm": {"model": "openai/test-model"}} schema = DummySchema instance = ScriptCreatorMultiGraph(prompt, source, config, schema) instance.llm_model = {"model": "openai/test-model"} instance.schema = {"type": "dummy"} dummy_final_state = {"other_key": "no script"} dummy_execution_info = {"info": "dummy"} instance.graph = DummyGraph(dummy_final_state, dummy_execution_info) result = instance.run() assert result == "Failed to generate the script." def test_create_graph_structure(self): """Test _create_graph() returns a BaseGraph with the correct graph name and structure.""" prompt = "Test prompt" source = [] config = {"llm": {"model": "openai/test-model"}} schema = DummySchema instance = ScriptCreatorMultiGraph(prompt, source, config, schema) # Manually assign llm_model and schema for node configuration in the graph. instance.llm_model = {"model": "openai/test-model"} instance.schema = {"type": "dummy"} graph = instance._create_graph() assert isinstance(graph, BaseGraph) assert hasattr(graph, "graph_name") assert graph.graph_name == "ScriptCreatorMultiGraph" # Check that the graph has two nodes. assert len(graph.nodes) == 2 # Optional: Check that the edges list is correctly formed. assert len(graph.edges) == 1 def test_config_deepcopy(self): """Test that the configuration is deep copied during initialization.""" prompt = "Test prompt" source = [] config = {"llm": {"model": "openai/test-model"}, "other": [1, 2, 3]} schema = DummySchema instance = ScriptCreatorMultiGraph(prompt, source, config, schema) # Modify the original config. config["llm"]["model"] = "changed-model" config["other"].append(4) # Verify that the config copied within instance remains unchanged. assert instance.copy_config["llm"]["model"] == "openai/test-model" assert instance.copy_config["other"] == [1, 2, 3] def test_init_attributes(self): """Test that initial attributes are set correctly upon initialization.""" prompt = "Initialization test" source = ["http://init.com"] config = {"llm": {"model": "openai/init-model"}, "param": [1, 2]} schema = DummySchema instance = ScriptCreatorMultiGraph(prompt, source, config, schema) # Check that basic attributes are set correctly assert instance.prompt == prompt assert instance.source == source # Check that copy_config is a deep copy and equals the original config assert instance.copy_config == { "llm": {"model": "openai/init-model"}, "param": [1, 2], } # For classes, deepcopy returns the same object, so the copy_schema should equal schema assert instance.copy_schema == DummySchema def test_run_no_schema(self): """Test run() when schema is None.""" prompt = "No schema prompt" source = ["http://noschema.com"] config = {"llm": {"model": "openai/no-schema"}} instance = ScriptCreatorMultiGraph(prompt, source, config, schema=None) instance.llm_model = {"model": "openai/no-schema"} instance.schema = None dummy_final_state = {"merged_script": "print('No Schema Script')"} dummy_execution_info = {"info": "no schema"} instance.graph = DummyGraph(dummy_final_state, dummy_execution_info) result = instance.run() assert result == "print('No Schema Script')" def test_create_graph_node_configs(self): """Test that _create_graph() sets correct node configurations for its nodes.""" prompt = "Graph config test" source = ["http://graphconfig.com"] config = {"llm": {"model": "openai/graph-model"}, "extra": [10]} schema = DummySchema instance = ScriptCreatorMultiGraph(prompt, source, config, schema) # Manually assign llm_model and schema for node configuration instance.llm_model = {"model": "openai/graph-model"} instance.schema = {"type": "graph-dummy"} graph = instance._create_graph() # Validate configuration of the first node (GraphIteratorNode) node1 = graph.nodes[0] assert node1.node_config["graph_instance"] == ScriptCreatorGraph assert node1.node_config["scraper_config"] == instance.copy_config # Validate configuration of the second node (MergeGeneratedScriptsNode) node2 = graph.nodes[1] assert node2.node_config["llm_model"] == instance.llm_model assert node2.node_config["schema"] == instance.schema def test_entry_point_node(self): """Test that the graph entry point is the GraphIteratorNode (the first node).""" prompt = "Entry point test" source = ["http://entrypoint.com"] config = {"llm": {"model": "openai/test-model"}} schema = DummySchema instance = ScriptCreatorMultiGraph(prompt, source, config, schema) instance.llm_model = {"model": "openai/test-model"} instance.schema = {"type": "dummy"} graph = instance._create_graph() assert graph.entry_point == graph.nodes[0] def test_run_exception(self): """Test that run() propagates exceptions raised by graph.execute.""" prompt = "Exception test" source = ["http://exception.com"] config = {"llm": {"model": "openai/test-model"}} schema = DummySchema instance = ScriptCreatorMultiGraph(prompt, source, config, schema) instance.llm_model = {"model": "openai/test-model"} instance.schema = {"type": "dummy"} # Create a dummy graph that raises an exception when execute is called. class ExceptionGraph: def execute(self, inputs): raise ValueError("Testing exception") instance.graph = ExceptionGraph() with pytest.raises(ValueError, match="Testing exception"): instance.run() def test_run_with_empty_prompt(self): """Test run() method with an empty prompt.""" prompt = "" source = ["http://emptyprompt.com"] config = {"llm": {"model": "openai/test-model"}} schema = DummySchema instance = ScriptCreatorMultiGraph(prompt, source, config, schema) instance.llm_model = {"model": "openai/test-model"} instance.schema = {"type": "dummy"} dummy_final_state = {"merged_script": "print('Empty prompt')"} dummy_execution_info = {"info": "empty prompt"} instance.graph = DummyGraph(dummy_final_state, dummy_execution_info) result = instance.run() assert result == "print('Empty prompt')" def test_run_called_twice(self): """Test that running run() twice returns consistent and updated results.""" prompt = "Twice test" source = ["http://twicetest.com"] config = {"llm": {"model": "openai/test-model"}} schema = DummySchema instance = ScriptCreatorMultiGraph(prompt, source, config, schema) instance.llm_model = {"model": "openai/test-model"} instance.schema = {"type": "dummy"} dummy_final_state = {"merged_script": "print('First run')"} dummy_execution_info = {"info": "first run"} dummy_graph = DummyGraph(dummy_final_state, dummy_execution_info) instance.graph = dummy_graph result1 = instance.run() # Modify dummy graph's state for the second run. dummy_graph.final_state["merged_script"] = "print('Second run')" dummy_graph.execution_info = {"info": "second run"} result2 = instance.run() assert result1 == "print('First run')" assert result2 == "print('Second run')"