Your Name
415 lines
18 KiB
Python
415 lines
18 KiB
Python
"""Unit tests for CleverClaude swarm coordination."""
|
|
|
|
from datetime import datetime, timedelta
|
|
from unittest.mock import AsyncMock, patch
|
|
|
|
import pytest
|
|
|
|
from cleverclaude.agents.types import AgentType
|
|
from cleverclaude.config.settings import Settings
|
|
from cleverclaude.coordination.swarm import SwarmCoordinator
|
|
from cleverclaude.coordination.types import SwarmConfig, SwarmState, SwarmTask, SwarmTopology, TaskPriority
|
|
|
|
|
|
@pytest.mark.unit
|
|
@pytest.mark.async_test
|
|
class TestSwarmCoordinator:
|
|
"""Test suite for SwarmCoordinator class."""
|
|
|
|
async def test_initialization(self, test_settings: Settings, async_session, mock_redis):
|
|
"""Test SwarmCoordinator initialization."""
|
|
mock_agent_manager = AsyncMock()
|
|
coordinator = SwarmCoordinator(test_settings.swarm, async_session, mock_agent_manager, mock_redis)
|
|
|
|
assert coordinator.config == test_settings.swarm
|
|
assert coordinator.session == async_session
|
|
assert coordinator.agent_manager == mock_agent_manager
|
|
assert coordinator.redis == mock_redis
|
|
assert coordinator.swarms == {}
|
|
assert coordinator._initialized is False
|
|
|
|
async def test_initialize_coordinator(self, test_settings: Settings, async_session, mock_redis):
|
|
"""Test coordinator initialization process."""
|
|
mock_agent_manager = AsyncMock()
|
|
coordinator = SwarmCoordinator(test_settings.swarm, async_session, mock_agent_manager, mock_redis)
|
|
|
|
with patch("cleverclaude.coordination.swarm.structlog.get_logger") as mock_logger:
|
|
await coordinator.initialize()
|
|
|
|
assert coordinator._initialized is True
|
|
mock_logger.assert_called_once()
|
|
|
|
async def test_create_swarm_success(self, test_settings: Settings, async_session, mock_redis):
|
|
"""Test successful swarm creation."""
|
|
mock_agent_manager = AsyncMock()
|
|
coordinator = SwarmCoordinator(test_settings.swarm, async_session, mock_agent_manager, mock_redis)
|
|
await coordinator.initialize()
|
|
|
|
swarm_id = await coordinator.create_swarm(name="Test Swarm", topology=SwarmTopology.MESH, max_agents=10)
|
|
|
|
assert swarm_id in coordinator.swarms
|
|
swarm = coordinator.swarms[swarm_id]
|
|
assert swarm["name"] == "Test Swarm"
|
|
assert swarm["topology"] == SwarmTopology.MESH
|
|
assert swarm["state"] == SwarmState.ACTIVE
|
|
assert swarm["max_agents"] == 10
|
|
|
|
async def test_create_swarm_with_agents(self, test_settings: Settings, async_session, mock_redis):
|
|
"""Test swarm creation with initial agents."""
|
|
mock_agent_manager = AsyncMock()
|
|
coordinator = SwarmCoordinator(test_settings.swarm, async_session, mock_agent_manager, mock_redis)
|
|
await coordinator.initialize()
|
|
|
|
# Mock agent creation
|
|
mock_agent_ids = ["agent_1", "agent_2", "agent_3"]
|
|
mock_agent_manager.create_agent.side_effect = mock_agent_ids
|
|
|
|
swarm_id = await coordinator.create_swarm(
|
|
name="Test Swarm",
|
|
topology=SwarmTopology.HIERARCHICAL,
|
|
initial_agents=[
|
|
{"type": AgentType.COORDINATOR, "name": "coordinator"},
|
|
{"type": AgentType.RESEARCHER, "name": "researcher_1"},
|
|
{"type": AgentType.ANALYST, "name": "analyst_1"},
|
|
],
|
|
)
|
|
|
|
swarm = coordinator.swarms[swarm_id]
|
|
assert len(swarm["agents"]) == 3
|
|
assert mock_agent_manager.create_agent.call_count == 3
|
|
|
|
async def test_add_agent_to_swarm(self, test_settings: Settings, async_session, mock_redis):
|
|
"""Test adding an agent to an existing swarm."""
|
|
mock_agent_manager = AsyncMock()
|
|
coordinator = SwarmCoordinator(test_settings.swarm, async_session, mock_agent_manager, mock_redis)
|
|
await coordinator.initialize()
|
|
|
|
# Create swarm
|
|
swarm_id = await coordinator.create_swarm("Test Swarm", SwarmTopology.MESH)
|
|
|
|
# Add agent
|
|
agent_id = "test_agent_123"
|
|
await coordinator.add_agent(swarm_id, agent_id, role="worker")
|
|
|
|
swarm = coordinator.swarms[swarm_id]
|
|
assert len(swarm["agents"]) == 1
|
|
assert swarm["agents"][0]["agent_id"] == agent_id
|
|
assert swarm["agents"][0]["role"] == "worker"
|
|
|
|
async def test_add_agent_max_limit(self, test_settings: Settings, async_session, mock_redis):
|
|
"""Test adding agent when max limit is reached."""
|
|
mock_agent_manager = AsyncMock()
|
|
coordinator = SwarmCoordinator(test_settings.swarm, async_session, mock_agent_manager, mock_redis)
|
|
await coordinator.initialize()
|
|
|
|
# Create swarm with low max limit
|
|
swarm_id = await coordinator.create_swarm("Test Swarm", SwarmTopology.MESH, max_agents=1)
|
|
|
|
# Add first agent (should succeed)
|
|
await coordinator.add_agent(swarm_id, "agent_1", role="worker")
|
|
|
|
# Try to add second agent (should fail)
|
|
with pytest.raises(ValueError, match="Maximum number of agents reached"):
|
|
await coordinator.add_agent(swarm_id, "agent_2", role="worker")
|
|
|
|
async def test_remove_agent_from_swarm(self, test_settings: Settings, async_session, mock_redis):
|
|
"""Test removing an agent from swarm."""
|
|
mock_agent_manager = AsyncMock()
|
|
coordinator = SwarmCoordinator(test_settings.swarm, async_session, mock_agent_manager, mock_redis)
|
|
await coordinator.initialize()
|
|
|
|
# Create swarm and add agent
|
|
swarm_id = await coordinator.create_swarm("Test Swarm", SwarmTopology.MESH)
|
|
agent_id = "test_agent_123"
|
|
await coordinator.add_agent(swarm_id, agent_id, role="worker")
|
|
|
|
# Remove agent
|
|
await coordinator.remove_agent(swarm_id, agent_id)
|
|
|
|
swarm = coordinator.swarms[swarm_id]
|
|
assert len(swarm["agents"]) == 0
|
|
|
|
async def test_submit_task_to_swarm(self, test_settings: Settings, async_session, mock_redis):
|
|
"""Test submitting a task to swarm."""
|
|
mock_agent_manager = AsyncMock()
|
|
mock_agent_manager.execute_task.return_value = {"status": "completed", "result": "task done"}
|
|
|
|
coordinator = SwarmCoordinator(test_settings.swarm, async_session, mock_agent_manager, mock_redis)
|
|
await coordinator.initialize()
|
|
|
|
# Create swarm with agents
|
|
swarm_id = await coordinator.create_swarm("Test Swarm", SwarmTopology.MESH)
|
|
await coordinator.add_agent(swarm_id, "agent_1", role="worker")
|
|
|
|
# Submit task
|
|
task = SwarmTask(
|
|
task_type="analysis",
|
|
priority=TaskPriority.NORMAL,
|
|
data={"analysis_type": "data_analysis", "dataset": {"records": ["data_1", "data_2"]}},
|
|
)
|
|
|
|
task_id = await coordinator.submit_task(swarm_id, task)
|
|
|
|
assert task_id in coordinator.swarms[swarm_id]["tasks"]
|
|
task_info = coordinator.swarms[swarm_id]["tasks"][task_id]
|
|
assert task_info["task_type"] == "analysis"
|
|
assert task_info["status"] == "submitted"
|
|
|
|
async def test_task_distribution_mesh(self, test_settings: Settings, async_session, mock_redis):
|
|
"""Test task distribution in mesh topology."""
|
|
mock_agent_manager = AsyncMock()
|
|
mock_agent_manager.execute_task.return_value = {"status": "completed"}
|
|
|
|
coordinator = SwarmCoordinator(test_settings.swarm, async_session, mock_agent_manager, mock_redis)
|
|
await coordinator.initialize()
|
|
|
|
# Create mesh swarm with multiple agents
|
|
swarm_id = await coordinator.create_swarm("Mesh Swarm", SwarmTopology.MESH)
|
|
for i in range(3):
|
|
await coordinator.add_agent(swarm_id, f"agent_{i}", role="worker")
|
|
|
|
# Submit multiple tasks
|
|
tasks = []
|
|
for i in range(5):
|
|
task = SwarmTask(task_type="processing", priority=TaskPriority.NORMAL, data={"task_id": i})
|
|
task_id = await coordinator.submit_task(swarm_id, task)
|
|
tasks.append(task_id)
|
|
|
|
# Process tasks
|
|
await coordinator._process_pending_tasks(swarm_id)
|
|
|
|
# Verify tasks were distributed
|
|
swarm = coordinator.swarms[swarm_id]
|
|
completed_tasks = [t for t in swarm["tasks"].values() if t["status"] == "completed"]
|
|
assert len(completed_tasks) == 5
|
|
|
|
async def test_task_distribution_hierarchical(self, test_settings: Settings, async_session, mock_redis):
|
|
"""Test task distribution in hierarchical topology."""
|
|
mock_agent_manager = AsyncMock()
|
|
mock_agent_manager.execute_task.return_value = {"status": "completed"}
|
|
|
|
coordinator = SwarmCoordinator(test_settings.swarm, async_session, mock_agent_manager, mock_redis)
|
|
await coordinator.initialize()
|
|
|
|
# Create hierarchical swarm
|
|
swarm_id = await coordinator.create_swarm("Hierarchical Swarm", SwarmTopology.HIERARCHICAL)
|
|
|
|
# Add agents with hierarchy
|
|
await coordinator.add_agent(swarm_id, "coordinator", role="coordinator", level=0)
|
|
await coordinator.add_agent(swarm_id, "team_lead_1", role="team_lead", level=1)
|
|
await coordinator.add_agent(swarm_id, "worker_1", role="worker", level=2)
|
|
await coordinator.add_agent(swarm_id, "worker_2", role="worker", level=2)
|
|
|
|
# Submit complex task
|
|
task = SwarmTask(
|
|
task_type="complex_analysis",
|
|
priority=TaskPriority.HIGH,
|
|
data={"requires_coordination": True, "subtasks": 3},
|
|
)
|
|
|
|
task_id = await coordinator.submit_task(swarm_id, task)
|
|
|
|
# Process task
|
|
await coordinator._process_pending_tasks(swarm_id)
|
|
|
|
# Verify hierarchical processing
|
|
swarm = coordinator.swarms[swarm_id]
|
|
task_info = swarm["tasks"][task_id]
|
|
assert task_info["assigned_coordinator"] == "coordinator"
|
|
|
|
async def test_get_swarm_metrics(self, test_settings: Settings, async_session, mock_redis):
|
|
"""Test swarm metrics collection."""
|
|
mock_agent_manager = AsyncMock()
|
|
coordinator = SwarmCoordinator(test_settings.swarm, async_session, mock_agent_manager, mock_redis)
|
|
await coordinator.initialize()
|
|
|
|
# Create swarm with agents and tasks
|
|
swarm_id = await coordinator.create_swarm("Metrics Swarm", SwarmTopology.MESH)
|
|
await coordinator.add_agent(swarm_id, "agent_1", role="worker")
|
|
await coordinator.add_agent(swarm_id, "agent_2", role="worker")
|
|
|
|
# Simulate completed tasks
|
|
swarm = coordinator.swarms[swarm_id]
|
|
swarm["tasks"]["task_1"] = {"status": "completed", "completed_at": datetime.utcnow() - timedelta(minutes=5)}
|
|
swarm["tasks"]["task_2"] = {"status": "completed", "completed_at": datetime.utcnow() - timedelta(minutes=3)}
|
|
swarm["tasks"]["task_3"] = {"status": "running", "started_at": datetime.utcnow() - timedelta(minutes=1)}
|
|
|
|
metrics = await coordinator.get_swarm_metrics(swarm_id)
|
|
|
|
assert metrics.total_agents == 2
|
|
assert metrics.active_agents == 2
|
|
assert metrics.completed_tasks == 2
|
|
assert metrics.running_tasks == 1
|
|
assert metrics.efficiency_score > 0
|
|
|
|
async def test_scale_swarm_up(self, test_settings: Settings, async_session, mock_redis):
|
|
"""Test scaling swarm up (adding agents)."""
|
|
mock_agent_manager = AsyncMock()
|
|
mock_agent_manager.create_agent.return_value = "new_agent_123"
|
|
|
|
coordinator = SwarmCoordinator(test_settings.swarm, async_session, mock_agent_manager, mock_redis)
|
|
await coordinator.initialize()
|
|
|
|
# Create swarm with initial agents
|
|
swarm_id = await coordinator.create_swarm("Scalable Swarm", SwarmTopology.MESH, max_agents=10)
|
|
await coordinator.add_agent(swarm_id, "agent_1", role="worker")
|
|
|
|
# Scale up
|
|
await coordinator.scale_swarm(swarm_id, target_size=3)
|
|
|
|
swarm = coordinator.swarms[swarm_id]
|
|
assert len(swarm["agents"]) == 3
|
|
assert mock_agent_manager.create_agent.call_count == 2 # 2 new agents created
|
|
|
|
async def test_scale_swarm_down(self, test_settings: Settings, async_session, mock_redis):
|
|
"""Test scaling swarm down (removing agents)."""
|
|
mock_agent_manager = AsyncMock()
|
|
|
|
coordinator = SwarmCoordinator(test_settings.swarm, async_session, mock_agent_manager, mock_redis)
|
|
await coordinator.initialize()
|
|
|
|
# Create swarm with multiple agents
|
|
swarm_id = await coordinator.create_swarm("Scalable Swarm", SwarmTopology.MESH)
|
|
for i in range(5):
|
|
await coordinator.add_agent(swarm_id, f"agent_{i}", role="worker")
|
|
|
|
# Scale down
|
|
await coordinator.scale_swarm(swarm_id, target_size=2)
|
|
|
|
swarm = coordinator.swarms[swarm_id]
|
|
assert len(swarm["agents"]) == 2
|
|
|
|
async def test_swarm_health_monitoring(self, test_settings: Settings, async_session, mock_redis):
|
|
"""Test swarm health monitoring."""
|
|
mock_agent_manager = AsyncMock()
|
|
mock_agent_manager.get_agent_status.return_value = {
|
|
"status": "active",
|
|
"health": {"cpu_usage": 50, "memory_usage": 200, "task_count": 2},
|
|
}
|
|
|
|
coordinator = SwarmCoordinator(test_settings.swarm, async_session, mock_agent_manager, mock_redis)
|
|
await coordinator.initialize()
|
|
|
|
# Create swarm with agents
|
|
swarm_id = await coordinator.create_swarm("Health Swarm", SwarmTopology.MESH)
|
|
await coordinator.add_agent(swarm_id, "agent_1", role="worker")
|
|
await coordinator.add_agent(swarm_id, "agent_2", role="worker")
|
|
|
|
# Check health
|
|
health_report = await coordinator.check_swarm_health(swarm_id)
|
|
|
|
assert health_report["swarm_id"] == swarm_id
|
|
assert health_report["total_agents"] == 2
|
|
assert len(health_report["agent_health"]) == 2
|
|
assert health_report["overall_health"] in ["healthy", "degraded", "critical"]
|
|
|
|
async def test_handle_agent_failure(self, test_settings: Settings, async_session, mock_redis):
|
|
"""Test handling agent failure in swarm."""
|
|
mock_agent_manager = AsyncMock()
|
|
|
|
coordinator = SwarmCoordinator(test_settings.swarm, async_session, mock_agent_manager, mock_redis)
|
|
await coordinator.initialize()
|
|
|
|
# Create swarm with agents
|
|
swarm_id = await coordinator.create_swarm("Resilient Swarm", SwarmTopology.MESH)
|
|
await coordinator.add_agent(swarm_id, "agent_1", role="worker")
|
|
await coordinator.add_agent(swarm_id, "agent_2", role="worker")
|
|
|
|
# Simulate agent failure
|
|
await coordinator.handle_agent_failure(swarm_id, "agent_1")
|
|
|
|
swarm = coordinator.swarms[swarm_id]
|
|
failed_agent = next((a for a in swarm["agents"] if a["agent_id"] == "agent_1"), None)
|
|
assert failed_agent["status"] == "failed"
|
|
|
|
# Verify tasks are redistributed
|
|
# This would check that tasks assigned to failed agent are reassigned
|
|
|
|
async def test_destroy_swarm(self, test_settings: Settings, async_session, mock_redis):
|
|
"""Test swarm destruction and cleanup."""
|
|
mock_agent_manager = AsyncMock()
|
|
|
|
coordinator = SwarmCoordinator(test_settings.swarm, async_session, mock_agent_manager, mock_redis)
|
|
await coordinator.initialize()
|
|
|
|
# Create swarm with agents
|
|
swarm_id = await coordinator.create_swarm("Doomed Swarm", SwarmTopology.MESH)
|
|
await coordinator.add_agent(swarm_id, "agent_1", role="worker")
|
|
await coordinator.add_agent(swarm_id, "agent_2", role="worker")
|
|
|
|
# Destroy swarm
|
|
await coordinator.destroy_swarm(swarm_id)
|
|
|
|
# Verify cleanup
|
|
assert swarm_id not in coordinator.swarms
|
|
# Verify agents were destroyed/removed
|
|
assert mock_agent_manager.destroy_agent.call_count == 2
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestSwarmTypes:
|
|
"""Test suite for swarm type definitions."""
|
|
|
|
def test_swarm_topology_enum(self):
|
|
"""Test SwarmTopology enumeration."""
|
|
assert SwarmTopology.MESH == "mesh"
|
|
assert SwarmTopology.HIERARCHICAL == "hierarchical"
|
|
assert SwarmTopology.STAR == "star"
|
|
assert SwarmTopology.RING == "ring"
|
|
|
|
def test_swarm_state_enum(self):
|
|
"""Test SwarmState enumeration."""
|
|
assert SwarmState.INITIALIZING == "initializing"
|
|
assert SwarmState.ACTIVE == "active"
|
|
assert SwarmState.IDLE == "idle"
|
|
assert SwarmState.PAUSED == "paused"
|
|
assert SwarmState.SCALING == "scaling"
|
|
assert SwarmState.TERMINATED == "terminated"
|
|
|
|
def test_task_priority_enum(self):
|
|
"""Test TaskPriority enumeration."""
|
|
assert TaskPriority.LOW == "low"
|
|
assert TaskPriority.NORMAL == "normal"
|
|
assert TaskPriority.HIGH == "high"
|
|
assert TaskPriority.CRITICAL == "critical"
|
|
|
|
def test_swarm_task_creation(self):
|
|
"""Test SwarmTask creation and validation."""
|
|
task = SwarmTask(
|
|
task_type="analysis",
|
|
priority=TaskPriority.HIGH,
|
|
data={"analysis_type": "sentiment", "dataset": "customer_reviews"},
|
|
timeout=600,
|
|
dependencies=["task_1", "task_2"],
|
|
)
|
|
|
|
assert task.task_type == "analysis"
|
|
assert task.priority == TaskPriority.HIGH
|
|
assert task.data["analysis_type"] == "sentiment"
|
|
assert task.timeout == 600
|
|
assert task.dependencies == ["task_1", "task_2"]
|
|
|
|
def test_swarm_config_creation(self):
|
|
"""Test SwarmConfig creation and validation."""
|
|
config = SwarmConfig(
|
|
name="Test Swarm",
|
|
topology=SwarmTopology.HIERARCHICAL,
|
|
max_agents=20,
|
|
coordination_timeout=120,
|
|
load_balancing=True,
|
|
)
|
|
|
|
assert config.name == "Test Swarm"
|
|
assert config.topology == SwarmTopology.HIERARCHICAL
|
|
assert config.max_agents == 20
|
|
assert config.coordination_timeout == 120
|
|
assert config.load_balancing is True
|
|
|
|
def test_swarm_config_defaults(self):
|
|
"""Test SwarmConfig default values."""
|
|
config = SwarmConfig(name="Default Swarm", topology=SwarmTopology.MESH)
|
|
|
|
assert config.max_agents == 50 # Default
|
|
assert config.coordination_timeout == 60 # Default
|
|
assert config.load_balancing is True # Default
|