Scrapegraph-ai/tests/nodes/search_link_node_test.py

from unittest.mock import MagicMock, patch

import pytest
from langchain_community.chat_models import ChatOllama

from scrapegraphai.nodes import SearchLinkNode


@pytest.fixture
def setup():
    """
    Setup the SearchLinkNode and initial state for testing.
    """
    # Define the configuration for the graph
    graph_config = {
        "llm": {"model_name": "ollama/llama3", "temperature": 0, "streaming": True},
    }

    # Instantiate the LLM model with the configuration
    llm_model = ChatOllama(graph_config["llm"])

    # Define the SearchLinkNode with necessary configurations
    search_link_node = SearchLinkNode(
        input=["user_prompt", "parsed_content_chunks"],
        output=["relevant_links"],
        node_config={"llm_model": llm_model, "verbose": False},
    )

    # Define the initial state for the node
    initial_state = {
        "user_prompt": "Example user prompt",
        "parsed_content_chunks": [
            {"page_content": "Example page content 1"},
            {"page_content": "Example page content 2"},
            # Add more example page content dictionaries as needed
        ],
    }

    return search_link_node, initial_state


def test_search_link_node(setup):
    """
    Test the SearchLinkNode execution.
    """
    search_link_node, initial_state = setup

    # Patch the execute method to avoid actual network calls and return a mock response
    with patch.object(
        SearchLinkNode,
        "execute",
        return_value={"relevant_links": ["http://example.com"]},
    ) as mock_execute:
        result = search_link_node.execute(initial_state)

        # Check if the result is not None
        assert result is not None
        # Additional assertion to check the returned value
        assert "relevant_links" in result
        assert isinstance(result["relevant_links"], list)
        assert len(result["relevant_links"]) > 0
        # Ensure the execute method was called once
        mock_execute.assert_called_once_with(initial_state)