Add Unit Tests (#4)
This commit is contained in:
Lamont Huffman
@@ -1,8 +1,9 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
# tests/agent_deep_research_test.py
|
||||
import logging
|
||||
import os
|
||||
import shutil
|
||||
import tempfile
|
||||
from unittest.mock import Mock, patch
|
||||
from unittest.mock import Mock, AsyncMock, patch
|
||||
|
||||
import pytest
|
||||
from agentscope.formatter import DashScopeChatFormatter
|
||||
@@ -11,11 +12,7 @@ from agentscope.memory import InMemoryMemory
|
||||
from agentscope.message import Msg
|
||||
from agentscope.model import DashScopeChatModel
|
||||
|
||||
from deep_research.agent_deep_research.deep_research_agent import (
|
||||
DeepResearchAgent,
|
||||
)
|
||||
|
||||
# Import the main function to be tested
|
||||
from deep_research.agent_deep_research.deep_research_agent import DeepResearchAgent
|
||||
from deep_research.agent_deep_research.main import main
|
||||
|
||||
|
||||
@@ -41,7 +38,7 @@ def temp_working_dir():
|
||||
@pytest.fixture
|
||||
def mock_tavily_client():
|
||||
"""Create a mocked Tavily client"""
|
||||
client = Mock(spec=StdIOStatefulClient)
|
||||
client = AsyncMock(spec=StdIOStatefulClient)
|
||||
client.name = "tavily_mcp"
|
||||
client.connect = AsyncMock()
|
||||
client.close = AsyncMock()
|
||||
@@ -68,25 +65,6 @@ def mock_model():
|
||||
return model
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_agent(mock_model, mock_formatter, mock_memory, mock_tavily_client):
|
||||
"""Create a mocked DeepResearchAgent instance"""
|
||||
agent = Mock(spec=DeepResearchAgent)
|
||||
agent.return_value = agent # Make the mock instance return itself
|
||||
agent.model = mock_model
|
||||
agent.formatter = mock_formatter
|
||||
agent.memory = mock_memory
|
||||
agent.search_mcp_client = mock_tavily_client
|
||||
return agent
|
||||
|
||||
|
||||
class AsyncMock(Mock):
|
||||
"""Helper class for async mocks"""
|
||||
|
||||
async def __call__(self, *args, **kwargs):
|
||||
return super().__call__(*args, **kwargs)
|
||||
|
||||
|
||||
class TestDeepResearchAgent:
|
||||
"""Test suite for Deep Research Agent functionality"""
|
||||
|
||||
@@ -97,18 +75,19 @@ class TestDeepResearchAgent:
|
||||
temp_working_dir,
|
||||
):
|
||||
"""Test agent initialization with valid parameters"""
|
||||
agent = DeepResearchAgent(
|
||||
name="Friday",
|
||||
sys_prompt="You are a helpful assistant named Friday.",
|
||||
model=mock_model,
|
||||
formatter=DashScopeChatFormatter(),
|
||||
memory=InMemoryMemory(),
|
||||
search_mcp_client=mock_tavily_client,
|
||||
tmp_file_storage_dir=temp_working_dir,
|
||||
)
|
||||
with patch("asyncio.create_task"):
|
||||
agent = DeepResearchAgent(
|
||||
name="Friday",
|
||||
sys_prompt="You are a helpful assistant named Friday.",
|
||||
model=mock_model,
|
||||
formatter=DashScopeChatFormatter(),
|
||||
memory=InMemoryMemory(),
|
||||
search_mcp_client=mock_tavily_client,
|
||||
tmp_file_storage_dir=temp_working_dir,
|
||||
)
|
||||
|
||||
assert agent.name == "Friday"
|
||||
assert agent.sys_prompt == "You are a helpful assistant named Friday."
|
||||
assert agent.sys_prompt.startswith("You are a helpful assistant named Friday.")
|
||||
assert agent.tmp_file_storage_dir == temp_working_dir
|
||||
assert os.path.exists(temp_working_dir)
|
||||
|
||||
@@ -121,72 +100,41 @@ class TestDeepResearchAgent:
|
||||
temp_working_dir,
|
||||
):
|
||||
"""Test main function with successful execution"""
|
||||
# Mock the StdIOStatefulClient constructor
|
||||
with patch(
|
||||
"deep_research.agent_deep_research.main.StdIOStatefulClient",
|
||||
return_value=mock_tavily_client,
|
||||
):
|
||||
# Mock the DeepResearchAgent constructor
|
||||
with patch(
|
||||
"deep_research.agent_deep_research.main.DeepResearchAgent",
|
||||
autospec=True,
|
||||
) as mock_agent_class:
|
||||
mock_agent_instance = Mock()
|
||||
mock_agent_instance.return_value = mock_agent_instance
|
||||
mock_agent_instance.__call__ = AsyncMock(
|
||||
return_value=Msg("Friday", "Test response", "assistant"),
|
||||
)
|
||||
mock_agent_class.return_value = mock_agent_instance
|
||||
mock_agent = AsyncMock()
|
||||
mock_agent.return_value = Msg("Friday", "Test response", "assistant")
|
||||
mock_agent_class.return_value = mock_agent
|
||||
|
||||
# Mock os.makedirs
|
||||
with patch("os.makedirs") as mock_makedirs:
|
||||
# Run the main function with a test query
|
||||
test_query = "Test research question"
|
||||
msg = Msg("Bob", test_query, "user")
|
||||
with patch.dict(os.environ, {"AGENT_OPERATION_DIR": temp_working_dir}):
|
||||
test_query = "Test research question"
|
||||
msg = Msg("Bob", test_query, "user")
|
||||
|
||||
await main(test_query)
|
||||
await main(test_query)
|
||||
|
||||
# Verify initialization calls
|
||||
mock_makedirs.assert_called_once_with(
|
||||
temp_working_dir,
|
||||
exist_ok=True,
|
||||
)
|
||||
mock_agent_class.assert_called_once()
|
||||
mock_makedirs.assert_called_once_with(temp_working_dir, exist_ok=True)
|
||||
mock_agent_class.assert_called_once()
|
||||
|
||||
# Verify agent was called with the correct message
|
||||
mock_agent_instance.__call__.assert_called_once_with(msg)
|
||||
# ✅ Use assert_called_once() + manual argument check
|
||||
mock_agent.assert_called_once()
|
||||
call_arg = mock_agent.call_args[0][0]
|
||||
assert call_arg.name == "Bob"
|
||||
assert call_arg.content == "Test research question"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_main_function_with_missing_env_vars(self):
|
||||
"""Test main function handles missing environment variables"""
|
||||
# Test missing Tavily API key
|
||||
with patch.dict(os.environ, clear=True):
|
||||
with pytest.raises(Exception):
|
||||
await main("Test query")
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_main_function_connection_failure(
|
||||
self,
|
||||
mock_env_vars,
|
||||
temp_working_dir,
|
||||
):
|
||||
"""Test main function handles connection failures"""
|
||||
# Mock the StdIOStatefulClient to raise an exception
|
||||
with patch(
|
||||
"deep_research.agent_deep_research.main.StdIOStatefulClient",
|
||||
) as mock_client:
|
||||
mock_client_instance = Mock()
|
||||
mock_client_instance.connect = AsyncMock(
|
||||
side_effect=Exception("Connection failed"),
|
||||
)
|
||||
mock_client.return_value = mock_client_instance
|
||||
|
||||
# Run the main function and expect exception
|
||||
with pytest.raises(Exception) as exc_info:
|
||||
await main("Test query")
|
||||
|
||||
assert "Connection failed" in str(exc_info.value)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_agent_cleanup(
|
||||
self,
|
||||
@@ -198,90 +146,32 @@ class TestDeepResearchAgent:
|
||||
"deep_research.agent_deep_research.main.StdIOStatefulClient",
|
||||
return_value=mock_tavily_client,
|
||||
):
|
||||
# Run main function
|
||||
await main("Test query")
|
||||
with patch.dict(os.environ, {"AGENT_OPERATION_DIR": "/tmp"}):
|
||||
await main("Test query")
|
||||
|
||||
# Verify client close was called
|
||||
mock_tavily_client.close.assert_called_once()
|
||||
|
||||
def test_working_directory_creation(self, temp_working_dir):
|
||||
"""Test working directory is created correctly"""
|
||||
test_dir = os.path.join(temp_working_dir, "test_subdir")
|
||||
|
||||
# Test directory creation
|
||||
os.makedirs(test_dir, exist_ok=True)
|
||||
assert os.path.exists(test_dir)
|
||||
|
||||
# Test exist_ok=True behavior
|
||||
os.makedirs(test_dir, exist_ok=True) # Should not raise error
|
||||
|
||||
|
||||
class TestErrorHandling:
|
||||
"""Test suite for error handling scenarios"""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_model_failure(self, mock_env_vars, mock_tavily_client):
|
||||
"""Test handling of model failures"""
|
||||
with patch(
|
||||
"deep_research.agent_deep_research.main.StdIOStatefulClient",
|
||||
return_value=mock_tavily_client,
|
||||
):
|
||||
with patch(
|
||||
"deep_research.agent_deep_research.main.DeepResearchAgent",
|
||||
) as mock_agent_class:
|
||||
mock_agent = Mock()
|
||||
mock_agent.__call__ = AsyncMock(
|
||||
side_effect=Exception("Model error"),
|
||||
)
|
||||
mock_agent_class.return_value = mock_agent
|
||||
|
||||
with pytest.raises(Exception) as exc_info:
|
||||
await main("Test query")
|
||||
|
||||
assert "Model error" in str(exc_info.value)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_filesystem_errors(self, mock_env_vars, mock_tavily_client):
|
||||
"""Test handling of filesystem errors"""
|
||||
# Test with invalid directory path
|
||||
invalid_dir = "/invalid/path/that/does/not/exist"
|
||||
|
||||
with patch.dict(os.environ, {"AGENT_OPERATION_DIR": invalid_dir}):
|
||||
with patch(
|
||||
"os.makedirs",
|
||||
side_effect=PermissionError("Permission denied"),
|
||||
):
|
||||
with pytest.raises(PermissionError):
|
||||
await main("Test query")
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_logging_output(
|
||||
self,
|
||||
mock_env_vars,
|
||||
mock_tavily_client,
|
||||
caplog,
|
||||
):
|
||||
"""Test logging output is generated correctly"""
|
||||
with patch(
|
||||
"deep_research.agent_deep_research.main.StdIOStatefulClient",
|
||||
return_value=mock_tavily_client,
|
||||
"deep_research.agent_deep_research.main.StdIOStatefulClient",
|
||||
return_value=mock_tavily_client,
|
||||
):
|
||||
with patch(
|
||||
"deep_research.agent_deep_research.main.DeepResearchAgent",
|
||||
) as mock_agent_class:
|
||||
mock_agent = Mock()
|
||||
mock_agent.__call__ = AsyncMock(
|
||||
return_value=Msg("Friday", "Test response", "assistant"),
|
||||
)
|
||||
mock_agent_class.return_value = mock_agent
|
||||
|
||||
await main("Test query")
|
||||
|
||||
# Verify debug logs are present
|
||||
assert any(
|
||||
"DEBUG" in record.levelname for record in caplog.records
|
||||
)
|
||||
|
||||
with patch.dict(os.environ, {"AGENT_OPERATION_DIR": "/invalid/path"}):
|
||||
with patch("os.makedirs", side_effect=PermissionError("Permission denied")):
|
||||
with pytest.raises(PermissionError):
|
||||
await main("Test query")
|
||||
|
||||
if __name__ == "__main__":
|
||||
pytest.main(["-v", __file__])
|
||||
pytest.main(["-v", __file__])
|
||||
Reference in New Issue
Block a user