AutoGPT

1
from __future__ import annotations
2

3
import os
4
import shutil
5
import sys
6
from pathlib import Path
7
from typing import TYPE_CHECKING
8

9
import pytest
10

11
if TYPE_CHECKING:
12
    from autogpt.agents import Agent, BaseAgent
13

14
from autogpt.core.utils.json_schema import JSONSchema
15
from autogpt.models.command import Command, CommandParameter
16
from autogpt.models.command_registry import CommandRegistry
17

18
PARAMETERS = [
19
    CommandParameter(
20
        "arg1",
21
        spec=JSONSchema(
22
            type=JSONSchema.Type.INTEGER,
23
            description="Argument 1",
24
            required=True,
25
        ),
26
    ),
27
    CommandParameter(
28
        "arg2",
29
        spec=JSONSchema(
30
            type=JSONSchema.Type.STRING,
31
            description="Argument 2",
32
            required=False,
33
        ),
34
    ),
35
]
36

37

38
def example_command_method(arg1: int, arg2: str, agent: BaseAgent) -> str:
39
    """Example function for testing the Command class."""
40
    # This function is static because it is not used by any other test cases.
41
    return f"{arg1} - {arg2}"
42

43

44
def test_command_creation():
45
    """Test that a Command object can be created with the correct attributes."""
46
    cmd = Command(
47
        name="example",
48
        description="Example command",
49
        method=example_command_method,
50
        parameters=PARAMETERS,
51
    )
52

53
    assert cmd.name == "example"
54
    assert cmd.description == "Example command"
55
    assert cmd.method == example_command_method
56
    assert (
57
        str(cmd)
58
        == "example: Example command. Params: (arg1: integer, arg2: Optional[string])"
59
    )
60

61

62
@pytest.fixture
63
def example_command():
64
    yield Command(
65
        name="example",
66
        description="Example command",
67
        method=example_command_method,
68
        parameters=PARAMETERS,
69
    )
70

71

72
def test_command_call(example_command: Command, agent: Agent):
73
    """Test that Command(*args) calls and returns the result of method(*args)."""
74
    result = example_command(arg1=1, arg2="test", agent=agent)
75
    assert result == "1 - test"
76

77

78
def test_command_call_with_invalid_arguments(example_command: Command, agent: Agent):
79
    """Test that calling a Command object with invalid arguments raises a TypeError."""
80
    with pytest.raises(TypeError):
81
        example_command(arg1="invalid", does_not_exist="test", agent=agent)
82

83

84
def test_register_command(example_command: Command):
85
    """Test that a command can be registered to the registry."""
86
    registry = CommandRegistry()
87

88
    registry.register(example_command)
89

90
    assert registry.get_command(example_command.name) == example_command
91
    assert len(registry.commands) == 1
92

93

94
def test_unregister_command(example_command: Command):
95
    """Test that a command can be unregistered from the registry."""
96
    registry = CommandRegistry()
97

98
    registry.register(example_command)
99
    registry.unregister(example_command)
100

101
    assert len(registry.commands) == 0
102
    assert example_command.name not in registry
103

104

105
@pytest.fixture
106
def example_command_with_aliases(example_command: Command):
107
    example_command.aliases = ["example_alias", "example_alias_2"]
108
    return example_command
109

110

111
def test_register_command_aliases(example_command_with_aliases: Command):
112
    """Test that a command can be registered to the registry."""
113
    registry = CommandRegistry()
114
    command = example_command_with_aliases
115

116
    registry.register(command)
117

118
    assert command.name in registry
119
    assert registry.get_command(command.name) == command
120
    for alias in command.aliases:
121
        assert registry.get_command(alias) == command
122
    assert len(registry.commands) == 1
123

124

125
def test_unregister_command_aliases(example_command_with_aliases: Command):
126
    """Test that a command can be unregistered from the registry."""
127
    registry = CommandRegistry()
128
    command = example_command_with_aliases
129

130
    registry.register(command)
131
    registry.unregister(command)
132

133
    assert len(registry.commands) == 0
134
    assert command.name not in registry
135
    for alias in command.aliases:
136
        assert alias not in registry
137

138

139
def test_command_in_registry(example_command_with_aliases: Command):
140
    """Test that `command_name in registry` works."""
141
    registry = CommandRegistry()
142
    command = example_command_with_aliases
143

144
    assert command.name not in registry
145
    assert "nonexistent_command" not in registry
146

147
    registry.register(command)
148

149
    assert command.name in registry
150
    assert "nonexistent_command" not in registry
151
    for alias in command.aliases:
152
        assert alias in registry
153

154

155
def test_get_command(example_command: Command):
156
    """Test that a command can be retrieved from the registry."""
157
    registry = CommandRegistry()
158

159
    registry.register(example_command)
160
    retrieved_cmd = registry.get_command(example_command.name)
161

162
    assert retrieved_cmd == example_command
163

164

165
def test_get_nonexistent_command():
166
    """Test that attempting to get a nonexistent command raises a KeyError."""
167
    registry = CommandRegistry()
168

169
    assert registry.get_command("nonexistent_command") is None
170
    assert "nonexistent_command" not in registry
171

172

173
def test_call_command(agent: Agent):
174
    """Test that a command can be called through the registry."""
175
    registry = CommandRegistry()
176
    cmd = Command(
177
        name="example",
178
        description="Example command",
179
        method=example_command_method,
180
        parameters=PARAMETERS,
181
    )
182

183
    registry.register(cmd)
184
    result = registry.call("example", arg1=1, arg2="test", agent=agent)
185

186
    assert result == "1 - test"
187

188

189
def test_call_nonexistent_command(agent: Agent):
190
    """Test that attempting to call a nonexistent command raises a KeyError."""
191
    registry = CommandRegistry()
192

193
    with pytest.raises(KeyError):
194
        registry.call("nonexistent_command", arg1=1, arg2="test", agent=agent)
195

196

197
def test_import_mock_commands_module():
198
    """Test that the registry can import a module with mock command plugins."""
199
    registry = CommandRegistry()
200
    mock_commands_module = "tests.mocks.mock_commands"
201

202
    registry.import_command_module(mock_commands_module)
203

204
    assert "function_based_cmd" in registry
205
    assert registry.commands["function_based_cmd"].name == "function_based_cmd"
206
    assert (
207
        registry.commands["function_based_cmd"].description
208
        == "Function-based test command"
209
    )
210

211

212
def test_import_temp_command_file_module(tmp_path: Path):
213
    """
214
    Test that the registry can import a command plugins module from a temp file.
215
    Args:
216
        tmp_path (pathlib.Path): Path to a temporary directory.
217
    """
218
    registry = CommandRegistry()
219

220
    # Create a temp command file
221
    src = Path(os.getcwd()) / "tests/mocks/mock_commands.py"
222
    temp_commands_file = tmp_path / "mock_commands.py"
223
    shutil.copyfile(src, temp_commands_file)
224

225
    # Add the temp directory to sys.path to make the module importable
226
    sys.path.append(str(tmp_path))
227

228
    temp_commands_module = "mock_commands"
229
    registry.import_command_module(temp_commands_module)
230

231
    # Remove the temp directory from sys.path
232
    sys.path.remove(str(tmp_path))
233

234
    assert "function_based_cmd" in registry
235
    assert registry.commands["function_based_cmd"].name == "function_based_cmd"
236
    assert (
237
        registry.commands["function_based_cmd"].description
238
        == "Function-based test command"
239
    )
240