AutoGPT

1
import json
2
import os
3
from pathlib import Path
4
from unittest.mock import patch
5

6
import pytest
7
import requests
8
from git import InvalidGitRepositoryError
9

10
import autogpt.app.utils
11
from autogpt.app.utils import (
12
    get_bulletin_from_web,
13
    get_current_git_branch,
14
    get_latest_bulletin,
15
    set_env_config_value,
16
)
17
from autogpt.core.utils.json_utils import extract_dict_from_json
18
from autogpt.utils import validate_yaml_file
19
from tests.utils import skip_in_ci
20

21

22
@pytest.fixture
23
def valid_json_response() -> dict:
24
    return {
25
        "thoughts": {
26
            "text": "My task is complete. I will use the 'task_complete' command "
27
            "to shut down.",
28
            "reasoning": "I will use the 'task_complete' command because it allows me "
29
            "to shut down and signal that my task is complete.",
30
            "plan": "I will use the 'task_complete' command with the reason "
31
            "'Task complete: retrieved Tesla's revenue in 2022.' to shut down.",
32
            "criticism": "I need to ensure that I have completed all necessary tasks "
33
            "before shutting down.",
34
            "speak": "All done!",
35
        },
36
        "command": {
37
            "name": "task_complete",
38
            "args": {"reason": "Task complete: retrieved Tesla's revenue in 2022."},
39
        },
40
    }
41

42

43
@pytest.fixture
44
def invalid_json_response() -> dict:
45
    return {
46
        "thoughts": {
47
            "text": "My task is complete. I will use the 'task_complete' command "
48
            "to shut down.",
49
            "reasoning": "I will use the 'task_complete' command because it allows me "
50
            "to shut down and signal that my task is complete.",
51
            "plan": "I will use the 'task_complete' command with the reason "
52
            "'Task complete: retrieved Tesla's revenue in 2022.' to shut down.",
53
            "criticism": "I need to ensure that I have completed all necessary tasks "
54
            "before shutting down.",
55
            "speak": "",
56
        },
57
        "command": {"name": "", "args": {}},
58
    }
59

60

61
def test_validate_yaml_file_valid():
62
    with open("valid_test_file.yaml", "w") as f:
63
        f.write("setting: value")
64
    result, message = validate_yaml_file("valid_test_file.yaml")
65
    os.remove("valid_test_file.yaml")
66

67
    assert result is True
68
    assert "Successfully validated" in message
69

70

71
def test_validate_yaml_file_not_found():
72
    result, message = validate_yaml_file("non_existent_file.yaml")
73

74
    assert result is False
75
    assert "wasn't found" in message
76

77

78
def test_validate_yaml_file_invalid():
79
    with open("invalid_test_file.yaml", "w") as f:
80
        f.write(
81
            "settings:\n"
82
            "  first_setting: value\n"
83
            "  second_setting: value\n"
84
            "    nested_setting: value\n"
85
            "  third_setting: value\n"
86
            "unindented_setting: value"
87
        )
88
    result, message = validate_yaml_file("invalid_test_file.yaml")
89
    os.remove("invalid_test_file.yaml")
90
    print(result)
91
    print(message)
92
    assert result is False
93
    assert "There was an issue while trying to read" in message
94

95

96
@patch("requests.get")
97
def test_get_bulletin_from_web_success(mock_get):
98
    expected_content = "Test bulletin from web"
99

100
    mock_get.return_value.status_code = 200
101
    mock_get.return_value.text = expected_content
102
    bulletin = get_bulletin_from_web()
103

104
    assert expected_content in bulletin
105
    mock_get.assert_called_with(
106
        "https://raw.githubusercontent.com/Significant-Gravitas/AutoGPT/master/autogpts/autogpt/BULLETIN.md"  # noqa: E501
107
    )
108

109

110
@patch("requests.get")
111
def test_get_bulletin_from_web_failure(mock_get):
112
    mock_get.return_value.status_code = 404
113
    bulletin = get_bulletin_from_web()
114

115
    assert bulletin == ""
116

117

118
@patch("requests.get")
119
def test_get_bulletin_from_web_exception(mock_get):
120
    mock_get.side_effect = requests.exceptions.RequestException()
