Refactored test_spreadsheet.py and test_swarm_architectures.py to use pytest with real agent executionpull/1168/head
Steve-Dusty
1 week ago
parent
a9be38ea9c
commit
d105482347
@ -1,226 +1,563 @@
|
|||||||
|
"""
|
||||||
|
SpreadSheetSwarm Test Suite
|
||||||
|
|
||||||
|
Tests for the SpreadSheetSwarm class, which manages multiple agents to execute tasks
|
||||||
|
concurrently with support for CSV-based agent loading and automatic metadata tracking.
|
||||||
|
|
||||||
|
The SpreadSheetSwarm processes tasks across multiple agents in parallel, tracks outputs,
|
||||||
|
and provides data export capabilities in both CSV and JSON formats.
|
||||||
|
"""
|
||||||
|
|
||||||
import os
|
import os
|
||||||
import asyncio
|
import json
|
||||||
from loguru import logger
|
import csv
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
|
||||||
from swarms.structs.agent import Agent
|
from swarms.structs.agent import Agent
|
||||||
from swarms.structs.spreadsheet_swarm import SpreadSheetSwarm
|
from swarms.structs.spreadsheet_swarm import SpreadSheetSwarm
|
||||||
|
|
||||||
|
|
||||||
def create_test_csv() -> str:
|
@pytest.fixture
|
||||||
"""Create a test CSV file with agent configurations."""
|
def temp_workspace(tmp_path):
|
||||||
print("\nStarting creation of test CSV file")
|
"""Create a temporary workspace directory for test isolation."""
|
||||||
try:
|
workspace = tmp_path / "test_workspace"
|
||||||
csv_content = """agent_name,description,system_prompt,task
|
workspace.mkdir()
|
||||||
test_agent_1,Test Agent 1,System prompt 1,Task 1
|
return str(workspace)
|
||||||
test_agent_2,Test Agent 2,System prompt 2,Task 2"""
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def sample_csv_file(tmp_path):
|
||||||
|
"""Create a sample CSV file with agent configurations."""
|
||||||
|
csv_path = tmp_path / "test_agents.csv"
|
||||||
|
csv_content = [
|
||||||
|
["agent_name", "description", "system_prompt", "task", "model_name"],
|
||||||
|
["agent_1", "First test agent", "You are a helpful assistant. Respond with exactly 'Task completed.'", "Say hello", "gpt-4o-mini"],
|
||||||
|
["agent_2", "Second test agent", "You are a code reviewer. Respond with exactly 'Review done.'", "Review this: print('hello')", "gpt-4o-mini"],
|
||||||
|
]
|
||||||
|
|
||||||
file_path = "test_agents.csv"
|
with open(csv_path, "w", newline="") as f:
|
||||||
with open(file_path, "w") as f:
|
writer = csv.writer(f)
|
||||||
f.write(csv_content)
|
writer.writerows(csv_content)
|
||||||
|
|
||||||
print(f"Created CSV with content:\n{csv_content}")
|
return str(csv_path)
|
||||||
print(f"CSV file created at: {file_path}")
|
|
||||||
return file_path
|
|
||||||
except Exception as e:
|
|
||||||
logger.error(f"Failed to create test CSV: {str(e)}")
|
|
||||||
raise
|
|
||||||
|
|
||||||
|
|
||||||
def create_test_agent(name: str) -> Agent:
|
def test_swarm_initialization_basic(temp_workspace):
|
||||||
"""Create a test agent with specified name."""
|
"""Test basic swarm initialization with required parameters."""
|
||||||
print(f"\nCreating test agent: {name}")
|
|
||||||
try:
|
|
||||||
agent = Agent(
|
agent = Agent(
|
||||||
agent_name=name,
|
agent_name="test_agent_1",
|
||||||
system_prompt=f"Test prompt for {name}",
|
system_prompt="You are a helpful assistant",
|
||||||
model_name="gpt-4o-mini",
|
model_name="gpt-4o-mini",
|
||||||
max_loops=1,
|
max_loops=1,
|
||||||
autosave=True,
|
|
||||||
verbose=True,
|
|
||||||
)
|
)
|
||||||
print(f"Created agent: {name}")
|
|
||||||
return agent
|
|
||||||
except Exception as e:
|
|
||||||
logger.error(f"Failed to create agent {name}: {str(e)}")
|
|
||||||
raise
|
|
||||||
|
|
||||||
|
swarm = SpreadSheetSwarm(
|
||||||
|
name="Test Swarm",
|
||||||
|
description="Test swarm description",
|
||||||
|
agents=[agent],
|
||||||
|
workspace_dir=temp_workspace,
|
||||||
|
)
|
||||||
|
|
||||||
|
assert swarm.name == "Test Swarm"
|
||||||
|
assert swarm.description == "Test swarm description"
|
||||||
|
assert len(swarm.agents) == 1
|
||||||
|
assert swarm.max_loops == 1
|
||||||
|
assert swarm.autosave is True
|
||||||
|
assert swarm.tasks_completed == 0
|
||||||
|
assert swarm.outputs == []
|
||||||
|
|
||||||
def test_swarm_initialization() -> None:
|
|
||||||
"""Test basic swarm initialization."""
|
def test_swarm_initialization_multiple_agents(temp_workspace):
|
||||||
print("\n[TEST] Starting swarm initialization test")
|
"""Test swarm initialization with multiple agents."""
|
||||||
try:
|
|
||||||
print("Creating test agents...")
|
|
||||||
agents = [
|
agents = [
|
||||||
create_test_agent("agent1"),
|
Agent(
|
||||||
create_test_agent("agent2"),
|
agent_name="agent_1",
|
||||||
|
system_prompt="You are agent 1",
|
||||||
|
model_name="gpt-4o-mini",
|
||||||
|
max_loops=1,
|
||||||
|
),
|
||||||
|
Agent(
|
||||||
|
agent_name="agent_2",
|
||||||
|
system_prompt="You are agent 2",
|
||||||
|
model_name="gpt-4o-mini",
|
||||||
|
max_loops=1,
|
||||||
|
),
|
||||||
]
|
]
|
||||||
|
|
||||||
print("Initializing swarm...")
|
|
||||||
swarm = SpreadSheetSwarm(
|
swarm = SpreadSheetSwarm(
|
||||||
name="Test Swarm",
|
name="Multi Agent Swarm",
|
||||||
description="Test Description",
|
|
||||||
agents=agents,
|
agents=agents,
|
||||||
max_loops=2,
|
workspace_dir=temp_workspace,
|
||||||
)
|
)
|
||||||
|
|
||||||
print("Verifying swarm configuration...")
|
|
||||||
assert swarm.name == "Test Swarm"
|
|
||||||
assert swarm.description == "Test Description"
|
|
||||||
assert len(swarm.agents) == 2
|
assert len(swarm.agents) == 2
|
||||||
assert swarm.max_loops == 2
|
assert swarm.agents[0].agent_name == "agent_1"
|
||||||
|
assert swarm.agents[1].agent_name == "agent_2"
|
||||||
|
|
||||||
|
|
||||||
|
def test_swarm_initialization_custom_max_loops(temp_workspace):
|
||||||
|
"""Test initialization with custom max_loops setting."""
|
||||||
|
agent = Agent(
|
||||||
|
agent_name="test_agent",
|
||||||
|
system_prompt="Test prompt",
|
||||||
|
model_name="gpt-4o-mini",
|
||||||
|
max_loops=1,
|
||||||
|
)
|
||||||
|
|
||||||
print("✅ Swarm initialization test PASSED")
|
swarm = SpreadSheetSwarm(
|
||||||
except Exception as e:
|
name="Custom Loop Swarm",
|
||||||
logger.error(f"❌ Swarm initialization test FAILED: {str(e)}")
|
agents=[agent],
|
||||||
raise
|
max_loops=3,
|
||||||
|
workspace_dir=temp_workspace,
|
||||||
|
)
|
||||||
|
|
||||||
|
assert swarm.max_loops == 3
|
||||||
|
|
||||||
async def test_load_from_csv() -> None:
|
|
||||||
"""Test loading agent configurations from CSV."""
|
|
||||||
print("\n[TEST] Starting CSV loading test")
|
|
||||||
try:
|
|
||||||
csv_path = create_test_csv()
|
|
||||||
print("Initializing swarm with CSV...")
|
|
||||||
swarm = SpreadSheetSwarm(load_path=csv_path)
|
|
||||||
|
|
||||||
print("Loading configurations...")
|
def test_swarm_initialization_autosave_disabled(temp_workspace):
|
||||||
await swarm._load_from_csv()
|
"""Test initialization with autosave disabled."""
|
||||||
|
agent = Agent(
|
||||||
|
agent_name="test_agent",
|
||||||
|
system_prompt="Test prompt",
|
||||||
|
model_name="gpt-4o-mini",
|
||||||
|
max_loops=1,
|
||||||
|
)
|
||||||
|
|
||||||
print("Verifying loaded configurations...")
|
swarm = SpreadSheetSwarm(
|
||||||
assert len(swarm.agents) == 2
|
agents=[agent],
|
||||||
assert len(swarm.agent_configs) == 2
|
autosave=False,
|
||||||
assert "test_agent_1" in swarm.agent_configs
|
workspace_dir=temp_workspace,
|
||||||
assert "test_agent_2" in swarm.agent_configs
|
)
|
||||||
|
|
||||||
|
assert swarm.autosave is False
|
||||||
|
|
||||||
|
|
||||||
|
def test_swarm_save_file_path_generation(temp_workspace):
|
||||||
|
"""Test that save file path is correctly generated with workspace_dir."""
|
||||||
|
agent = Agent(
|
||||||
|
agent_name="test_agent",
|
||||||
|
system_prompt="Test prompt",
|
||||||
|
model_name="gpt-4o-mini",
|
||||||
|
max_loops=1,
|
||||||
|
)
|
||||||
|
|
||||||
|
swarm = SpreadSheetSwarm(
|
||||||
|
agents=[agent],
|
||||||
|
workspace_dir=temp_workspace,
|
||||||
|
)
|
||||||
|
|
||||||
|
assert swarm.save_file_path is not None
|
||||||
|
assert "spreadsheet_swarm_run_id_" in swarm.save_file_path
|
||||||
|
assert swarm.save_file_path.startswith(temp_workspace)
|
||||||
|
assert swarm.save_file_path.endswith(".csv")
|
||||||
|
|
||||||
|
|
||||||
|
def test_swarm_initialization_no_agents_raises_error():
|
||||||
|
"""Test that initialization without agents raises ValueError."""
|
||||||
|
with pytest.raises(ValueError, match="No agents are provided"):
|
||||||
|
SpreadSheetSwarm(agents=None)
|
||||||
|
|
||||||
|
|
||||||
|
def test_swarm_initialization_no_max_loops_raises_error():
|
||||||
|
"""Test that initialization without max_loops raises ValueError."""
|
||||||
|
agent = Agent(
|
||||||
|
agent_name="test_agent",
|
||||||
|
system_prompt="Test prompt",
|
||||||
|
model_name="gpt-4o-mini",
|
||||||
|
max_loops=1,
|
||||||
|
)
|
||||||
|
|
||||||
|
with pytest.raises(ValueError, match="No max loops are provided"):
|
||||||
|
SpreadSheetSwarm(agents=[agent], max_loops=None)
|
||||||
|
|
||||||
|
|
||||||
|
def test_track_output_single(temp_workspace):
|
||||||
|
"""Test tracking a single task output."""
|
||||||
|
agent = Agent(
|
||||||
|
agent_name="test_agent",
|
||||||
|
system_prompt="Test prompt",
|
||||||
|
model_name="gpt-4o-mini",
|
||||||
|
max_loops=1,
|
||||||
|
)
|
||||||
|
|
||||||
|
swarm = SpreadSheetSwarm(
|
||||||
|
agents=[agent],
|
||||||
|
workspace_dir=temp_workspace,
|
||||||
|
)
|
||||||
|
|
||||||
os.remove(csv_path)
|
swarm._track_output("test_agent", "Test task", "Test result")
|
||||||
print(f"Cleaned up test file: {csv_path}")
|
|
||||||
|
|
||||||
print("✅ CSV loading test PASSED")
|
assert swarm.tasks_completed == 1
|
||||||
except Exception as e:
|
assert len(swarm.outputs) == 1
|
||||||
logger.error(f"❌ CSV loading test FAILED: {str(e)}")
|
assert swarm.outputs[0]["agent_name"] == "test_agent"
|
||||||
raise
|
assert swarm.outputs[0]["task"] == "Test task"
|
||||||
|
assert swarm.outputs[0]["result"] == "Test result"
|
||||||
|
assert "timestamp" in swarm.outputs[0]
|
||||||
|
|
||||||
|
|
||||||
async def test_run_tasks() -> None:
|
def test_track_output_multiple(temp_workspace):
|
||||||
"""Test running tasks with multiple agents."""
|
"""Test tracking multiple task outputs."""
|
||||||
print("\n[TEST] Starting task execution test")
|
|
||||||
try:
|
|
||||||
print("Setting up test swarm...")
|
|
||||||
agents = [
|
agents = [
|
||||||
create_test_agent("agent1"),
|
Agent(
|
||||||
create_test_agent("agent2"),
|
agent_name=f"agent_{i}",
|
||||||
|
system_prompt="Test prompt",
|
||||||
|
model_name="gpt-4o-mini",
|
||||||
|
max_loops=1,
|
||||||
|
)
|
||||||
|
for i in range(1, 4)
|
||||||
]
|
]
|
||||||
swarm = SpreadSheetSwarm(agents=agents, max_loops=1)
|
|
||||||
|
|
||||||
test_task = "Test task for all agents"
|
swarm = SpreadSheetSwarm(
|
||||||
print(f"Running test task: {test_task}")
|
agents=agents,
|
||||||
await swarm._run_tasks(test_task)
|
workspace_dir=temp_workspace,
|
||||||
|
)
|
||||||
|
|
||||||
|
for i in range(1, 4):
|
||||||
|
swarm._track_output(f"agent_{i}", f"Task {i}", f"Result {i}")
|
||||||
|
|
||||||
|
assert swarm.tasks_completed == 3
|
||||||
|
assert len(swarm.outputs) == 3
|
||||||
|
|
||||||
|
for i, output in enumerate(swarm.outputs, 1):
|
||||||
|
assert output["agent_name"] == f"agent_{i}"
|
||||||
|
assert output["task"] == f"Task {i}"
|
||||||
|
assert output["result"] == f"Result {i}"
|
||||||
|
|
||||||
|
|
||||||
|
def test_track_output_increments_counter(temp_workspace):
|
||||||
|
"""Test that tasks_completed counter increments correctly."""
|
||||||
|
agent = Agent(
|
||||||
|
agent_name="test_agent",
|
||||||
|
system_prompt="Test prompt",
|
||||||
|
model_name="gpt-4o-mini",
|
||||||
|
max_loops=1,
|
||||||
|
)
|
||||||
|
|
||||||
|
swarm = SpreadSheetSwarm(
|
||||||
|
agents=[agent],
|
||||||
|
workspace_dir=temp_workspace,
|
||||||
|
)
|
||||||
|
|
||||||
|
initial_count = swarm.tasks_completed
|
||||||
|
swarm._track_output("test_agent", "Task 1", "Result 1")
|
||||||
|
assert swarm.tasks_completed == initial_count + 1
|
||||||
|
|
||||||
|
swarm._track_output("test_agent", "Task 2", "Result 2")
|
||||||
|
assert swarm.tasks_completed == initial_count + 2
|
||||||
|
|
||||||
|
|
||||||
|
def test_load_from_csv_basic(sample_csv_file, temp_workspace):
|
||||||
|
"""Test loading agents from a CSV file."""
|
||||||
|
# Initialize with empty agents list, will be populated from CSV
|
||||||
|
agent = Agent(
|
||||||
|
agent_name="placeholder",
|
||||||
|
system_prompt="placeholder",
|
||||||
|
model_name="gpt-4o-mini",
|
||||||
|
max_loops=1,
|
||||||
|
)
|
||||||
|
|
||||||
print("Verifying task execution...")
|
swarm = SpreadSheetSwarm(
|
||||||
assert swarm.metadata.tasks_completed == 2
|
agents=[agent],
|
||||||
assert len(swarm.metadata.outputs) == 2
|
load_path=sample_csv_file,
|
||||||
|
workspace_dir=temp_workspace,
|
||||||
|
)
|
||||||
|
|
||||||
|
swarm._load_from_csv()
|
||||||
|
|
||||||
print("✅ Task execution test PASSED")
|
# Should have loaded 2 agents from CSV plus the initial placeholder
|
||||||
except Exception as e:
|
assert len(swarm.agents) == 3
|
||||||
logger.error(f"❌ Task execution test FAILED: {str(e)}")
|
assert len(swarm.agent_tasks) == 2
|
||||||
raise
|
assert "agent_1" in swarm.agent_tasks
|
||||||
|
assert "agent_2" in swarm.agent_tasks
|
||||||
|
assert swarm.agent_tasks["agent_1"] == "Say hello"
|
||||||
|
assert swarm.agent_tasks["agent_2"] == "Review this: print('hello')"
|
||||||
|
|
||||||
|
|
||||||
def test_output_tracking() -> None:
|
def test_load_from_csv_creates_agents(sample_csv_file, temp_workspace):
|
||||||
"""Test tracking of task outputs."""
|
"""Test that CSV loading creates proper Agent objects."""
|
||||||
print("\n[TEST] Starting output tracking test")
|
agent = Agent(
|
||||||
try:
|
agent_name="placeholder",
|
||||||
print("Creating test swarm...")
|
system_prompt="placeholder",
|
||||||
swarm = SpreadSheetSwarm(agents=[create_test_agent("agent1")])
|
model_name="gpt-4o-mini",
|
||||||
|
max_loops=1,
|
||||||
|
)
|
||||||
|
|
||||||
|
swarm = SpreadSheetSwarm(
|
||||||
|
agents=[agent],
|
||||||
|
load_path=sample_csv_file,
|
||||||
|
workspace_dir=temp_workspace,
|
||||||
|
)
|
||||||
|
|
||||||
|
swarm._load_from_csv()
|
||||||
|
|
||||||
|
# Verify the loaded agents are proper Agent instances
|
||||||
|
for agent in swarm.agents[1:]: # Skip placeholder
|
||||||
|
assert isinstance(agent, Agent)
|
||||||
|
assert hasattr(agent, "agent_name")
|
||||||
|
assert hasattr(agent, "system_prompt")
|
||||||
|
|
||||||
|
|
||||||
|
def test_load_from_nonexistent_csv(temp_workspace):
|
||||||
|
"""Test loading from non-existent CSV file."""
|
||||||
|
agent = Agent(
|
||||||
|
agent_name="test_agent",
|
||||||
|
system_prompt="Test prompt",
|
||||||
|
model_name="gpt-4o-mini",
|
||||||
|
max_loops=1,
|
||||||
|
)
|
||||||
|
|
||||||
print("Tracking test output...")
|
swarm = SpreadSheetSwarm(
|
||||||
swarm._track_output("agent1", "Test task", "Test result")
|
agents=[agent],
|
||||||
|
load_path="nonexistent_file.csv",
|
||||||
|
workspace_dir=temp_workspace,
|
||||||
|
)
|
||||||
|
|
||||||
print("Verifying output tracking...")
|
# Error is caught and logged, not raised
|
||||||
assert swarm.metadata.tasks_completed == 1
|
swarm._load_from_csv()
|
||||||
assert len(swarm.metadata.outputs) == 1
|
# Should still have the original agent
|
||||||
assert swarm.metadata.outputs[0].agent_name == "agent1"
|
assert len(swarm.agents) == 1
|
||||||
|
|
||||||
print("✅ Output tracking test PASSED")
|
|
||||||
except Exception as e:
|
|
||||||
logger.error(f"❌ Output tracking test FAILED: {str(e)}")
|
|
||||||
raise
|
|
||||||
|
|
||||||
|
def test_save_to_csv_creates_file(temp_workspace):
|
||||||
|
"""Test that saving to CSV creates the output file."""
|
||||||
|
agent = Agent(
|
||||||
|
agent_name="test_agent",
|
||||||
|
system_prompt="Test prompt",
|
||||||
|
model_name="gpt-4o-mini",
|
||||||
|
max_loops=1,
|
||||||
|
)
|
||||||
|
|
||||||
async def test_save_to_csv() -> None:
|
|
||||||
"""Test saving metadata to CSV."""
|
|
||||||
print("\n[TEST] Starting CSV saving test")
|
|
||||||
try:
|
|
||||||
print("Setting up test data...")
|
|
||||||
swarm = SpreadSheetSwarm(
|
swarm = SpreadSheetSwarm(
|
||||||
agents=[create_test_agent("agent1")],
|
agents=[agent],
|
||||||
save_file_path="test_output.csv",
|
workspace_dir=temp_workspace,
|
||||||
)
|
)
|
||||||
swarm._track_output("agent1", "Test task", "Test result")
|
|
||||||
|
|
||||||
print("Saving to CSV...")
|
swarm._track_output("test_agent", "Test task", "Test result")
|
||||||
await swarm._save_to_csv()
|
swarm._save_to_csv()
|
||||||
|
|
||||||
print("Verifying file creation...")
|
|
||||||
assert os.path.exists(swarm.save_file_path)
|
assert os.path.exists(swarm.save_file_path)
|
||||||
|
|
||||||
os.remove(swarm.save_file_path)
|
|
||||||
print("Cleaned up test file")
|
|
||||||
|
|
||||||
print("✅ CSV saving test PASSED")
|
def test_save_to_csv_headers(temp_workspace):
|
||||||
except Exception as e:
|
"""Test that CSV file includes proper headers."""
|
||||||
logger.error(f"❌ CSV saving test FAILED: {str(e)}")
|
agent = Agent(
|
||||||
raise
|
agent_name="test_agent",
|
||||||
|
system_prompt="Test prompt",
|
||||||
|
model_name="gpt-4o-mini",
|
||||||
|
max_loops=1,
|
||||||
|
)
|
||||||
|
|
||||||
|
swarm = SpreadSheetSwarm(
|
||||||
|
agents=[agent],
|
||||||
|
workspace_dir=temp_workspace,
|
||||||
|
)
|
||||||
|
|
||||||
|
swarm._track_output("test_agent", "Test task", "Test result")
|
||||||
|
swarm._save_to_csv()
|
||||||
|
|
||||||
|
with open(swarm.save_file_path, "r") as f:
|
||||||
|
reader = csv.reader(f)
|
||||||
|
headers = next(reader)
|
||||||
|
assert headers == ["Run ID", "Agent Name", "Task", "Result", "Timestamp"]
|
||||||
|
|
||||||
|
|
||||||
|
def test_save_to_csv_data(temp_workspace):
|
||||||
|
"""Test that CSV file includes the tracked output data."""
|
||||||
|
agent = Agent(
|
||||||
|
agent_name="test_agent",
|
||||||
|
system_prompt="Test prompt",
|
||||||
|
model_name="gpt-4o-mini",
|
||||||
|
max_loops=1,
|
||||||
|
)
|
||||||
|
|
||||||
|
swarm = SpreadSheetSwarm(
|
||||||
|
agents=[agent],
|
||||||
|
workspace_dir=temp_workspace,
|
||||||
|
)
|
||||||
|
|
||||||
|
swarm._track_output("test_agent", "Test task", "Test result")
|
||||||
|
swarm._save_to_csv()
|
||||||
|
|
||||||
|
with open(swarm.save_file_path, "r") as f:
|
||||||
|
reader = csv.DictReader(f)
|
||||||
|
rows = list(reader)
|
||||||
|
|
||||||
|
assert len(rows) == 1
|
||||||
|
assert rows[0]["Agent Name"] == "test_agent"
|
||||||
|
assert rows[0]["Task"] == "Test task"
|
||||||
|
assert rows[0]["Result"] == "Test result"
|
||||||
|
|
||||||
|
|
||||||
|
def test_save_to_csv_appends(temp_workspace):
|
||||||
|
"""Test that multiple saves append to the same CSV file.
|
||||||
|
|
||||||
|
Note: _save_to_csv() saves ALL outputs in swarm.outputs each time,
|
||||||
|
so calling it twice will result in duplicates. This tests the actual behavior.
|
||||||
|
"""
|
||||||
|
agent = Agent(
|
||||||
|
agent_name="test_agent",
|
||||||
|
system_prompt="Test prompt",
|
||||||
|
model_name="gpt-4o-mini",
|
||||||
|
max_loops=1,
|
||||||
|
)
|
||||||
|
|
||||||
|
swarm = SpreadSheetSwarm(
|
||||||
|
agents=[agent],
|
||||||
|
workspace_dir=temp_workspace,
|
||||||
|
)
|
||||||
|
|
||||||
|
swarm._track_output("test_agent", "Task 1", "Result 1")
|
||||||
|
swarm._save_to_csv()
|
||||||
|
|
||||||
|
swarm._track_output("test_agent", "Task 2", "Result 2")
|
||||||
|
swarm._save_to_csv()
|
||||||
|
|
||||||
|
# After first save: Task 1 (1 row)
|
||||||
|
# After second save: Task 1 + Task 2 (2 more rows)
|
||||||
|
# Total: 3 rows (Task 1 appears twice, Task 2 appears once)
|
||||||
|
with open(swarm.save_file_path, "r") as f:
|
||||||
|
reader = csv.DictReader(f)
|
||||||
|
rows = list(reader)
|
||||||
|
assert len(rows) == 3
|
||||||
|
|
||||||
|
# Verify the data
|
||||||
|
assert rows[0]["Task"] == "Task 1"
|
||||||
|
assert rows[1]["Task"] == "Task 1"
|
||||||
|
assert rows[2]["Task"] == "Task 2"
|
||||||
|
|
||||||
|
|
||||||
|
def test_export_to_json_structure(temp_workspace):
|
||||||
|
"""Test that JSON export contains expected structure."""
|
||||||
|
agent = Agent(
|
||||||
|
agent_name="test_agent",
|
||||||
|
system_prompt="Test prompt",
|
||||||
|
model_name="gpt-4o-mini",
|
||||||
|
max_loops=1,
|
||||||
|
)
|
||||||
|
|
||||||
|
swarm = SpreadSheetSwarm(
|
||||||
|
name="JSON Test Swarm",
|
||||||
|
description="Testing JSON export",
|
||||||
|
agents=[agent],
|
||||||
|
workspace_dir=temp_workspace,
|
||||||
|
)
|
||||||
|
|
||||||
|
swarm._track_output("test_agent", "Test task", "Test result")
|
||||||
|
json_output = swarm.export_to_json()
|
||||||
|
data = json.loads(json_output)
|
||||||
|
|
||||||
|
assert "run_id" in data
|
||||||
|
assert "name" in data
|
||||||
|
assert "description" in data
|
||||||
|
assert "tasks_completed" in data
|
||||||
|
assert "number_of_agents" in data
|
||||||
|
assert "outputs" in data
|
||||||
|
|
||||||
|
assert data["name"] == "JSON Test Swarm"
|
||||||
|
assert data["description"] == "Testing JSON export"
|
||||||
|
assert data["tasks_completed"] == 1
|
||||||
|
assert data["number_of_agents"] == 1
|
||||||
|
|
||||||
|
|
||||||
|
def test_export_to_json_outputs(temp_workspace):
|
||||||
|
"""Test that JSON export includes all tracked outputs."""
|
||||||
|
agent = Agent(
|
||||||
|
agent_name="test_agent",
|
||||||
|
system_prompt="Test prompt",
|
||||||
|
model_name="gpt-4o-mini",
|
||||||
|
max_loops=1,
|
||||||
|
)
|
||||||
|
|
||||||
|
swarm = SpreadSheetSwarm(
|
||||||
|
agents=[agent],
|
||||||
|
workspace_dir=temp_workspace,
|
||||||
|
)
|
||||||
|
|
||||||
|
swarm._track_output("test_agent", "Task 1", "Result 1")
|
||||||
|
swarm._track_output("test_agent", "Task 2", "Result 2")
|
||||||
|
|
||||||
|
json_output = swarm.export_to_json()
|
||||||
|
data = json.loads(json_output)
|
||||||
|
|
||||||
|
assert len(data["outputs"]) == 2
|
||||||
|
assert data["outputs"][0]["agent_name"] == "test_agent"
|
||||||
|
assert data["outputs"][0]["task"] == "Task 1"
|
||||||
|
assert data["outputs"][1]["task"] == "Task 2"
|
||||||
|
|
||||||
|
|
||||||
def test_json_export() -> None:
|
def test_export_to_json_valid_format(temp_workspace):
|
||||||
"""Test JSON export functionality."""
|
"""Test that JSON export is valid JSON."""
|
||||||
print("\n[TEST] Starting JSON export test")
|
agent = Agent(
|
||||||
try:
|
agent_name="test_agent",
|
||||||
print("Creating test data...")
|
system_prompt="Test prompt",
|
||||||
swarm = SpreadSheetSwarm(agents=[create_test_agent("agent1")])
|
model_name="gpt-4o-mini",
|
||||||
swarm._track_output("agent1", "Test task", "Test result")
|
max_loops=1,
|
||||||
|
)
|
||||||
|
|
||||||
|
swarm = SpreadSheetSwarm(
|
||||||
|
agents=[agent],
|
||||||
|
workspace_dir=temp_workspace,
|
||||||
|
)
|
||||||
|
|
||||||
|
json_output = swarm.export_to_json()
|
||||||
|
data = json.loads(json_output)
|
||||||
|
assert isinstance(data, dict)
|
||||||
|
|
||||||
|
|
||||||
|
def test_export_empty_swarm_to_json(temp_workspace):
|
||||||
|
"""Test JSON export with no completed tasks."""
|
||||||
|
agent = Agent(
|
||||||
|
agent_name="test_agent",
|
||||||
|
system_prompt="Test prompt",
|
||||||
|
model_name="gpt-4o-mini",
|
||||||
|
max_loops=1,
|
||||||
|
)
|
||||||
|
|
||||||
|
swarm = SpreadSheetSwarm(
|
||||||
|
agents=[agent],
|
||||||
|
workspace_dir=temp_workspace,
|
||||||
|
)
|
||||||
|
|
||||||
print("Exporting to JSON...")
|
|
||||||
json_output = swarm.export_to_json()
|
json_output = swarm.export_to_json()
|
||||||
|
data = json.loads(json_output)
|
||||||
|
|
||||||
|
assert data["tasks_completed"] == 0
|
||||||
|
assert data["outputs"] == []
|
||||||
|
assert data["number_of_agents"] == 1
|
||||||
|
|
||||||
|
|
||||||
|
def test_reliability_check_passes(temp_workspace):
|
||||||
|
"""Test that reliability check passes with valid configuration."""
|
||||||
|
agent = Agent(
|
||||||
|
agent_name="test_agent",
|
||||||
|
system_prompt="Test prompt",
|
||||||
|
model_name="gpt-4o-mini",
|
||||||
|
max_loops=1,
|
||||||
|
)
|
||||||
|
|
||||||
|
swarm = SpreadSheetSwarm(
|
||||||
|
agents=[agent],
|
||||||
|
max_loops=1,
|
||||||
|
workspace_dir=temp_workspace,
|
||||||
|
)
|
||||||
|
|
||||||
|
assert swarm is not None
|
||||||
|
|
||||||
|
|
||||||
|
def test_reliability_check_verbose(temp_workspace):
|
||||||
|
"""Test verbose mode during initialization."""
|
||||||
|
agent = Agent(
|
||||||
|
agent_name="test_agent",
|
||||||
|
system_prompt="Test prompt",
|
||||||
|
model_name="gpt-4o-mini",
|
||||||
|
max_loops=1,
|
||||||
|
)
|
||||||
|
|
||||||
|
swarm = SpreadSheetSwarm(
|
||||||
|
agents=[agent],
|
||||||
|
verbose=True,
|
||||||
|
workspace_dir=temp_workspace,
|
||||||
|
)
|
||||||
|
|
||||||
print("Verifying JSON output...")
|
assert swarm.verbose is True
|
||||||
assert isinstance(json_output, str)
|
|
||||||
assert "run_id" in json_output
|
|
||||||
assert "tasks_completed" in json_output
|
|
||||||
|
|
||||||
print("✅ JSON export test PASSED")
|
|
||||||
except Exception as e:
|
|
||||||
logger.error(f"❌ JSON export test FAILED: {str(e)}")
|
|
||||||
raise
|
|
||||||
|
|
||||||
|
|
||||||
async def run_all_tests() -> None:
|
|
||||||
"""Run all test functions."""
|
|
||||||
print("\n" + "=" * 50)
|
|
||||||
print("Starting SpreadsheetSwarm Test Suite")
|
|
||||||
print("=" * 50 + "\n")
|
|
||||||
|
|
||||||
try:
|
|
||||||
# Run synchronous tests
|
|
||||||
print("Running synchronous tests...")
|
|
||||||
test_swarm_initialization()
|
|
||||||
test_output_tracking()
|
|
||||||
test_json_export()
|
|
||||||
|
|
||||||
# Run asynchronous tests
|
|
||||||
print("\nRunning asynchronous tests...")
|
|
||||||
await test_load_from_csv()
|
|
||||||
await test_run_tasks()
|
|
||||||
await test_save_to_csv()
|
|
||||||
|
|
||||||
print("\n🎉 All tests completed successfully!")
|
|
||||||
print("=" * 50)
|
|
||||||
except Exception as e:
|
|
||||||
logger.error(f"\n❌ Test suite failed: {str(e)}")
|
|
||||||
print("=" * 50)
|
|
||||||
raise
|
|
||||||
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
|
||||||
# Run all tests
|
|
||||||
asyncio.run(run_all_tests())
|
|
||||||
|
|||||||
Loading…
Reference in new issue