adding unit test for spreadsheet_swarm fixing issues with linting tests ruff more linting tests linting fixes and QA on spreadsheet_swarm.py spreadsheet_swarm.py confirmed working on windows in tests with Harshal multi platform spreadsheet_swarm, deps build clusterops from github source spreadsheet and syntax fix to auto_test_eval.py spreadsheet and finance tests confirmed to work align repo with upstreampull/710/head
Patrick Devaney
2 weeks ago
committed by
mike dupont
parent
5365915ddd
commit
8258483017
@ -0,0 +1,253 @@
|
||||
import os
|
||||
import json
|
||||
import asyncio
|
||||
from typing import Optional, Dict, Any, List, Tuple
|
||||
from datetime import datetime
|
||||
from loguru import logger
|
||||
from swarms import Agent
|
||||
from swarms.prompts.finance_agent_sys_prompt import FINANCIAL_AGENT_SYS_PROMPT
|
||||
|
||||
# Configure Loguru logger
|
||||
logger.remove() # Remove default handler
|
||||
logger.add(
|
||||
"financial_agent_tests_{time}.log",
|
||||
rotation="1 day",
|
||||
retention="7 days",
|
||||
level="DEBUG",
|
||||
format="{time:YYYY-MM-DD HH:mm:ss} | {level} | {message}",
|
||||
)
|
||||
|
||||
class FinancialAgentTestSuite:
|
||||
"""
|
||||
Production-grade test suite for Financial Analysis Agent.
|
||||
|
||||
This test suite provides comprehensive testing of the Financial Analysis Agent's
|
||||
functionality, including initialization, configuration, and response validation.
|
||||
|
||||
Attributes:
|
||||
test_data_path (str): Path to store test data and outputs
|
||||
agent_config (Dict[str, Any]): Default configuration for test agents
|
||||
"""
|
||||
|
||||
def __init__(self, test_data_path: str = "./test_data"):
|
||||
"""
|
||||
Initialize the test suite with configuration and setup.
|
||||
|
||||
Args:
|
||||
test_data_path (str): Directory to store test data and outputs
|
||||
"""
|
||||
self.test_data_path = test_data_path
|
||||
self.agent_config = {
|
||||
"agent_name": "Test-Financial-Analysis-Agent",
|
||||
"system_prompt": FINANCIAL_AGENT_SYS_PROMPT,
|
||||
"model_name": "gpt-4o-mini",
|
||||
"max_loops": 1,
|
||||
"autosave": True,
|
||||
"dashboard": False,
|
||||
"verbose": True,
|
||||
"dynamic_temperature_enabled": True,
|
||||
"saved_state_path": "test_finance_agent.json",
|
||||
"user_name": "test_user",
|
||||
"retry_attempts": 1,
|
||||
"context_length": 200000,
|
||||
"return_step_meta": False,
|
||||
"output_type": "string",
|
||||
"streaming_on": False,
|
||||
}
|
||||
self._setup_test_environment()
|
||||
|
||||
def _setup_test_environment(self) -> None:
|
||||
"""Create necessary directories and files for testing."""
|
||||
try:
|
||||
os.makedirs(self.test_data_path, exist_ok=True)
|
||||
logger.info(f"Test environment setup completed at {self.test_data_path}")
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to setup test environment: {str(e)}")
|
||||
raise
|
||||
|
||||
async def _create_test_agent(self, config_override: Optional[Dict[str, Any]] = None) -> Agent:
|
||||
"""
|
||||
Create a test agent with specified or default configuration.
|
||||
|
||||
Args:
|
||||
config_override (Optional[Dict[str, Any]]): Override default config values
|
||||
|
||||
Returns:
|
||||
Agent: Configured test agent instance
|
||||
"""
|
||||
try:
|
||||
test_config = self.agent_config.copy()
|
||||
if config_override:
|
||||
test_config.update(config_override)
|
||||
|
||||
agent = Agent(**test_config)
|
||||
logger.debug(f"Created test agent with config: {test_config}")
|
||||
return agent
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to create test agent: {str(e)}")
|
||||
raise
|
||||
|
||||
async def test_agent_initialization(self) -> Tuple[bool, str]:
|
||||
"""
|
||||
Test agent initialization with various configurations.
|
||||
|
||||
Returns:
|
||||
Tuple[bool, str]: Success status and result message
|
||||
"""
|
||||
try:
|
||||
logger.info("Starting agent initialization test")
|
||||
|
||||
# Test default initialization
|
||||
agent = await self._create_test_agent()
|
||||
assert isinstance(agent, Agent), "Agent initialization failed"
|
||||
|
||||
# Test with modified configuration
|
||||
custom_config = {"max_loops": 2, "context_length": 150000}
|
||||
agent_custom = await self._create_test_agent(custom_config)
|
||||
assert agent_custom.max_loops == 2, "Custom configuration not applied"
|
||||
|
||||
logger.info("Agent initialization test passed")
|
||||
return True, "Agent initialization successful"
|
||||
except Exception as e:
|
||||
logger.error(f"Agent initialization test failed: {str(e)}")
|
||||
return False, f"Agent initialization failed: {str(e)}"
|
||||
|
||||
async def test_agent_response(self) -> Tuple[bool, str]:
|
||||
"""
|
||||
Test agent's response functionality with various queries.
|
||||
|
||||
Returns:
|
||||
Tuple[bool, str]: Success status and result message
|
||||
"""
|
||||
try:
|
||||
logger.info("Starting agent response test")
|
||||
agent = await self._create_test_agent()
|
||||
|
||||
test_queries = [
|
||||
"How can I establish a ROTH IRA?",
|
||||
"What are the tax implications of stock trading?",
|
||||
"Explain mutual fund investment strategies"
|
||||
]
|
||||
|
||||
for query in test_queries:
|
||||
response = agent.run(query)
|
||||
assert isinstance(response, str), "Response type mismatch"
|
||||
assert len(response) > 0, "Empty response received"
|
||||
logger.debug(f"Query: {query[:50]}... | Response length: {len(response)}")
|
||||
|
||||
logger.info("Agent response test passed")
|
||||
return True, "Agent response test successful"
|
||||
except Exception as e:
|
||||
logger.error(f"Agent response test failed: {str(e)}")
|
||||
return False, f"Agent response test failed: {str(e)}"
|
||||
|
||||
async def test_agent_persistence(self) -> Tuple[bool, str]:
|
||||
"""
|
||||
Test agent's state persistence and recovery.
|
||||
|
||||
Returns:
|
||||
Tuple[bool, str]: Success status and result message
|
||||
"""
|
||||
try:
|
||||
logger.info("Starting agent persistence test")
|
||||
|
||||
# Test state saving
|
||||
save_path = os.path.join(self.test_data_path, "test_state.json")
|
||||
agent = await self._create_test_agent({"saved_state_path": save_path})
|
||||
|
||||
test_query = "What is a 401k plan?"
|
||||
agent.run(test_query)
|
||||
|
||||
assert os.path.exists(save_path), "State file not created"
|
||||
|
||||
# Verify state content
|
||||
with open(save_path, 'r') as f:
|
||||
saved_state = json.load(f)
|
||||
assert "agent_name" in saved_state, "Invalid state file content"
|
||||
|
||||
logger.info("Agent persistence test passed")
|
||||
return True, "Agent persistence test successful"
|
||||
except Exception as e:
|
||||
logger.error(f"Agent persistence test failed: {str(e)}")
|
||||
return False, f"Agent persistence test failed: {str(e)}"
|
||||
|
||||
async def run_all_tests(self) -> Dict[str, Any]:
|
||||
"""
|
||||
Run all test cases and generate a comprehensive report.
|
||||
|
||||
Returns:
|
||||
Dict[str, Any]: Test results and statistics
|
||||
"""
|
||||
start_time = datetime.now()
|
||||
results = []
|
||||
|
||||
test_cases = [
|
||||
("Agent Initialization", self.test_agent_initialization),
|
||||
("Agent Response", self.test_agent_response),
|
||||
("Agent Persistence", self.test_agent_persistence)
|
||||
]
|
||||
|
||||
for test_name, test_func in test_cases:
|
||||
try:
|
||||
success, message = await test_func()
|
||||
results.append({
|
||||
"test_name": test_name,
|
||||
"success": success,
|
||||
"message": message,
|
||||
"timestamp": datetime.now().isoformat()
|
||||
})
|
||||
except Exception as e:
|
||||
logger.error(f"Test {test_name} failed with unexpected error: {str(e)}")
|
||||
results.append({
|
||||
"test_name": test_name,
|
||||
"success": False,
|
||||
"message": f"Unexpected error: {str(e)}",
|
||||
"timestamp": datetime.now().isoformat()
|
||||
})
|
||||
|
||||
end_time = datetime.now()
|
||||
duration = (end_time - start_time).total_seconds()
|
||||
|
||||
# Generate report
|
||||
total_tests = len(results)
|
||||
passed_tests = sum(1 for r in results if r["success"])
|
||||
|
||||
report = {
|
||||
"summary": {
|
||||
"total_tests": total_tests,
|
||||
"passed_tests": passed_tests,
|
||||
"failed_tests": total_tests - passed_tests,
|
||||
"success_rate": f"{(passed_tests/total_tests)*100:.2f}%",
|
||||
"duration_seconds": duration
|
||||
},
|
||||
"test_results": results,
|
||||
"timestamp": datetime.now().isoformat()
|
||||
}
|
||||
|
||||
# Save report
|
||||
report_path = os.path.join(self.test_data_path, f"test_report_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json")
|
||||
with open(report_path, 'w') as f:
|
||||
json.dump(report, f, indent=2)
|
||||
|
||||
logger.info(f"Test suite completed. Report saved to {report_path}")
|
||||
return report
|
||||
|
||||
async def main():
|
||||
"""Main entry point for running the test suite."""
|
||||
logger.info("Starting Financial Agent Test Suite")
|
||||
test_suite = FinancialAgentTestSuite()
|
||||
report = await test_suite.run_all_tests()
|
||||
|
||||
# Print summary to console
|
||||
print("\n" + "="*50)
|
||||
print("Financial Agent Test Suite Results")
|
||||
print("="*50)
|
||||
print(f"Total Tests: {report['summary']['total_tests']}")
|
||||
print(f"Passed Tests: {report['summary']['passed_tests']}")
|
||||
print(f"Failed Tests: {report['summary']['failed_tests']}")
|
||||
print(f"Success Rate: {report['summary']['success_rate']}")
|
||||
print(f"Duration: {report['summary']['duration_seconds']:.2f} seconds")
|
||||
print("="*50)
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
@ -0,0 +1,65 @@
|
||||
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()
|
||||
Loading…
Reference in new issue