Your Name
384 lines
14 KiB
Python
384 lines
14 KiB
Python
"""Unit tests for CleverClaude agent management."""
|
|
|
|
from unittest.mock import AsyncMock, patch
|
|
|
|
import pytest
|
|
|
|
from cleverclaude.agents.agent import BaseAgent
|
|
from cleverclaude.agents.manager import AgentManager
|
|
from cleverclaude.agents.types import AgentConfig, AgentStatus, AgentType
|
|
from cleverclaude.config.settings import Settings
|
|
|
|
|
|
@pytest.mark.unit
|
|
@pytest.mark.async_test
|
|
class TestAgentManager:
|
|
"""Test suite for AgentManager class."""
|
|
|
|
async def test_initialization(self, test_settings: Settings, async_session, mock_redis):
|
|
"""Test AgentManager initialization."""
|
|
manager = AgentManager(test_settings.agents, async_session, mock_redis)
|
|
|
|
assert manager.config == test_settings.agents
|
|
assert manager.session == async_session
|
|
assert manager.redis == mock_redis
|
|
assert manager.agents == {}
|
|
assert manager._initialized is False
|
|
|
|
async def test_initialize_manager(self, test_settings: Settings, async_session, mock_redis):
|
|
"""Test manager initialization process."""
|
|
manager = AgentManager(test_settings.agents, async_session, mock_redis)
|
|
|
|
with patch("cleverclaude.agents.manager.structlog.get_logger") as mock_logger:
|
|
await manager.initialize()
|
|
|
|
assert manager._initialized is True
|
|
mock_logger.assert_called_once()
|
|
|
|
async def test_create_agent_success(self, test_settings: Settings, async_session, mock_redis):
|
|
"""Test successful agent creation."""
|
|
manager = AgentManager(test_settings.agents, async_session, mock_redis)
|
|
await manager.initialize()
|
|
|
|
with patch("cleverclaude.agents.factory.AgentFactory.create_agent") as mock_factory:
|
|
mock_agent = AsyncMock()
|
|
mock_agent.agent_id = "test_agent_123"
|
|
mock_agent.name = "Test Agent"
|
|
mock_agent.agent_type = AgentType.RESEARCHER
|
|
mock_agent.status = AgentStatus.ACTIVE
|
|
mock_factory.return_value = mock_agent
|
|
|
|
agent_id = await manager.create_agent(
|
|
agent_type=AgentType.RESEARCHER, name="Test Agent", capabilities={"research", "analysis"}
|
|
)
|
|
|
|
assert agent_id == "test_agent_123"
|
|
assert agent_id in manager.agents
|
|
mock_factory.assert_called_once()
|
|
|
|
async def test_create_agent_max_limit(self, test_settings: Settings, async_session, mock_redis):
|
|
"""Test agent creation with max limit exceeded."""
|
|
# Set very low limit for testing
|
|
test_settings.agents.max_agents = 1
|
|
manager = AgentManager(test_settings.agents, async_session, mock_redis)
|
|
await manager.initialize()
|
|
|
|
# Create first agent (should succeed)
|
|
with patch("cleverclaude.agents.factory.AgentFactory.create_agent") as mock_factory:
|
|
mock_agent = AsyncMock()
|
|
mock_agent.agent_id = "test_agent_1"
|
|
mock_factory.return_value = mock_agent
|
|
|
|
agent_id_1 = await manager.create_agent(AgentType.RESEARCHER, "Agent 1")
|
|
assert agent_id_1 == "test_agent_1"
|
|
|
|
# Try to create second agent (should fail)
|
|
with pytest.raises(ValueError, match="Maximum number of agents reached"):
|
|
await manager.create_agent(AgentType.RESEARCHER, "Agent 2")
|
|
|
|
async def test_get_agent_status(self, test_settings: Settings, async_session, mock_redis, mock_agent):
|
|
"""Test getting agent status."""
|
|
manager = AgentManager(test_settings.agents, async_session, mock_redis)
|
|
await manager.initialize()
|
|
|
|
# Add mock agent to manager
|
|
manager.agents[mock_agent.agent_id] = mock_agent
|
|
|
|
status = await manager.get_agent_status(mock_agent.agent_id)
|
|
|
|
assert status["agent_id"] == mock_agent.agent_id
|
|
assert status["name"] == mock_agent.name
|
|
assert status["type"] == mock_agent.agent_type
|
|
assert status["status"] == mock_agent.status
|
|
|
|
async def test_get_agent_status_not_found(self, test_settings: Settings, async_session, mock_redis):
|
|
"""Test getting status for non-existent agent."""
|
|
manager = AgentManager(test_settings.agents, async_session, mock_redis)
|
|
await manager.initialize()
|
|
|
|
with pytest.raises(ValueError, match="Agent not found"):
|
|
await manager.get_agent_status("non_existent_agent")
|
|
|
|
async def test_list_agents(self, test_settings: Settings, async_session, mock_redis):
|
|
"""Test listing all agents."""
|
|
manager = AgentManager(test_settings.agents, async_session, mock_redis)
|
|
await manager.initialize()
|
|
|
|
# Add multiple mock agents
|
|
mock_agents = []
|
|
for i in range(3):
|
|
mock_agent = AsyncMock()
|
|
mock_agent.agent_id = f"agent_{i}"
|
|
mock_agent.name = f"Agent {i}"
|
|
mock_agent.agent_type = AgentType.RESEARCHER
|
|
mock_agent.status = AgentStatus.ACTIVE
|
|
mock_agents.append(mock_agent)
|
|
manager.agents[mock_agent.agent_id] = mock_agent
|
|
|
|
agents_list = await manager.list_agents()
|
|
|
|
assert len(agents_list) == 3
|
|
for i, agent_info in enumerate(agents_list):
|
|
assert agent_info["agent_id"] == f"agent_{i}"
|
|
assert agent_info["name"] == f"Agent {i}"
|
|
|
|
async def test_execute_task(self, test_settings: Settings, async_session, mock_redis, mock_agent, mock_task):
|
|
"""Test task execution on agent."""
|
|
manager = AgentManager(test_settings.agents, async_session, mock_redis)
|
|
await manager.initialize()
|
|
|
|
# Set up mock agent
|
|
manager.agents[mock_agent.agent_id] = mock_agent
|
|
mock_agent.execute_task.return_value = {"status": "completed", "result": "task completed"}
|
|
|
|
result = await manager.execute_task(mock_task, agent_id=mock_agent.agent_id)
|
|
|
|
assert result["status"] == "completed"
|
|
mock_agent.execute_task.assert_called_once_with(mock_task)
|
|
|
|
async def test_execute_task_auto_assign(self, test_settings: Settings, async_session, mock_redis, mock_task):
|
|
"""Test task execution with automatic agent assignment."""
|
|
manager = AgentManager(test_settings.agents, async_session, mock_redis)
|
|
await manager.initialize()
|
|
|
|
# Add multiple mock agents
|
|
for i in range(2):
|
|
mock_agent = AsyncMock()
|
|
mock_agent.agent_id = f"agent_{i}"
|
|
mock_agent.agent_type = AgentType.RESEARCHER
|
|
mock_agent.status = AgentStatus.ACTIVE
|
|
mock_agent.capabilities = {"research", "analysis"}
|
|
mock_agent.execute_task.return_value = {"status": "completed"}
|
|
manager.agents[mock_agent.agent_id] = mock_agent
|
|
|
|
with patch.object(manager, "_find_suitable_agent") as mock_find:
|
|
mock_find.return_value = "agent_0"
|
|
|
|
result = await manager.execute_task(mock_task)
|
|
|
|
assert result["status"] == "completed"
|
|
mock_find.assert_called_once_with(mock_task)
|
|
|
|
async def test_destroy_agent(self, test_settings: Settings, async_session, mock_redis, mock_agent):
|
|
"""Test agent destruction."""
|
|
manager = AgentManager(test_settings.agents, async_session, mock_redis)
|
|
await manager.initialize()
|
|
|
|
# Add mock agent
|
|
manager.agents[mock_agent.agent_id] = mock_agent
|
|
mock_agent.shutdown = AsyncMock()
|
|
|
|
await manager.destroy_agent(mock_agent.agent_id)
|
|
|
|
assert mock_agent.agent_id not in manager.agents
|
|
mock_agent.shutdown.assert_called_once()
|
|
|
|
async def test_health_check(self, test_settings: Settings, async_session, mock_redis, mock_agent):
|
|
"""Test agent health check."""
|
|
manager = AgentManager(test_settings.agents, async_session, mock_redis)
|
|
await manager.initialize()
|
|
|
|
# Add mock agent
|
|
manager.agents[mock_agent.agent_id] = mock_agent
|
|
mock_agent.get_health_status.return_value = mock_agent.health_status
|
|
|
|
health_report = await manager.health_check()
|
|
|
|
assert len(health_report["agents"]) == 1
|
|
assert health_report["agents"][0]["agent_id"] == mock_agent.agent_id
|
|
assert health_report["total_agents"] == 1
|
|
assert health_report["healthy_agents"] == 1
|
|
|
|
async def test_shutdown(self, test_settings: Settings, async_session, mock_redis):
|
|
"""Test manager shutdown."""
|
|
manager = AgentManager(test_settings.agents, async_session, mock_redis)
|
|
await manager.initialize()
|
|
|
|
# Add mock agents
|
|
mock_agents = []
|
|
for i in range(2):
|
|
mock_agent = AsyncMock()
|
|
mock_agent.agent_id = f"agent_{i}"
|
|
mock_agent.shutdown = AsyncMock()
|
|
mock_agents.append(mock_agent)
|
|
manager.agents[mock_agent.agent_id] = mock_agent
|
|
|
|
await manager.shutdown()
|
|
|
|
# Verify all agents were shut down
|
|
for mock_agent in mock_agents:
|
|
mock_agent.shutdown.assert_called_once()
|
|
|
|
assert len(manager.agents) == 0
|
|
|
|
|
|
@pytest.mark.unit
|
|
@pytest.mark.async_test
|
|
class TestBaseAgent:
|
|
"""Test suite for BaseAgent class."""
|
|
|
|
def test_agent_initialization(self):
|
|
"""Test agent initialization with config."""
|
|
config = AgentConfig(
|
|
agent_id="test_agent",
|
|
name="Test Agent",
|
|
agent_type=AgentType.RESEARCHER,
|
|
capabilities={"research", "analysis"},
|
|
timeout=300,
|
|
)
|
|
|
|
agent = BaseAgent(config)
|
|
|
|
assert agent.agent_id == "test_agent"
|
|
assert agent.name == "Test Agent"
|
|
assert agent.agent_type == AgentType.RESEARCHER
|
|
assert agent.capabilities == {"research", "analysis"}
|
|
assert agent.status == AgentStatus.INITIALIZING
|
|
assert agent.timeout == 300
|
|
|
|
async def test_agent_startup(self):
|
|
"""Test agent startup process."""
|
|
config = AgentConfig(agent_id="test_agent", name="Test Agent", agent_type=AgentType.RESEARCHER)
|
|
|
|
agent = BaseAgent(config)
|
|
await agent.startup()
|
|
|
|
assert agent.status == AgentStatus.ACTIVE
|
|
assert agent.created_at is not None
|
|
|
|
async def test_agent_execute_task(self, mock_task):
|
|
"""Test basic task execution."""
|
|
config = AgentConfig(
|
|
agent_id="test_agent", name="Test Agent", agent_type=AgentType.RESEARCHER, capabilities={"research"}
|
|
)
|
|
|
|
agent = BaseAgent(config)
|
|
await agent.startup()
|
|
|
|
with patch.object(agent, "_process_task") as mock_process:
|
|
mock_process.return_value = {"status": "completed", "result": "processed"}
|
|
|
|
result = await agent.execute_task(mock_task)
|
|
|
|
assert result["status"] == "completed"
|
|
assert agent.task_count == 1
|
|
mock_process.assert_called_once_with(mock_task)
|
|
|
|
async def test_agent_execute_task_timeout(self, mock_task):
|
|
"""Test task execution with timeout."""
|
|
config = AgentConfig(
|
|
agent_id="test_agent",
|
|
name="Test Agent",
|
|
agent_type=AgentType.RESEARCHER,
|
|
timeout=1, # Very short timeout
|
|
)
|
|
|
|
agent = BaseAgent(config)
|
|
await agent.startup()
|
|
|
|
# Mock a slow task
|
|
async def slow_task(task):
|
|
import asyncio
|
|
|
|
await asyncio.sleep(2) # Longer than timeout
|
|
return {"status": "completed"}
|
|
|
|
with patch.object(agent, "_process_task", side_effect=slow_task):
|
|
result = await agent.execute_task(mock_task)
|
|
|
|
assert result["status"] == "error"
|
|
assert "timeout" in result["error"].lower()
|
|
|
|
def test_agent_get_health_status(self):
|
|
"""Test agent health status reporting."""
|
|
config = AgentConfig(agent_id="test_agent", name="Test Agent", agent_type=AgentType.RESEARCHER)
|
|
|
|
agent = BaseAgent(config)
|
|
agent.task_count = 5
|
|
agent.error_count = 1
|
|
|
|
health = agent.get_health_status()
|
|
|
|
assert health["agent_id"] == "test_agent"
|
|
assert health["status"] == AgentStatus.INITIALIZING
|
|
assert health["task_count"] == 5
|
|
assert health["error_count"] == 1
|
|
assert "cpu_usage" in health
|
|
assert "memory_usage" in health
|
|
|
|
async def test_agent_pause_resume(self):
|
|
"""Test agent pause and resume functionality."""
|
|
config = AgentConfig(agent_id="test_agent", name="Test Agent", agent_type=AgentType.RESEARCHER)
|
|
|
|
agent = BaseAgent(config)
|
|
await agent.startup()
|
|
|
|
# Test pause
|
|
await agent.pause()
|
|
assert agent.status == AgentStatus.PAUSED
|
|
|
|
# Test resume
|
|
await agent.resume()
|
|
assert agent.status == AgentStatus.ACTIVE
|
|
|
|
async def test_agent_shutdown(self):
|
|
"""Test agent shutdown process."""
|
|
config = AgentConfig(agent_id="test_agent", name="Test Agent", agent_type=AgentType.RESEARCHER)
|
|
|
|
agent = BaseAgent(config)
|
|
await agent.startup()
|
|
|
|
await agent.shutdown()
|
|
|
|
assert agent.status == AgentStatus.TERMINATED
|
|
assert agent.shutdown_at is not None
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestAgentTypes:
|
|
"""Test suite for agent type definitions."""
|
|
|
|
def test_agent_type_enum(self):
|
|
"""Test AgentType enumeration."""
|
|
assert AgentType.RESEARCHER == "researcher"
|
|
assert AgentType.CODER == "coder"
|
|
assert AgentType.ANALYST == "analyst"
|
|
assert AgentType.COORDINATOR == "coordinator"
|
|
assert AgentType.REVIEWER == "reviewer"
|
|
assert AgentType.TESTER == "tester"
|
|
|
|
def test_agent_status_enum(self):
|
|
"""Test AgentStatus enumeration."""
|
|
assert AgentStatus.INITIALIZING == "initializing"
|
|
assert AgentStatus.ACTIVE == "active"
|
|
assert AgentStatus.BUSY == "busy"
|
|
assert AgentStatus.PAUSED == "paused"
|
|
assert AgentStatus.ERROR == "error"
|
|
assert AgentStatus.TERMINATED == "terminated"
|
|
|
|
def test_agent_config_creation(self):
|
|
"""Test AgentConfig creation and validation."""
|
|
config = AgentConfig(
|
|
agent_id="test_agent",
|
|
name="Test Agent",
|
|
agent_type=AgentType.RESEARCHER,
|
|
capabilities={"research", "analysis"},
|
|
timeout=300,
|
|
max_concurrent_tasks=5,
|
|
)
|
|
|
|
assert config.agent_id == "test_agent"
|
|
assert config.name == "Test Agent"
|
|
assert config.agent_type == AgentType.RESEARCHER
|
|
assert config.capabilities == {"research", "analysis"}
|
|
assert config.timeout == 300
|
|
assert config.max_concurrent_tasks == 5
|
|
|
|
def test_agent_config_defaults(self):
|
|
"""Test AgentConfig default values."""
|
|
config = AgentConfig(agent_id="test_agent", name="Test Agent", agent_type=AgentType.RESEARCHER)
|
|
|
|
assert config.capabilities == set()
|
|
assert config.timeout == 300 # Default timeout
|
|
assert config.max_concurrent_tasks == 1 # Default concurrency
|