import unittest
from unittest.mock import patch
from swarms import create_agents_from_yaml
import os
class TestCreateAgentsFromYaml(unittest.TestCase):
def setUp(self):
# Mock the environment variable for API key
os.environ['OPENAI_API_KEY'] = 'fake-api-key'
# Mock agent configuration YAML content
self.valid_yaml_content = """
agents:
- agent_name: "Financial-Analysis-Agent"
model:
openai_api_key: "fake-api-key"
model_name: "gpt-4o-mini"
temperature: 0.1
max_tokens: 2000
system_prompt: "financial_agent_sys_prompt"
max_loops: 1
autosave: true
dashboard: false
verbose: true
dynamic_temperature_enabled: true
saved_state_path: "finance_agent.json"
user_name: "swarms_corp"
retry_attempts: 1
context_length: 200000
return_step_meta: false
output_type: "str"
task: "How can I establish a ROTH IRA to buy stocks and get a tax break?"
- agent_name: "Stock-Analysis-Agent"
model:
openai_api_key: "fake-api-key"
model_name: "gpt-4o-mini"
temperature: 0.2
max_tokens: 1500
system_prompt: "stock_agent_sys_prompt"
max_loops: 2
autosave: true
dashboard: false
verbose: true
dynamic_temperature_enabled: false
saved_state_path: "stock_agent.json"
user_name: "stock_user"
retry_attempts: 3
context_length: 150000
return_step_meta: true
output_type: "json"
task: "What is the best strategy for long-term stock investment?"
"""
@patch('builtins.open', new_callable=unittest.mock.mock_open, read_data="")
@patch('yaml.safe_load')
def test_create_agents_return_agents(self, mock_safe_load, mock_open):
# Mock YAML content parsing
mock_safe_load.return_value = {
"agents": [
{
"agent_name": "Financial-Analysis-Agent",
"model": {
"openai_api_key": "fake-api-key",
"model_name": "gpt-4o-mini",
"temperature": 0.1,
"max_tokens": 2000
},
"system_prompt": "financial_agent_sys_prompt",
"max_loops": 1,
"autosave": True,
"dashboard": False,
"verbose": True,
"dynamic_temperature_enabled": True,
"saved_state_path": "finance_agent.json",
"user_name": "swarms_corp",
"retry_attempts": 1,
"context_length": 200000,
"return_step_meta": False,
"output_type": "str",
"task": "How can I establish a ROTH IRA to buy stocks and get a tax break?"
}
]
}
# Test if agents are returned correctly
agents = create_agents_from_yaml('fake_yaml_path.yaml', return_type="agents")
self.assertEqual(len(agents), 1)
self.assertEqual(agents[0].agent_name, "Financial-Analysis-Agent")
@patch('builtins.open', new_callable=unittest.mock.mock_open, read_data="")
@patch('yaml.safe_load')
@patch('swarms.Agent.run', return_value="Task completed successfully")
def test_create_agents_return_tasks(self, mock_agent_run, mock_safe_load, mock_open):
# Mock YAML content parsing
mock_safe_load.return_value = {
"agents": [
{
"agent_name": "Financial-Analysis-Agent",
"model": {
"openai_api_key": "fake-api-key",
"model_name": "gpt-4o-mini",
"temperature": 0.1,
"max_tokens": 2000
},
"system_prompt": "financial_agent_sys_prompt",
"max_loops": 1,
"autosave": True,
"dashboard": False,
"verbose": True,
"dynamic_temperature_enabled": True,
"saved_state_path": "finance_agent.json",
"user_name": "swarms_corp",
"retry_attempts": 1,
"context_length": 200000,
"return_step_meta": False,
"output_type": "str",
"task": "How can I establish a ROTH IRA to buy stocks and get a tax break?"
}
]
}
# Test if tasks are executed and results are returned
task_results = create_agents_from_yaml('fake_yaml_path.yaml', return_type="tasks")
self.assertEqual(len(task_results), 1)
self.assertEqual(task_results[0]['agent_name'], "Financial-Analysis-Agent")
self.assertIsNotNone(task_results[0]['output'])
@patch('builtins.open', new_callable=unittest.mock.mock_open, read_data="")
@patch('yaml.safe_load')
def test_create_agents_return_both(self, mock_safe_load, mock_open):
# Mock YAML content parsing
mock_safe_load.return_value = {
"agents": [
{
"agent_name": "Financial-Analysis-Agent",
"model": {
"openai_api_key": "fake-api-key",
"model_name": "gpt-4o-mini",
"temperature": 0.1,
"max_tokens": 2000
},
"system_prompt": "financial_agent_sys_prompt",
"max_loops": 1,
"autosave": True,
"dashboard": False,
"verbose": True,
"dynamic_temperature_enabled": True,
"saved_state_path": "finance_agent.json",
"user_name": "swarms_corp",
"retry_attempts": 1,
"context_length": 200000,
"return_step_meta": False,
"output_type": "str",
"task": "How can I establish a ROTH IRA to buy stocks and get a tax break?"
}
]
}
# Test if both agents and tasks are returned
agents, task_results = create_agents_from_yaml('fake_yaml_path.yaml', return_type="both")
self.assertEqual(len(agents), 1)
self.assertEqual(len(task_results), 1)
self.assertEqual(agents[0].agent_name, "Financial-Analysis-Agent")
self.assertIsNotNone(task_results[0]['output'])
@patch('builtins.open', new_callable=unittest.mock.mock_open, read_data="")
@patch('yaml.safe_load')
def test_missing_agents_in_yaml(self, mock_safe_load, mock_open):
# Mock YAML content with missing "agents" key
mock_safe_load.return_value = {}
# Test if the function raises an error for missing "agents" key
with self.assertRaises(ValueError) as context:
create_agents_from_yaml('fake_yaml_path.yaml', return_type="agents")
self.assertTrue("The YAML configuration does not contain 'agents'." in str(context.exception))
@patch('builtins.open', new_callable=unittest.mock.mock_open, read_data="")
@patch('yaml.safe_load')
def test_invalid_return_type(self, mock_safe_load, mock_open):
# Mock YAML content parsing
mock_safe_load.return_value = {
"agents": [
{
"agent_name": "Financial-Analysis-Agent",
"model": {
"openai_api_key": "fake-api-key",
"model_name": "gpt-4o-mini",
"temperature": 0.1,
"max_tokens": 2000
},
"system_prompt": "financial_agent_sys_prompt",
"max_loops": 1,
"autosave": True,
"dashboard": False,
"verbose": True,
"dynamic_temperature_enabled": True,
"saved_state_path": "finance_agent.json",
"user_name": "swarms_corp",
"retry_attempts": 1,
"context_length": 200000,
"return_step_meta": False,
"output_type": "str",
"task": "How can I establish a ROTH IRA to buy stocks and get a tax break?"
}
]
}
# Test if an error is raised for invalid return_type
with self.assertRaises(ValueError) as context:
create_agents_from_yaml('fake_yaml_path.yaml', return_type="invalid_type")
self.assertTrue("Invalid return_type" in str(context.exception))
if __name__ == '__main__':
unittest.main()