Skip to content

Commit

Permalink
adding unit test for spreadsheet_swarm
Browse files Browse the repository at this point in the history
  • Loading branch information
patrickbdevaney committed Dec 23, 2024
1 parent 6aa2522 commit a359c05
Show file tree
Hide file tree
Showing 2 changed files with 87 additions and 3 deletions.
12 changes: 9 additions & 3 deletions swarms/structs/spreadsheet_swarm.py
Original file line number Diff line number Diff line change
Expand Up @@ -150,8 +150,14 @@ def run(self, task: str, *args, **kwargs):
logger.info(f"Running the swarm with task: {task}")
self.metadata.start_time = time

# Run the asyncio event loop
asyncio.run(self._run_tasks(task, *args, **kwargs))
# Check if we're already in an event loop
if asyncio.get_event_loop().is_running():
# If so, create and run tasks directly using `create_task` without `asyncio.run`
task_future = asyncio.create_task(self._run_tasks(task, *args, **kwargs))
asyncio.get_event_loop().run_until_complete(task_future)
else:
# If no event loop is running, run using `asyncio.run`
asyncio.run(self._run_tasks(task, *args, **kwargs))

self.metadata.end_time = time

Expand Down Expand Up @@ -298,4 +304,4 @@ async def _save_to_csv(self):
output.result,
output.timestamp,
]
)
)
78 changes: 78 additions & 0 deletions tests/structs/test_spreadsheet_swarm.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,78 @@
import os
from datetime import datetime
from uuid import uuid4

# Import necessary classes from your swarm module
from swarms.structs.agent import Agent
from swarms.structs.base_swarm import BaseSwarm
from swarms.telemetry.capture_sys_data import log_agent_data
from swarms.utils.file_processing import create_file_in_folder
from swarms import SpreadSheetSwarm

# Ensure you have an environment variable or default workspace dir
workspace_dir = os.getenv("WORKSPACE_DIR", "./workspace")

def create_agents(num_agents: int):
"""
Create a list of agent instances.
Args:
num_agents (int): The number of agents to create.
Returns:
List[Agent]: List of created Agent objects.
"""
agents = []
for i in range(num_agents):
agent_name = f"Agent-{i + 1}"
agents.append(Agent(agent_name=agent_name))
return agents

def main():
# Number of agents to create
num_agents = 5

# Create the agents
agents = create_agents(num_agents)

# Initialize the swarm with agents and other configurations
swarm = SpreadSheetSwarm(
name="Test-Swarm",
description="A swarm for testing purposes.",
agents=agents,
autosave_on=True,
max_loops=2,
workspace_dir=workspace_dir
)

# Run a sample task in the swarm (synchronously)
task = "process_data"

# Ensure the run method is synchronous
swarm_metadata = swarm.run(task) # Assuming this is made synchronous

# Print swarm metadata after task completion
print("Swarm Metadata:")
print(swarm_metadata)

# Check if CSV file has been created and saved
if os.path.exists(swarm.save_file_path):
print(f"Metadata saved to: {swarm.save_file_path}")
else:
print(f"Metadata not saved correctly. Check the save path.")

# Test saving metadata to JSON file
swarm.data_to_json_file()

# Test exporting metadata to JSON
swarm_json = swarm.export_to_json()
print("Exported JSON metadata:")
print(swarm_json)

# Log agent data
print("Logging agent data:")
print(log_agent_data(swarm.metadata.model_dump()))

# Run the synchronous main function
if __name__ == "__main__":
main()

0 comments on commit a359c05

Please sign in to comment.