From a406264a125318d39234cdbdfc6cfaa540b20464 Mon Sep 17 00:00:00 2001 From: "codebeaver-ai[bot]" <192081515+codebeaver-ai[bot]@users.noreply.github.com> Date: Tue, 28 Jan 2025 11:02:14 +0000 Subject: [PATCH 1/3] test: Update coverage improvement test for tests/graphs/abstract_graph_test.py --- tests/graphs/abstract_graph_test.py | 55 ++++++++++++++++++++++++----- 1 file changed, 47 insertions(+), 8 deletions(-) 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 From 9919e7c12211039f03381b6b7cc0167fb268a3fb Mon Sep 17 00:00:00 2001 From: "codebeaver-ai[bot]" <192081515+codebeaver-ai[bot]@users.noreply.github.com> Date: Tue, 28 Jan 2025 11:02:15 +0000 Subject: [PATCH 2/3] test: Update coverage improvement test for tests/test_json_scraper_graph.py --- tests/test_json_scraper_graph.py | 46 ++++++++++++++++++++++++++++++++ 1 file changed, 46 insertions(+) diff --git a/tests/test_json_scraper_graph.py b/tests/test_json_scraper_graph.py index 7ef72c71..0f00691f 100644 --- a/tests/test_json_scraper_graph.py +++ b/tests/test_json_scraper_graph.py @@ -49,4 +49,50 @@ class TestJSONScraperGraph: mock_execute.assert_called_once_with({"user_prompt": "Summarize the data from all JSON files", "json_dir": "path/to/json/directory"}) mock_fetch_node.assert_called_once() mock_generate_answer_node.assert_called_once() + mock_create_llm.assert_called_once_with({"model": "test-model", "temperature": 0}) + + @pytest.fixture + def mock_llm_model(self): + return Mock() + + @pytest.fixture + def mock_embedder_model(self): + return Mock() + + @patch('scrapegraphai.graphs.json_scraper_graph.FetchNode') + @patch('scrapegraphai.graphs.json_scraper_graph.GenerateAnswerNode') + @patch.object(JSONScraperGraph, '_create_llm') + def test_json_scraper_graph_with_single_file(self, mock_create_llm, mock_generate_answer_node, mock_fetch_node, mock_llm_model, mock_embedder_model): + """ + Test JSONScraperGraph with a single JSON file. + This test checks if the graph correctly handles a single JSON file input + and processes it to generate an answer. + """ + # Mock the _create_llm method to return a mock LLM model + mock_create_llm.return_value = mock_llm_model + + # Mock the execute method of BaseGraph + with patch('scrapegraphai.graphs.json_scraper_graph.BaseGraph.execute') as mock_execute: + mock_execute.return_value = ({"answer": "Mocked answer for single JSON file"}, {}) + + # Create a JSONScraperGraph instance with a single JSON file + graph = JSONScraperGraph( + prompt="Analyze the data from the JSON file", + source="path/to/single/file.json", + config={"llm": {"model": "test-model", "temperature": 0}}, + schema=BaseModel + ) + + # Set mocked embedder model + graph.embedder_model = mock_embedder_model + + # Run the graph + result = graph.run() + + # Assertions + assert result == "Mocked answer for single JSON file" + assert graph.input_key == "json" + mock_execute.assert_called_once_with({"user_prompt": "Analyze the data from the JSON file", "json": "path/to/single/file.json"}) + mock_fetch_node.assert_called_once() + mock_generate_answer_node.assert_called_once() mock_create_llm.assert_called_once_with({"model": "test-model", "temperature": 0}) \ No newline at end of file From 16c688f090559497175677010bbb285c9d53cf22 Mon Sep 17 00:00:00 2001 From: "codebeaver-ai[bot]" <192081515+codebeaver-ai[bot]@users.noreply.github.com> Date: Tue, 28 Jan 2025 11:02:17 +0000 Subject: [PATCH 3/3] test: Update coverage improvement test for tests/test_search_graph.py --- tests/test_search_graph.py | 26 +++++++++++++++++++++++++- 1 file changed, 25 insertions(+), 1 deletion(-) diff --git a/tests/test_search_graph.py b/tests/test_search_graph.py index 1384988d..f8e51c9c 100644 --- a/tests/test_search_graph.py +++ b/tests/test_search_graph.py @@ -33,4 +33,28 @@ class TestSearchGraph: search_graph.run() # Assert - assert search_graph.get_considered_urls() == urls \ No newline at end of file + assert search_graph.get_considered_urls() == urls + + @patch('scrapegraphai.graphs.search_graph.BaseGraph') + @patch('scrapegraphai.graphs.abstract_graph.AbstractGraph._create_llm') + def test_run_no_answer_found(self, mock_create_llm, mock_base_graph): + """ + Test that the run() method returns "No answer found." when the final state + doesn't contain an "answer" key. + """ + # Arrange + prompt = "Test prompt" + config = {"llm": {"model": "test-model"}} + + # Mock the _create_llm method to return a MagicMock + mock_create_llm.return_value = MagicMock() + + # Mock the execute method to set the final_state without an "answer" key + mock_base_graph.return_value.execute.return_value = ({"urls": []}, {}) + + # Act + search_graph = SearchGraph(prompt, config) + result = search_graph.run() + + # Assert + assert result == "No answer found." \ No newline at end of file