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()