Spaces:
Running
Running
File size: 3,567 Bytes
c690006 |
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 95 96 97 98 99 100 |
"""Unit tests for HypothesisAgent."""
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from agent_framework import AgentRunResponse
from src.agents.hypothesis_agent import HypothesisAgent
from src.utils.models import Citation, Evidence, HypothesisAssessment, MechanismHypothesis
@pytest.fixture
def sample_evidence():
return [
Evidence(
content="Metformin activates AMPK, which inhibits mTOR signaling...",
citation=Citation(
source="pubmed",
title="Metformin and AMPK",
url="https://pubmed.ncbi.nlm.nih.gov/12345/",
date="2023",
),
)
]
@pytest.fixture
def mock_assessment():
return HypothesisAssessment(
hypotheses=[
MechanismHypothesis(
drug="Metformin",
target="AMPK",
pathway="mTOR inhibition",
effect="Reduced cancer cell proliferation",
confidence=0.75,
search_suggestions=["metformin AMPK cancer", "mTOR cancer therapy"],
)
],
primary_hypothesis=None,
knowledge_gaps=["Clinical trial data needed"],
recommended_searches=["metformin clinical trial cancer"],
)
@pytest.mark.asyncio
async def test_hypothesis_agent_generates_hypotheses(sample_evidence, mock_assessment):
"""HypothesisAgent should generate mechanistic hypotheses."""
store = {"current": sample_evidence, "hypotheses": []}
with patch("src.agents.hypothesis_agent.Agent") as mock_agent_class:
mock_result = MagicMock()
mock_result.output = mock_assessment
# pydantic-ai Agent returns an object with .output for structured output
mock_agent_class.return_value.run = AsyncMock(return_value=mock_result)
agent = HypothesisAgent(store)
response = await agent.run("metformin cancer")
assert isinstance(response, AgentRunResponse)
assert "AMPK" in response.messages[0].text
assert len(store["hypotheses"]) == 1
assert store["hypotheses"][0].drug == "Metformin"
@pytest.mark.asyncio
async def test_hypothesis_agent_no_evidence():
"""HypothesisAgent should handle empty evidence gracefully."""
store = {"current": [], "hypotheses": []}
with patch("src.agents.hypothesis_agent.Agent"):
agent = HypothesisAgent(store)
response = await agent.run("test query")
assert "No evidence" in response.messages[0].text
assert len(store["hypotheses"]) == 0
@pytest.mark.asyncio
async def test_hypothesis_agent_uses_embeddings(sample_evidence, mock_assessment):
"""HypothesisAgent should pass embeddings to prompt formatter."""
store = {"current": sample_evidence, "hypotheses": []}
mock_embeddings = MagicMock()
with patch("src.agents.hypothesis_agent.Agent") as mock_agent_class:
# Mock format_hypothesis_prompt to check if embeddings were passed
with patch("src.agents.hypothesis_agent.format_hypothesis_prompt") as mock_format:
mock_format.return_value = "Prompt"
mock_result = MagicMock()
mock_result.output = mock_assessment
mock_agent_class.return_value.run = AsyncMock(return_value=mock_result)
agent = HypothesisAgent(store, embedding_service=mock_embeddings)
await agent.run("query")
mock_format.assert_called_once()
_args, kwargs = mock_format.call_args
assert kwargs["embeddings"] == mock_embeddings
|