coding-alt/AutoGPT

1from unittest import TestCase2 3from autogpt.promptgenerator import PromptGenerator4 5 6class TestPromptGenerator(TestCase):7    """8    Test cases for the PromptGenerator class, which is responsible for generating9    prompts for the AI with constraints, commands, resources, and performance evaluations.10    """11 12    @classmethod13    def setUpClass(cls):14        """15        Set up the initial state for each test method by creating an instance of PromptGenerator.16        """17        cls.generator = PromptGenerator()18 19    # Test whether the add_constraint() method adds a constraint to the generator's constraints list20    def test_add_constraint(self):21        """22        Test if the add_constraint() method adds a constraint to the generator's constraints list.23        """24        constraint = "Constraint1"25        self.generator.add_constraint(constraint)26        self.assertIn(constraint, self.generator.constraints)27 28    # Test whether the add_command() method adds a command to the generator's commands list29    def test_add_command(self):30        """31        Test if the add_command() method adds a command to the generator's commands list.32        """33        command_label = "Command Label"34        command_name = "command_name"35        args = {"arg1": "value1", "arg2": "value2"}36        self.generator.add_command(command_label, command_name, args)37        command = {38            "label": command_label,39            "name": command_name,40            "args": args,41        }42        self.assertIn(command, self.generator.commands)43 44    def test_add_resource(self):45        """46        Test if the add_resource() method adds a resource to the generator's resources list.47        """48        resource = "Resource1"49        self.generator.add_resource(resource)50        self.assertIn(resource, self.generator.resources)51 52    def test_add_performance_evaluation(self):53        """54        Test if the add_performance_evaluation() method adds an evaluation to the generator's55        performance_evaluation list.56        """57        evaluation = "Evaluation1"58        self.generator.add_performance_evaluation(evaluation)59        self.assertIn(evaluation, self.generator.performance_evaluation)60 61    def test_generate_prompt_string(self):62        """63        Test if the generate_prompt_string() method generates a prompt string with all the added64        constraints, commands, resources, and evaluations.65        """66        # Define the test data67        constraints = ["Constraint1", "Constraint2"]68        commands = [69            {70                "label": "Command1",71                "name": "command_name1",72                "args": {"arg1": "value1"},73            },74            {75                "label": "Command2",76                "name": "command_name2",77                "args": {},78            },79        ]80        resources = ["Resource1", "Resource2"]81        evaluations = ["Evaluation1", "Evaluation2"]82 83        # Add test data to the generator84        for constraint in constraints:85            self.generator.add_constraint(constraint)86        for command in commands:87            self.generator.add_command(88                command["label"], command["name"], command["args"]89            )90        for resource in resources:91            self.generator.add_resource(resource)92        for evaluation in evaluations:93            self.generator.add_performance_evaluation(evaluation)94 95        # Generate the prompt string and verify its correctness96        prompt_string = self.generator.generate_prompt_string()97        self.assertIsNotNone(prompt_string)98 99        # Check if all constraints, commands, resources, and evaluations are present in the prompt string100        for constraint in constraints:101            self.assertIn(constraint, prompt_string)102        for command in commands:103            self.assertIn(command["name"], prompt_string)104            for key, value in command["args"].items():105                self.assertIn(f'"{key}": "{value}"', prompt_string)106        for resource in resources:107            self.assertIn(resource, prompt_string)108        for evaluation in evaluations:109            self.assertIn(evaluation, prompt_string)110 111        self.assertIn("constraints", prompt_string.lower())112        self.assertIn("commands", prompt_string.lower())113        self.assertIn("resources", prompt_string.lower())114        self.assertIn("performance evaluation", prompt_string.lower())115