tests/unit/graph/test_nodes.py

"""Unit tests for graph nodes."""

import pytest

from src.agents.graph.nodes import judge_node, search_node, supervisor_node
from src.agents.graph.state import ResearchState


@pytest.mark.asyncio
async def test_judge_node_initialization(mocker):
    """Test judge creates initial hypothesis if none exist."""
    # Mock get_model to avoid needing real API keys
    mocker.patch("src.agents.graph.nodes.get_model", return_value=mocker.Mock())

    # Create a mock assessment with attributes (sexual health domain)
    mock_hypothesis = mocker.Mock()
    mock_hypothesis.drug = "Testosterone"
    mock_hypothesis.target = "Androgen Receptor"
    mock_hypothesis.pathway = "HPG Axis"
    mock_hypothesis.effect = "Libido Enhancement"
    mock_hypothesis.confidence = 0.8

    mock_assessment = mocker.Mock()
    mock_assessment.hypotheses = [mock_hypothesis]

    mock_result = mocker.Mock()
    mock_result.output = mock_assessment

    # Mock the Agent class entirely
    mock_agent_instance = mocker.Mock()
    mock_agent_instance.run = mocker.AsyncMock(return_value=mock_result)
    mocker.patch("src.agents.graph.nodes.Agent", return_value=mock_agent_instance)

    state: ResearchState = {
        "query": "Does stress affect libido?",
        "hypotheses": [],
        "conflicts": [],
        "evidence_ids": [],
        "messages": [],
        "next_step": "judge",
        "iteration_count": 0,
        "max_iterations": 10,
    }

    update = await judge_node(state)

    assert "hypotheses" in update
    assert len(update["hypotheses"]) == 1
    assert update["hypotheses"][0].id == "Testosterone"
    assert update["hypotheses"][0].status == "proposed"


@pytest.mark.asyncio
async def test_supervisor_termination():
    """Test supervisor forces synthesis at max iterations."""
    state: ResearchState = {
        "query": "test",
        "hypotheses": [],
        "conflicts": [],
        "evidence_ids": [],
        "messages": [],
        "next_step": "search",
        "iteration_count": 10,  # Max reached
        "max_iterations": 10,
    }

    update = await supervisor_node(state)
    assert update["next_step"] == "synthesize"


@pytest.mark.asyncio
async def test_search_node_execution(mocker):
    """Test search node calls tools (mocked)."""
    # Mock the tools
    mocker.patch("src.tools.pubmed.PubMedTool.search", return_value=[])
    mocker.patch("src.tools.clinicaltrials.ClinicalTrialsTool.search", return_value=[])
    mocker.patch("src.tools.europepmc.EuropePMCTool.search", return_value=[])

    state: ResearchState = {
        "query": "test",
        "hypotheses": [],
        "conflicts": [],
        "evidence_ids": [],
        "messages": [],
        "next_step": "search",
        "iteration_count": 0,
        "max_iterations": 10,
    }

    update = await search_node(state)
    assert "messages" in update
    # Matches "Found 0 total, 0 unique new papers."
    assert "0 unique new papers" in update["messages"][0].content