You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
268 lines
9.3 KiB
268 lines
9.3 KiB
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()
|