121
    bulletin = get_bulletin_from_web()
122

123
    assert bulletin == ""
124

125

126
def test_get_latest_bulletin_no_file():
127
    if os.path.exists("data/CURRENT_BULLETIN.md"):
128
        os.remove("data/CURRENT_BULLETIN.md")
129

130
    bulletin, is_new = get_latest_bulletin()
131
    assert is_new
132

133

134
def test_get_latest_bulletin_with_file():
135
    expected_content = "Test bulletin"
136
    with open("data/CURRENT_BULLETIN.md", "w", encoding="utf-8") as f:
137
        f.write(expected_content)
138

139
    with patch("autogpt.app.utils.get_bulletin_from_web", return_value=""):
140
        bulletin, is_new = get_latest_bulletin()
141
        assert expected_content in bulletin
142
        assert is_new is False
143

144
    os.remove("data/CURRENT_BULLETIN.md")
145

146

147
def test_get_latest_bulletin_with_new_bulletin():
148
    with open("data/CURRENT_BULLETIN.md", "w", encoding="utf-8") as f:
149
        f.write("Old bulletin")
150

151
    expected_content = "New bulletin from web"
152
    with patch(
153
        "autogpt.app.utils.get_bulletin_from_web", return_value=expected_content
154
    ):
155
        bulletin, is_new = get_latest_bulletin()
156
        assert "::NEW BULLETIN::" in bulletin
157
        assert expected_content in bulletin
158
        assert is_new
159

160
    os.remove("data/CURRENT_BULLETIN.md")
161

162

163
def test_get_latest_bulletin_new_bulletin_same_as_old_bulletin():
164
    expected_content = "Current bulletin"
165
    with open("data/CURRENT_BULLETIN.md", "w", encoding="utf-8") as f:
166
        f.write(expected_content)
167

168
    with patch(
169
        "autogpt.app.utils.get_bulletin_from_web", return_value=expected_content
170
    ):
171
        bulletin, is_new = get_latest_bulletin()
172
        assert expected_content in bulletin
173
        assert is_new is False
174

175
    os.remove("data/CURRENT_BULLETIN.md")
176

177

178
@skip_in_ci
179
def test_get_current_git_branch():
180
    branch_name = get_current_git_branch()
181
    assert branch_name != ""
182

183

184
@patch("autogpt.app.utils.Repo")
185
def test_get_current_git_branch_success(mock_repo):
186
    mock_repo.return_value.active_branch.name = "test-branch"
187
    branch_name = get_current_git_branch()
188

189
    assert branch_name == "test-branch"
190

191

192
@patch("autogpt.app.utils.Repo")
193
def test_get_current_git_branch_failure(mock_repo):
194
    mock_repo.side_effect = InvalidGitRepositoryError()
195
    branch_name = get_current_git_branch()
196

197
    assert branch_name == ""
198

199

200
def test_extract_json_from_response(valid_json_response: dict):
201
    emulated_response_from_openai = json.dumps(valid_json_response)
202
    assert extract_dict_from_json(emulated_response_from_openai) == valid_json_response
203

204

205
def test_extract_json_from_response_wrapped_in_code_block(valid_json_response: dict):
206
    emulated_response_from_openai = "```" + json.dumps(valid_json_response) + "```"
207
    assert extract_dict_from_json(emulated_response_from_openai) == valid_json_response
208

209

210
def test_extract_json_from_response_wrapped_in_code_block_with_language(
211
    valid_json_response: dict,
212
):
213
    emulated_response_from_openai = "```json" + json.dumps(valid_json_response) + "```"
214
    assert extract_dict_from_json(emulated_response_from_openai) == valid_json_response
215

216

217
def test_extract_json_from_response_json_contained_in_string(valid_json_response: dict):
218
    emulated_response_from_openai = (
219
        "sentence1" + json.dumps(valid_json_response) + "sentence2"
220
    )
221
    assert extract_dict_from_json(emulated_response_from_openai) == valid_json_response
222

223

224
@pytest.fixture
225
def mock_env_file_path(tmp_path):
226
    return tmp_path / ".env"
227

228

