adding unit test for spreadsheet_swarm

pull/703/head
Patrick Devaney 4 months ago committed by mike dupont
parent 9c19ef31c9
commit 1d0f79c848

@ -1,14 +1,17 @@
import os import os
from datetime import datetime from datetime import datetime
from uuid import uuid4 from uuid import uuid4
# Import necessary classes from your swarm module # Import necessary classes from your swarm module
from swarms.structs.agent import Agent from swarms.structs.agent import Agent
from swarms.structs.base_swarm import BaseSwarm from swarms.structs.base_swarm import BaseSwarm
from swarms.telemetry.capture_sys_data import log_agent_data from swarms.telemetry.capture_sys_data import log_agent_data
from swarms.utils.file_processing import create_file_in_folder from swarms.utils.file_processing import create_file_in_folder
from swarms import SpreadSheetSwarm from swarms import SpreadSheetSwarm
# Ensure you have an environment variable or default workspace dir # Ensure you have an environment variable or default workspace dir
workspace_dir = os.getenv("WORKSPACE_DIR", "./workspace") workspace_dir = os.getenv("WORKSPACE_DIR", "./workspace")
def create_agents(num_agents: int): def create_agents(num_agents: int):
""" """
Create a list of agent instances. Create a list of agent instances.
@ -24,11 +27,14 @@ def create_agents(num_agents: int):
agent_name = f"Agent-{i + 1}" agent_name = f"Agent-{i + 1}"
agents.append(Agent(agent_name=agent_name)) agents.append(Agent(agent_name=agent_name))
return agents return agents
def main(): def main():
# Number of agents to create # Number of agents to create
num_agents = 5 num_agents = 5
# Create the agents # Create the agents
agents = create_agents(num_agents) agents = create_agents(num_agents)
# Initialize the swarm with agents and other configurations # Initialize the swarm with agents and other configurations
swarm = SpreadSheetSwarm( swarm = SpreadSheetSwarm(
name="Test-Swarm", name="Test-Swarm",
@ -38,28 +44,35 @@ def main():
max_loops=2, max_loops=2,
workspace_dir=workspace_dir workspace_dir=workspace_dir
) )
# Run a sample task in the swarm (synchronously) # Run a sample task in the swarm (synchronously)
task = "process_data" task = "process_data"
# Ensure the run method is synchronous # Ensure the run method is synchronous
swarm_metadata = swarm.run(task) # Assuming this is made synchronous swarm_metadata = swarm.run(task) # Assuming this is made synchronous
# Print swarm metadata after task completion # Print swarm metadata after task completion
print("Swarm Metadata:") print("Swarm Metadata:")
print(swarm_metadata) print(swarm_metadata)
# Check if CSV file has been created and saved # Check if CSV file has been created and saved
if os.path.exists(swarm.save_file_path): if os.path.exists(swarm.save_file_path):
print(f"Metadata saved to: {swarm.save_file_path}") print(f"Metadata saved to: {swarm.save_file_path}")
else: else:
print(f"Metadata not saved correctly. Check the save path.") print(f"Metadata not saved correctly. Check the save path.")
# Test saving metadata to JSON file # Test saving metadata to JSON file
swarm.data_to_json_file() swarm.data_to_json_file()
# Test exporting metadata to JSON # Test exporting metadata to JSON
swarm_json = swarm.export_to_json() swarm_json = swarm.export_to_json()
print("Exported JSON metadata:") print("Exported JSON metadata:")
print(swarm_json) print(swarm_json)
# Log agent data # Log agent data
print("Logging agent data:") print("Logging agent data:")
print(log_agent_data(swarm.metadata.model_dump())) print(log_agent_data(swarm.metadata.model_dump()))
# Run the synchronous main function # Run the synchronous main function
if __name__ == "__main__": if __name__ == "__main__":
main() main()
Loading…
Cancel
Save