test_link_result_basic_functionality

Author: wangpatrick57

util/tests/unittest_workspace.py

@@ -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):
@@ -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)
 
@@ -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)
@@ -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()

util/workspace.py

@@ -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
@@ -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:
@@ -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,