AutoGPT

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

5
import pytest
6
import pytest_asyncio
7
from google.auth.exceptions import GoogleAuthError
8
from google.cloud import storage
9
from google.cloud.exceptions import NotFound
10

11
from autogpt.file_storage.gcs import GCSFileStorage, GCSFileStorageConfiguration
12

13
try:
14
    storage.Client()
15
except GoogleAuthError:
16
    pytest.skip("Google Cloud Authentication not configured", allow_module_level=True)
17

18

19
@pytest.fixture(scope="module")
20
def gcs_bucket_name() -> str:
21
    return f"test-bucket-{str(uuid.uuid4())[:8]}"
22

23

24
@pytest.fixture(scope="module")
25
def gcs_root() -> Path:
26
    return Path("/workspaces/AutoGPT-some-unique-task-id")
27

28

29
@pytest.fixture(scope="module")
30
def gcs_storage_uninitialized(gcs_bucket_name: str, gcs_root: Path) -> GCSFileStorage:
31
    os.environ["STORAGE_BUCKET"] = gcs_bucket_name
32
    storage_config = GCSFileStorageConfiguration.from_env()
33
    storage_config.root = gcs_root
34
    storage = GCSFileStorage(storage_config)
35
    yield storage  # type: ignore
36
    del os.environ["STORAGE_BUCKET"]
37

38

39
def test_initialize(gcs_bucket_name: str, gcs_storage_uninitialized: GCSFileStorage):
40
    gcs = gcs_storage_uninitialized._gcs
41

42
    # test that the bucket doesn't exist yet
43
    with pytest.raises(NotFound):
44
        gcs.get_bucket(gcs_bucket_name)
45

46
    gcs_storage_uninitialized.initialize()
47

48
    # test that the bucket has been created
49
    bucket = gcs.get_bucket(gcs_bucket_name)
50

51
    # clean up
52
    bucket.delete(force=True)
53

54

55
@pytest.fixture(scope="module")
56
def gcs_storage(gcs_storage_uninitialized: GCSFileStorage) -> GCSFileStorage:
57
    (gcs_storage := gcs_storage_uninitialized).initialize()
58
    yield gcs_storage  # type: ignore
59

60
    # Empty & delete the test bucket
61
    gcs_storage._bucket.delete(force=True)
62

63

64
def test_workspace_bucket_name(
65
    gcs_storage: GCSFileStorage,
66
    gcs_bucket_name: str,
67
):
68
    assert gcs_storage._bucket.name == gcs_bucket_name
69

70

71
NESTED_DIR = "existing/test/dir"
72
TEST_FILES: list[tuple[str | Path, str]] = [
73
    ("existing_test_file_1", "test content 1"),
74
    ("existing_test_file_2.txt", "test content 2"),
75
    (Path("existing_test_file_3"), "test content 3"),
76
    (Path(f"{NESTED_DIR}/test_file_4"), "test content 4"),
77
]
78

79

80
@pytest_asyncio.fixture
81
async def gcs_storage_with_files(gcs_storage: GCSFileStorage) -> GCSFileStorage:
82
    for file_name, file_content in TEST_FILES:
83
        gcs_storage._bucket.blob(
84
            str(gcs_storage.get_path(file_name))
85
        ).upload_from_string(file_content)
86
    yield gcs_storage  # type: ignore
87

88

89
@pytest.mark.asyncio
90
async def test_read_file(gcs_storage_with_files: GCSFileStorage):
91
    for file_name, file_content in TEST_FILES:
92
        content = gcs_storage_with_files.read_file(file_name)
93
        assert content == file_content
94

95
    with pytest.raises(NotFound):
96
        gcs_storage_with_files.read_file("non_existent_file")
97

98

99
def test_list_files(gcs_storage_with_files: GCSFileStorage):
100
    # List at root level
101
    assert (
102
        files := gcs_storage_with_files.list_files()
103
    ) == gcs_storage_with_files.list_files()
104
    assert len(files) > 0
105
    assert set(files) == set(Path(file_name) for file_name, _ in TEST_FILES)
106

107
    # List at nested path
108
    assert (
109
        nested_files := gcs_storage_with_files.list_files(NESTED_DIR)
110
    ) == gcs_storage_with_files.list_files(NESTED_DIR)
111
    assert len(nested_files) > 0
112
    assert set(nested_files) == set(
113
        p.relative_to(NESTED_DIR)
114
        for file_name, _ in TEST_FILES
115
        if (p := Path(file_name)).is_relative_to(NESTED_DIR)
116
    )
117

118

119
def test_list_folders(gcs_storage_with_files: GCSFileStorage):
120
    # List recursive
121
    folders = gcs_storage_with_files.list_folders(recursive=True)
122
    assert len(folders) > 0
123
    assert set(folders) == {
124
        Path("existing"),
125
        Path("existing/test"),
126
        Path("existing/test/dir"),
127
    }
128
    # List non-recursive
129
    folders = gcs_storage_with_files.list_folders(recursive=False)
130
    assert len(folders) > 0
131
    assert set(folders) == {Path("existing")}
132

133

134
@pytest.mark.asyncio
135
async def test_write_read_file(gcs_storage: GCSFileStorage):
136
    await gcs_storage.write_file("test_file", "test_content")
137
    assert gcs_storage.read_file("test_file") == "test_content"
138

139

140
@pytest.mark.asyncio
141
async def test_overwrite_file(gcs_storage_with_files: GCSFileStorage):
142
    for file_name, _ in TEST_FILES:
143
        await gcs_storage_with_files.write_file(file_name, "new content")
144
        assert gcs_storage_with_files.read_file(file_name) == "new content"
145

146

147
def test_delete_file(gcs_storage_with_files: GCSFileStorage):
148
    for file_to_delete, _ in TEST_FILES:
149
        gcs_storage_with_files.delete_file(file_to_delete)
150
        assert not gcs_storage_with_files.exists(file_to_delete)
151

152

153
def test_exists(gcs_storage_with_files: GCSFileStorage):
154
    for file_name, _ in TEST_FILES:
155
        assert gcs_storage_with_files.exists(file_name)
156

157
    assert not gcs_storage_with_files.exists("non_existent_file")
158

159

160
def test_rename_file(gcs_storage_with_files: GCSFileStorage):
161
    for file_name, _ in TEST_FILES:
162
        new_name = str(file_name) + "_renamed"
163
        gcs_storage_with_files.rename(file_name, new_name)
164
        assert gcs_storage_with_files.exists(new_name)
165
        assert not gcs_storage_with_files.exists(file_name)
166

167

168
def test_rename_dir(gcs_storage_with_files: GCSFileStorage):
169
    gcs_storage_with_files.rename(NESTED_DIR, "existing/test/dir_renamed")
170
    assert gcs_storage_with_files.exists("existing/test/dir_renamed")
171
    assert not gcs_storage_with_files.exists(NESTED_DIR)
172

173

174
def test_clone(gcs_storage_with_files: GCSFileStorage, gcs_root: Path):
175
    cloned = gcs_storage_with_files.clone_with_subroot("existing/test")
176
    assert cloned.root == gcs_root / Path("existing/test")
177
    assert cloned._bucket.name == gcs_storage_with_files._bucket.name
178
    assert cloned.exists("dir")
179
    assert cloned.exists("dir/test_file_4")
180