AutoGPT

1
import os
2
import uuid
3
from pathlib import Path
4

5
import pytest
6
import pytest_asyncio
7
from botocore.exceptions import ClientError
8

9
from autogpt.file_storage.s3 import S3FileStorage, S3FileStorageConfiguration
10

11
if not (os.getenv("S3_ENDPOINT_URL") and os.getenv("AWS_ACCESS_KEY_ID")):
12
    pytest.skip("S3 environment variables are not set", allow_module_level=True)
13

14

15
@pytest.fixture
16
def s3_bucket_name() -> str:
17
    return f"test-bucket-{str(uuid.uuid4())[:8]}"
18

19

20
@pytest.fixture
21
def s3_root() -> Path:
22
    return Path("/workspaces/AutoGPT-some-unique-task-id")
23

24

25
@pytest.fixture
26
def s3_storage_uninitialized(s3_bucket_name: str, s3_root: Path) -> S3FileStorage:
27
    os.environ["STORAGE_BUCKET"] = s3_bucket_name
28
    storage_config = S3FileStorageConfiguration.from_env()
29
    storage_config.root = s3_root
30
    storage = S3FileStorage(storage_config)
31
    yield storage  # type: ignore
32
    del os.environ["STORAGE_BUCKET"]
33

34

35
def test_initialize(s3_bucket_name: str, s3_storage_uninitialized: S3FileStorage):
36
    s3 = s3_storage_uninitialized._s3
37

38
    # test that the bucket doesn't exist yet
39
    with pytest.raises(ClientError):
40
        s3.meta.client.head_bucket(Bucket=s3_bucket_name)
41

42
    s3_storage_uninitialized.initialize()
43

44
    # test that the bucket has been created
45
    s3.meta.client.head_bucket(Bucket=s3_bucket_name)
46

47

48
def test_workspace_bucket_name(
49
    s3_storage: S3FileStorage,
50
    s3_bucket_name: str,
51
):
52
    assert s3_storage._bucket.name == s3_bucket_name
53

54

55
@pytest.fixture
56
def s3_storage(s3_storage_uninitialized: S3FileStorage) -> S3FileStorage:
57
    (s3_storage := s3_storage_uninitialized).initialize()
58
    yield s3_storage  # type: ignore
59

60
    # Empty & delete the test bucket
61
    s3_storage._bucket.objects.all().delete()
62
    s3_storage._bucket.delete()
63

64

65
NESTED_DIR = "existing/test/dir"
66
TEST_FILES: list[tuple[str | Path, str]] = [
67
    ("existing_test_file_1", "test content 1"),
68
    ("existing_test_file_2.txt", "test content 2"),
69
    (Path("existing_test_file_3"), "test content 3"),
70
    (Path(f"{NESTED_DIR}/test_file_4"), "test content 4"),
71
]
72

73

74
@pytest_asyncio.fixture
75
async def s3_storage_with_files(s3_storage: S3FileStorage) -> S3FileStorage:
76
    for file_name, file_content in TEST_FILES:
77
        s3_storage._bucket.Object(str(s3_storage.get_path(file_name))).put(
78
            Body=file_content
79
        )
80
    yield s3_storage  # type: ignore
81

82

83
@pytest.mark.asyncio
84
async def test_read_file(s3_storage_with_files: S3FileStorage):
85
    for file_name, file_content in TEST_FILES:
86
        content = s3_storage_with_files.read_file(file_name)
87
        assert content == file_content
88

89
    with pytest.raises(ClientError):
90
        s3_storage_with_files.read_file("non_existent_file")
91

92

93
def test_list_files(s3_storage_with_files: S3FileStorage):
94
    # List at root level
95
    assert (
96
        files := s3_storage_with_files.list_files()
97
    ) == s3_storage_with_files.list_files()
98
    assert len(files) > 0
99
    assert set(files) == set(Path(file_name) for file_name, _ in TEST_FILES)
100

101
    # List at nested path
102
    assert (
103
        nested_files := s3_storage_with_files.list_files(NESTED_DIR)
104
    ) == s3_storage_with_files.list_files(NESTED_DIR)
105
    assert len(nested_files) > 0
106
    assert set(nested_files) == set(
107
        p.relative_to(NESTED_DIR)
108
        for file_name, _ in TEST_FILES
109
        if (p := Path(file_name)).is_relative_to(NESTED_DIR)
110
    )
111

112

113
def test_list_folders(s3_storage_with_files: S3FileStorage):
114
    # List recursive
115
    folders = s3_storage_with_files.list_folders(recursive=True)
116
    assert len(folders) > 0
117
    assert set(folders) == {
118
        Path("existing"),
119
        Path("existing/test"),
120
        Path("existing/test/dir"),
121
    }
122
    # List non-recursive
123
    folders = s3_storage_with_files.list_folders(recursive=False)
124
    assert len(folders) > 0
125
    assert set(folders) == {Path("existing")}
126

127

128
@pytest.mark.asyncio
129
async def test_write_read_file(s3_storage: S3FileStorage):
130
    await s3_storage.write_file("test_file", "test_content")
131
    assert s3_storage.read_file("test_file") == "test_content"
132

133

134
@pytest.mark.asyncio
135
async def test_overwrite_file(s3_storage_with_files: S3FileStorage):
136
    for file_name, _ in TEST_FILES:
137
        await s3_storage_with_files.write_file(file_name, "new content")
138
        assert s3_storage_with_files.read_file(file_name) == "new content"
139

140

141
def test_delete_file(s3_storage_with_files: S3FileStorage):
142
    for file_to_delete, _ in TEST_FILES:
143
        s3_storage_with_files.delete_file(file_to_delete)
144
        with pytest.raises(ClientError):
145
            s3_storage_with_files.read_file(file_to_delete)
146

147

148
def test_exists(s3_storage_with_files: S3FileStorage):
149
    for file_name, _ in TEST_FILES:
150
        assert s3_storage_with_files.exists(file_name)
151

152
    assert not s3_storage_with_files.exists("non_existent_file")
153

154

155
def test_rename_file(s3_storage_with_files: S3FileStorage):
156
    for file_name, _ in TEST_FILES:
157
        new_name = str(file_name) + "_renamed"
158
        s3_storage_with_files.rename(file_name, new_name)
159
        assert s3_storage_with_files.exists(new_name)
160
        assert not s3_storage_with_files.exists(file_name)
161

162

163
def test_rename_dir(s3_storage_with_files: S3FileStorage):
164
    s3_storage_with_files.rename(NESTED_DIR, "existing/test/dir_renamed")
165
    assert s3_storage_with_files.exists("existing/test/dir_renamed")
166
    assert not s3_storage_with_files.exists(NESTED_DIR)
167

168

169
def test_clone(s3_storage_with_files: S3FileStorage, s3_root: Path):
170
    cloned = s3_storage_with_files.clone_with_subroot("existing/test")
171
    assert cloned.root == s3_root / Path("existing/test")
172
    assert cloned._bucket.name == s3_storage_with_files._bucket.name
173
    assert cloned.exists("dir")
174
    assert cloned.exists("dir/test_file_4")
175