swarms/tests/agents/test_create_agents_from_yaml.py

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