diff --git a/tests/graphs/abstract_graph_test.py b/tests/graphs/abstract_graph_test.py index c17ef09a..280f1f77 100644 --- a/tests/graphs/abstract_graph_test.py +++ b/tests/graphs/abstract_graph_test.py @@ -1,18 +1,16 @@ -""" -Tests for the AbstractGraph. -""" - -from unittest.mock import patch - import pytest + from langchain_aws import ChatBedrock from langchain_ollama import ChatOllama from langchain_openai import AzureChatOpenAI, ChatOpenAI - from scrapegraphai.graphs import AbstractGraph, BaseGraph from scrapegraphai.models import DeepSeek, OneApi from scrapegraphai.nodes import FetchNode, ParseNode +from unittest.mock import Mock, patch +""" +Tests for the AbstractGraph. +""" class TestGraph(AbstractGraph): def __init__(self, prompt: str, config: dict): @@ -50,7 +48,6 @@ class TestGraph(AbstractGraph): return self.final_state.get("answer", "No answer found.") - class TestAbstractGraph: @pytest.mark.parametrize( "llm_config, expected_model", @@ -161,3 +158,45 @@ class TestAbstractGraph: result = await graph.run_safe_async() assert result == "Async result" mock_run.assert_called_once() + + def test_create_llm_with_custom_model_instance(self): + """ + Test that the _create_llm method correctly uses a custom model instance + when provided in the configuration. + """ + mock_model = Mock() + mock_model.model_name = "custom-model" + + config = { + "llm": { + "model_instance": mock_model, + "model_tokens": 1000, + "model": "custom/model" + } + } + + graph = TestGraph("Test prompt", config) + + assert graph.llm_model == mock_model + assert graph.model_token == 1000 + + def test_set_common_params(self): + """ + Test that the set_common_params method correctly updates the configuration + of all nodes in the graph. + """ + # Create a mock graph with mock nodes + mock_graph = Mock() + mock_node1 = Mock() + mock_node2 = Mock() + mock_graph.nodes = [mock_node1, mock_node2] + + # Create a TestGraph instance with the mock graph + with patch('scrapegraphai.graphs.abstract_graph.AbstractGraph._create_graph', return_value=mock_graph): + graph = TestGraph("Test prompt", {"llm": {"model": "openai/gpt-3.5-turbo", "openai_api_key": "sk-test"}}) + + # Call set_common_params with test parameters + test_params = {"param1": "value1", "param2": "value2"} + graph.set_common_params(test_params) + + # Assert that update_config was called on each node with the correct parameters \ No newline at end of file