229
env_file_initial_content = """
230
# This is a comment
231
EXISTING_KEY=EXISTING_VALUE
232

233
## This is also a comment
234
# DISABLED_KEY=DISABLED_VALUE
235

236
# Another comment
237
UNUSED_KEY=UNUSED_VALUE
238
"""
239

240

241
@pytest.fixture
242
def mock_env_file(mock_env_file_path: Path, monkeypatch: pytest.MonkeyPatch):
243
    mock_env_file_path.write_text(env_file_initial_content)
244
    monkeypatch.setattr(autogpt.app.utils, "ENV_FILE_PATH", mock_env_file_path)
245
    return mock_env_file_path
246

247

248
@pytest.fixture
249
def mock_environ(monkeypatch: pytest.MonkeyPatch):
250
    env = {}
251
    monkeypatch.setattr(os, "environ", env)
252
    return env
253

254

255
def test_set_env_config_value_updates_existing_key(
256
    mock_env_file: Path, mock_environ: dict
257
):
258
    # Before updating, ensure the original content is as expected
259
    with mock_env_file.open("r") as file:
260
        assert file.readlines() == env_file_initial_content.splitlines(True)
261

262
    set_env_config_value("EXISTING_KEY", "NEW_VALUE")
263
    with mock_env_file.open("r") as file:
264
        content = file.readlines()
265

266
    # Ensure only the relevant line is altered
267
    expected_content_lines = [
268
        "\n",
269
        "# This is a comment\n",
270
        "EXISTING_KEY=NEW_VALUE\n",  # existing key + new value
271
        "\n",
272
        "## This is also a comment\n",
273
        "# DISABLED_KEY=DISABLED_VALUE\n",
274
        "\n",
275
        "# Another comment\n",
276
        "UNUSED_KEY=UNUSED_VALUE\n",
277
    ]
278
    assert content == expected_content_lines
279
    assert mock_environ["EXISTING_KEY"] == "NEW_VALUE"
280

281

282
def test_set_env_config_value_uncomments_and_updates_disabled_key(
283
    mock_env_file: Path, mock_environ: dict
284
):
285
    # Before adding, ensure the original content is as expected
286
    with mock_env_file.open("r") as file:
287
        assert file.readlines() == env_file_initial_content.splitlines(True)
288

289
    set_env_config_value("DISABLED_KEY", "ENABLED_NEW_VALUE")
290
    with mock_env_file.open("r") as file:
291
        content = file.readlines()
292

293
    # Ensure only the relevant line is altered
294
    expected_content_lines = [
295
        "\n",
296
        "# This is a comment\n",
297
        "EXISTING_KEY=EXISTING_VALUE\n",
298
        "\n",
299
        "## This is also a comment\n",
300
        "DISABLED_KEY=ENABLED_NEW_VALUE\n",  # disabled -> enabled + new value
301
        "\n",
302
        "# Another comment\n",
303
        "UNUSED_KEY=UNUSED_VALUE\n",
304
    ]
305
    assert content == expected_content_lines
306
    assert mock_environ["DISABLED_KEY"] == "ENABLED_NEW_VALUE"
307

308

309
def test_set_env_config_value_adds_new_key(mock_env_file: Path, mock_environ: dict):
310
    # Before adding, ensure the original content is as expected
311
    with mock_env_file.open("r") as file:
312
        assert file.readlines() == env_file_initial_content.splitlines(True)
313

314
    set_env_config_value("NEW_KEY", "NEW_VALUE")
315
    with mock_env_file.open("r") as file:
316
        content = file.readlines()
317

318
    # Ensure the new key-value pair is added without altering the rest
319
    expected_content_lines = [
320
        "\n",
321
        "# This is a comment\n",
322
        "EXISTING_KEY=EXISTING_VALUE\n",
323
        "\n",
324
        "## This is also a comment\n",
325
        "# DISABLED_KEY=DISABLED_VALUE\n",
326
        "\n",
327
        "# Another comment\n",
328
        "UNUSED_KEY=UNUSED_VALUE\n",
329
        "NEW_KEY=NEW_VALUE\n",  # New key-value pair added at the end
330
    ]
331
    assert content == expected_content_lines
332
    assert mock_environ["NEW_KEY"] == "NEW_VALUE"
333