File size: 2,975 Bytes
a820b5b
 
 
 
 
 
 
 
 
f160233
a820b5b
46b9533
 
f160233
89f1173
f160233
89f1173
 
 
 
f160233
 
 
 
 
 
 
46b9533
 
 
 
 
f160233
a820b5b
fa696e8
a820b5b
 
 
 
 
 
 
 
 
 
 
 
 
89f1173
a820b5b
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
f160233
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
"""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