From 6b86c147aad65cef400b462db1e938e3ac1c6e11 Mon Sep 17 00:00:00 2001
From: Jo Basevi <johanna.basevi@anu.edu.au>
Date: Wed, 18 Sep 2024 12:46:31 +1000
Subject: [PATCH] Fix clone with relative restart paths

---
 payu/branch.py      | 31 ++++++++++++++++++++++++-------
 test/test_branch.py | 41 +++++++++++++++++++++++++++++++++++++++++
 2 files changed, 65 insertions(+), 7 deletions(-)

diff --git a/payu/branch.py b/payu/branch.py
index dd0c51b6..7181747f 100644
--- a/payu/branch.py
+++ b/payu/branch.py
@@ -36,10 +36,23 @@
 Where BRANCH_NAME is the name of the branch"""
 
 
-def check_restart(restart_path: Optional[Path],
-                  archive_path: Path) -> Optional[Path]:
-    """Checks for valid prior restart path. Returns resolved restart path
-    if valid, otherwise returns None"""
+def check_restart(restart_path: Path,
+                  archive_path: Optional[Path] = None) -> Optional[Path]:
+    """Checks if restart path exists and whether the archive already
+    has pre-existing restarts. Returns a resolved restart path
+
+    Parameters
+    ----------
+    restart_path: Path
+        Absolute, or relative, restart path to start experiment from
+    archive_path: Optional[Path], default None
+        Experiment archive directory to check for pre-existing restarts files
+
+    Returns
+    ----------
+    Optional[Path]
+        Absolute restart path if a valid path, otherwise None
+    """
 
     # Check for valid path
     if not restart_path.exists():
@@ -51,7 +64,7 @@ def check_restart(restart_path: Optional[Path],
     restart_path = restart_path.resolve()
 
     # Check for pre-existing restarts in archive
-    if archive_path.exists():
+    if archive_path and archive_path.exists():
         if len(list_archive_dirs(archive_path, dir_type="restart")) > 0:
             warnings.warn((
                 f"Pre-existing restarts found in archive: {archive_path}."
@@ -136,7 +149,7 @@ def checkout_branch(branch_name: str,
     start_point: Optional[str]
         Branch name or commit hash to start new branch from
     restart_path: Optional[Path]
-        Absolute restart path to start experiment from
+        Restart path to start experiment from
     config_path: Optional[Path]
         Path to configuration file - config.yaml
     control_path: Optional[Path]
@@ -243,7 +256,7 @@ def clone(repository: str,
         lab_path: Optional[Path]
             Path to laboratory directory
         restart_path: Optional[Path]
-            Absolute restart path to start experiment from
+            Restart path to start experiment from
         parent_experiment: Optional[str]
             Parent experiment UUID to add to generated metadata
 
@@ -262,6 +275,10 @@ def clone(repository: str,
     # git clone the repository
     repo = git_clone(repository, control_path, branch)
 
+    if restart_path:
+        # Check path exists and resolve to an absolute path
+        restart_path = check_restart(restart_path)
+
     owd = os.getcwd()
     try:
         # cd into cloned directory
diff --git a/test/test_branch.py b/test/test_branch.py
index f8fc9ae7..b762ba4c 100644
--- a/test/test_branch.py
+++ b/test/test_branch.py
@@ -109,6 +109,19 @@ def test_add_restart_to_config(test_config, expected_config):
     assert updated_config == expected_config
 
 
+def test_check_restart_relative_path():
+    """Test an relative restart path is resolved to an absolute path"""
+    restart_path = tmpdir / "archive" / "tmpRestart"
+    restart_path.mkdir(parents=True)
+
+    with cd(tmpdir):
+        relative_restart_path = Path("archive") / "tmpRestart"
+        assert not relative_restart_path.is_absolute()
+
+        resolved_path = check_restart(relative_restart_path)
+        assert resolved_path.is_absolute()
+
+
 def test_check_restart_with_non_existent_restart():
     """Test restart path that does not exist raises a warning"""
     restart_path = tmpdir / "restartDNE"
@@ -536,6 +549,34 @@ def test_clone(mock_uuid):
     assert [head.name for head in cloned_repo2.heads] == ["Branch1", "Branch2"]
 
 
+def test_clone_with_relative_restart_path():
+    """Test clone with a restart path that is relative with respect to
+    the directory in which the clone command is run from"""
+    # Create a repo to clone
+    source_repo_path = tmpdir / "sourceRepo"
+    source_repo_path.mkdir()
+    setup_control_repository(path=source_repo_path)
+
+    # Create restart path
+    restart_path = tmpdir / "archive" / "tmpRestart"
+    restart_path.mkdir(parents=True)
+    relative_restart_path = Path("archive") / "tmpRestart"
+
+    cloned_repo_path = tmpdir / "clonedRepo"
+    with cd(tmpdir):
+        # Run clone
+        clone(repository=str(source_repo_path),
+              directory=cloned_repo_path,
+              lab_path=labdir,
+              restart_path=relative_restart_path)
+
+    # Test restart was added to config.yaml file
+    with cd(cloned_repo_path):
+        config = read_config()
+
+    assert config["restart"] == str(restart_path)
+
+
 def add_and_commit_metadata(repo, metadata):
     """Helper function to create/update metadata file and commit"""
     metadata_path = ctrldir / "metadata.yaml"