Skip to content

Commit

Permalink
test_link_result_basic_functionality
Browse files Browse the repository at this point in the history
  • Loading branch information
wangpatrick57 committed Dec 27, 2024
1 parent 316ca02 commit 88f31b3
Show file tree
Hide file tree
Showing 2 changed files with 89 additions and 2 deletions.
36 changes: 35 additions & 1 deletion util/tests/unittest_workspace.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@
make_workspace_structure,
verify_structure,
)
from util.workspace import DBGymWorkspace, save_file
from util.workspace import DBGymWorkspace, link_result, save_file


class WorkspaceTests(unittest.TestCase):
Expand All @@ -33,6 +33,7 @@ def setUp(self) -> None:
DBGymWorkspace.num_times_created_this_run = 0

def tearDown(self) -> None:
# You can comment this out if you want to inspect the scratchspace after a test (often used for debugging).
if self.scratchspace_path.exists():
shutil.rmtree(self.scratchspace_path)

Expand Down Expand Up @@ -63,6 +64,10 @@ def get_updated_structure_from_workspace_init(
)
return structure

def test_init_fields(self) -> None:
workspace = DBGymWorkspace(self.workspace_path)
self.assertEqual(workspace.app_name, "dbgym")

def test_init_from_nonexistent_workspace(self) -> None:
starting_structure = FilesystemStructure({})
create_structure(self.scratchspace_path, starting_structure)
Expand Down Expand Up @@ -97,6 +102,35 @@ def test_init_from_already_initialized_workspace(self) -> None:

self.assertTrue(verify_structure(self.scratchspace_path, ending_structure))

def test_link_result_basic_functionality(self) -> None:
# Setup.
workspace = DBGymWorkspace(self.workspace_path)
ending_structure = WorkspaceTests.get_workspace_init_structure(workspace)

# Make a result file.
result_path = workspace.dbgym_this_run_path / "result.txt"
result_path.touch()
ending_structure["dbgym_workspace"]["task_runs"][
workspace.dbgym_this_run_path.name
]["result.txt"] = ("file",)

# Link the result file.
workspace.link_result(result_path)
ending_structure["dbgym_workspace"]["symlinks"]["dbgym"] = {}
ending_structure["dbgym_workspace"]["symlinks"]["dbgym"]["result.txt.link"] = (
"symlink",
f"dbgym_workspace/task_runs/{workspace.dbgym_this_run_path.name}/result.txt",
)

# Verify structure.
self.assertTrue(verify_structure(self.scratchspace_path, ending_structure))

# TODO: test overriding existing symlink
# TODO: test linking result from another run should raise
# TODO: test that it should link in the agent dir in the links
# TODO: test that it will ignore the directory structure (unlike save which keeps it)
# TODO: test linking a symlink or a non-fully-resolved path


if __name__ == "__main__":
unittest.main()
55 changes: 54 additions & 1 deletion util/workspace.py
Original file line number Diff line number Diff line change
Expand Up @@ -168,6 +168,9 @@ def __init__(self, dbgym_workspace_path: Path):

self.base_dbgym_repo_dpath = get_base_dbgym_repo_dpath()
self.cur_path_list: list[str] = ["dbgym"]
self.app_name = (
"dbgym" # TODO: discover this dynamically. app means dbgym or an agent
)

# Set and create paths.
self.dbgym_workspace_path = dbgym_workspace_path
Expand Down Expand Up @@ -220,6 +223,56 @@ def __init__(self, dbgym_workspace_path: Path):
try_remove_file(self.dbgym_latest_run_path)
try_create_symlink(self.dbgym_this_run_path, self.dbgym_latest_run_path)

# TODO(phw2): refactor our manual symlinking in postgres/cli.py to use link_result() instead
def link_result(
self,
result_fordpath: Path,
custom_result_name: Optional[str] = None,
) -> Path:
"""
result_fordpath must be a "result", meaning it was generated inside dbgym_workspace.dbgym_this_run_path.
Further, result_fordpath must have been generated by this invocation to task.py. This also means that
result_fordpath itself can be a file or a dir but not a symlink.
Given a file or directory in task_runs/run_*/[codebase]/[org], this will create a symlink inside
symlinks/[codebase]/[org]/.
Will override the old symlink if there is one, so that symlinks/ always contains the latest generated
version of a file.
This function will return the path to the symlink that was created.
"""
assert isinstance(result_fordpath, Path)
assert is_fully_resolved(
result_fordpath
), f"result_fordpath ({result_fordpath}) should be a fully resolved path"
assert is_child_path(result_fordpath, self.dbgym_this_run_path)
assert not os.path.islink(result_fordpath)

if type(custom_result_name) is str:
result_name = custom_result_name
else:
if os.path.isfile(result_fordpath):
result_name = basename_of_path(result_fordpath) + ".link"
elif os.path.isdir(result_fordpath):
result_name = basename_of_path(result_fordpath) + ".link"
else:
raise AssertionError("result_fordpath must be either a file or dir")

# TODO: check that it was generated in this run_*/.
symlink_parent_dpath = self.dbgym_symlinks_path / self.app_name
symlink_parent_dpath.mkdir(parents=True, exist_ok=True)

# Remove the old symlink ("old" meaning created in an earlier run) if there is one
# Note that in a multi-threaded setting, this might remove one created by a process in the same run,
# meaning it's not "old" by our definition of "old". However, we'll always end up with a symlink
# file of the current run regardless of the order of threads.
assert result_name.endswith(".link") and not result_name.endswith(
".link.link"
), f'result_name ({result_name}) should end with ".link"'
symlink_path = symlink_parent_dpath / result_name
try_remove_file(symlink_path)
try_create_symlink(result_fordpath, symlink_path)

return symlink_path

# `append_group()` is used to mark the "codebase path" of an invocation of the CLI. The "codebase path" is
# explained further in the documentation.
def append_group(self, name: str) -> None:
Expand Down Expand Up @@ -564,7 +617,7 @@ def save_file(dbgym_workspace: DBGymWorkspace, fpath: Path) -> None:
shutil.copy(fpath, copy_fpath)


# TODO(phw2): refactor our manual symlinking in postgres/cli.py to use link_result() instead
# TODO(phw2): deprecate this once I'm done with unittest_workspace.py
def link_result(
dbgym_workspace: DBGymWorkspace,
result_fordpath: Path,
Expand Down

0 comments on commit 88f31b3

Please sign in to comment.