mirror of
https://github.com/aljazceru/Auto-GPT.git
synced 2025-12-22 08:24:26 +01:00
Add workspace abstraction (#2982)
* Add workspace abstraction * Remove old workspace implementation * Extract path resolution to a helper function * Add api key requirements to new tests
This commit is contained in:
@@ -4,7 +4,6 @@ import unittest
|
||||
from pathlib import Path
|
||||
|
||||
from autogpt.commands.file_operations import (
|
||||
LOG_FILE_PATH,
|
||||
append_to_file,
|
||||
check_duplicate_operation,
|
||||
delete_file,
|
||||
@@ -15,7 +14,7 @@ from autogpt.commands.file_operations import (
|
||||
write_to_file,
|
||||
)
|
||||
from autogpt.config import Config
|
||||
from autogpt.workspace import path_in_workspace
|
||||
from autogpt.workspace import Workspace
|
||||
|
||||
|
||||
class TestFileOperations(unittest.TestCase):
|
||||
@@ -24,24 +23,24 @@ class TestFileOperations(unittest.TestCase):
|
||||
"""
|
||||
|
||||
def setUp(self):
|
||||
self.test_file = "test_file.txt"
|
||||
self.config = Config()
|
||||
workspace_path = os.path.join(os.path.dirname(__file__), "workspace")
|
||||
self.workspace_path = Workspace.make_workspace(workspace_path)
|
||||
self.config.workspace_path = workspace_path
|
||||
self.config.file_logger_path = os.path.join(workspace_path, "file_logger.txt")
|
||||
self.workspace = Workspace(workspace_path, restrict_to_workspace=True)
|
||||
|
||||
self.test_file = str(self.workspace.get_path("test_file.txt"))
|
||||
self.test_file2 = "test_file2.txt"
|
||||
self.test_directory = "test_directory"
|
||||
self.test_directory = str(self.workspace.get_path("test_directory"))
|
||||
self.file_content = "This is a test file.\n"
|
||||
self.file_logger_logs = "file_logger.txt"
|
||||
|
||||
with open(path_in_workspace(self.test_file), "w") as f:
|
||||
with open(self.test_file, "w") as f:
|
||||
f.write(self.file_content)
|
||||
|
||||
if os.path.exists(LOG_FILE_PATH):
|
||||
os.remove(LOG_FILE_PATH)
|
||||
|
||||
def tearDown(self):
|
||||
if os.path.exists(path_in_workspace(self.test_file)):
|
||||
os.remove(path_in_workspace(self.test_file))
|
||||
|
||||
if os.path.exists(self.test_directory):
|
||||
shutil.rmtree(self.test_directory)
|
||||
def tearDown(self) -> None:
|
||||
shutil.rmtree(self.workspace_path)
|
||||
|
||||
def test_check_duplicate_operation(self):
|
||||
log_operation("write", self.test_file)
|
||||
@@ -53,9 +52,9 @@ class TestFileOperations(unittest.TestCase):
|
||||
os.remove(self.file_logger_logs)
|
||||
|
||||
log_operation("log_test", self.test_file)
|
||||
with open(LOG_FILE_PATH, "r") as f:
|
||||
with open(self.config.file_logger_path, "r") as f:
|
||||
content = f.read()
|
||||
self.assertIn("log_test: test_file.txt", content)
|
||||
self.assertIn(f"log_test: {self.test_file}", content)
|
||||
|
||||
# Test splitting a file into chunks
|
||||
def test_split_file(self):
|
||||
@@ -71,80 +70,59 @@ class TestFileOperations(unittest.TestCase):
|
||||
def test_write_to_file(self):
|
||||
new_content = "This is new content.\n"
|
||||
write_to_file(self.test_file, new_content)
|
||||
with open(path_in_workspace(self.test_file), "r") as f:
|
||||
with open(self.test_file, "r") as f:
|
||||
content = f.read()
|
||||
self.assertEqual(content, new_content)
|
||||
|
||||
def test_append_to_file(self):
|
||||
with open(path_in_workspace(self.test_file), "r") as f:
|
||||
with open(self.test_file, "r") as f:
|
||||
content_before = f.read()
|
||||
|
||||
append_text = "This is appended text.\n"
|
||||
append_to_file(self.test_file, append_text)
|
||||
with open(path_in_workspace(self.test_file), "r") as f:
|
||||
with open(self.test_file, "r") as f:
|
||||
content = f.read()
|
||||
|
||||
self.assertEqual(content, content_before + append_text)
|
||||
|
||||
def test_delete_file(self):
|
||||
delete_file(self.test_file)
|
||||
self.assertFalse(os.path.exists(path_in_workspace(self.test_file)))
|
||||
self.assertFalse(os.path.exists(self.test_file))
|
||||
|
||||
def test_search_files(self):
|
||||
# Case 1: Create files A and B, search for A, and ensure we don't return A and B
|
||||
file_a = "file_a.txt"
|
||||
file_b = "file_b.txt"
|
||||
file_a = self.workspace.get_path("file_a.txt")
|
||||
file_b = self.workspace.get_path("file_b.txt")
|
||||
|
||||
with open(path_in_workspace(file_a), "w") as f:
|
||||
with open(file_a, "w") as f:
|
||||
f.write("This is file A.")
|
||||
|
||||
with open(path_in_workspace(file_b), "w") as f:
|
||||
with open(file_b, "w") as f:
|
||||
f.write("This is file B.")
|
||||
|
||||
# Create a subdirectory and place a copy of file_a in it
|
||||
if not os.path.exists(path_in_workspace(self.test_directory)):
|
||||
os.makedirs(path_in_workspace(self.test_directory))
|
||||
if not os.path.exists(self.test_directory):
|
||||
os.makedirs(self.test_directory)
|
||||
|
||||
with open(
|
||||
path_in_workspace(os.path.join(self.test_directory, file_a)), "w"
|
||||
) as f:
|
||||
with open(os.path.join(self.test_directory, file_a.name), "w") as f:
|
||||
f.write("This is file A in the subdirectory.")
|
||||
|
||||
files = search_files(path_in_workspace(""))
|
||||
self.assertIn(file_a, files)
|
||||
self.assertIn(file_b, files)
|
||||
self.assertIn(os.path.join(self.test_directory, file_a), files)
|
||||
files = search_files(str(self.workspace.root))
|
||||
self.assertIn(file_a.name, files)
|
||||
self.assertIn(file_b.name, files)
|
||||
self.assertIn(f"{Path(self.test_directory).name}/{file_a.name}", files)
|
||||
|
||||
# Clean up
|
||||
os.remove(path_in_workspace(file_a))
|
||||
os.remove(path_in_workspace(file_b))
|
||||
os.remove(path_in_workspace(os.path.join(self.test_directory, file_a)))
|
||||
os.rmdir(path_in_workspace(self.test_directory))
|
||||
os.remove(file_a)
|
||||
os.remove(file_b)
|
||||
os.remove(os.path.join(self.test_directory, file_a.name))
|
||||
os.rmdir(self.test_directory)
|
||||
|
||||
# Case 2: Search for a file that does not exist and make sure we don't throw
|
||||
non_existent_file = "non_existent_file.txt"
|
||||
files = search_files("")
|
||||
self.assertNotIn(non_existent_file, files)
|
||||
|
||||
# Test to ensure we cannot read files out of workspace
|
||||
def test_restrict_workspace(self):
|
||||
CFG = Config()
|
||||
with open(self.test_file2, "w+") as f:
|
||||
f.write("test text")
|
||||
|
||||
CFG.restrict_to_workspace = True
|
||||
|
||||
# Get the absolute path of self.test_file2
|
||||
test_file2_abs_path = os.path.abspath(self.test_file2)
|
||||
|
||||
with self.assertRaises(ValueError):
|
||||
read_file(test_file2_abs_path)
|
||||
|
||||
CFG.restrict_to_workspace = False
|
||||
read_file(test_file2_abs_path)
|
||||
|
||||
os.remove(test_file2_abs_path)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
Reference in New Issue
Block a user