#150- refactor work_dir creation

CABLE-LSM · Oct 5, 2023 · 4305062 · 4305062
1 parent de33ae0
commit 4305062
Show file tree

Hide file tree

Showing 3 changed files with 127 additions and 113 deletions.
diff --git a/benchcab/benchcab.py b/benchcab/benchcab.py
@@ -11,7 +11,7 @@
 from benchcab import internal
 from benchcab.internal import get_met_forcing_file_names
 from benchcab.config import read_config
-from benchcab.workdir import setup_fluxsite_directory_tree, setup_src_dir
+from benchcab.workdir import fluxsite_directory_tree_list, setup_fluxsite_directory_tree, setup_src_dir
 from benchcab.repository import CableRepository
 from benchcab.fluxsite import (
     get_fluxsite_tasks,
@@ -215,7 +215,9 @@ def fluxsite_setup_work_directory(self):
         """Endpoint for `benchcab fluxsite-setup-work-dir`."""
         tasks = self.tasks if self.tasks else self._initialise_tasks()
         print("Setting up run directory tree for fluxsite tests...")
-        setup_fluxsite_directory_tree(fluxsite_tasks=tasks, verbose=self.args.verbose)
+        fluxsite_paths = fluxsite_directory_tree_list(
+            fluxsite_tasks=tasks)
+        setup_fluxsite_directory_tree(fluxsite_paths=fluxsite_paths, verbose=self.args.verbose)
         print("Setting up tasks...")
         for task in tasks:
             task.setup_task(verbose=self.args.verbose)

diff --git a/benchcab/workdir.py b/benchcab/workdir.py
@@ -3,6 +3,7 @@
 from pathlib import Path
 import os
 import shutil
+from typing import Union
 
 from benchcab import internal
 from benchcab.fluxsite import Task
@@ -28,60 +29,59 @@ def setup_src_dir(root_dir=internal.CWD):
         os.makedirs(src_dir)
 
 
-def setup_fluxsite_directory_tree(
-    fluxsite_tasks: list[Task], root_dir=internal.CWD, verbose=False
-):
-    """Generate the directory structure used of `benchcab`."""
+def fluxsite_directory_tree_list(fluxsite_tasks: list[Task], root_dir=internal.CWD) -> set:
+    """Generate a list of all the work directories used for fluxsite tests"""
 
-    run_dir = Path(root_dir, internal.RUN_DIR)
-    if not run_dir.exists():
-        os.makedirs(run_dir)
-
-    fluxsite_run_dir = Path(root_dir, internal.FLUXSITE_RUN_DIR)
-    if not fluxsite_run_dir.exists():
-        os.makedirs(fluxsite_run_dir)
-
-    fluxsite_log_dir = Path(root_dir, internal.FLUXSITE_LOG_DIR)
-    if not fluxsite_log_dir.exists():
-        print(
-            f"Creating {fluxsite_log_dir.relative_to(root_dir)} directory: {fluxsite_log_dir}"
-        )
-        os.makedirs(fluxsite_log_dir)
-
-    fluxsite_output_dir = Path(root_dir, internal.FLUXSITE_OUTPUT_DIR)
-    if not fluxsite_output_dir.exists():
-        print(
-            f"Creating {fluxsite_output_dir.relative_to(root_dir)} directory: {fluxsite_output_dir}"
-        )
-        os.makedirs(fluxsite_output_dir)
-
-    fluxsite_tasks_dir = Path(root_dir, internal.FLUXSITE_TASKS_DIR)
-    if not fluxsite_tasks_dir.exists():
-        print(
-            f"Creating {fluxsite_tasks_dir.relative_to(root_dir)} directory: {fluxsite_tasks_dir}"
-        )
-        os.makedirs(fluxsite_tasks_dir)
-
-    fluxsite_analysis_dir = Path(root_dir, internal.FLUXSITE_ANALYSIS_DIR)
-    if not fluxsite_analysis_dir.exists():
-        print(
-            f"Creating {fluxsite_analysis_dir.relative_to(root_dir)} directory: "
-            f"{fluxsite_analysis_dir}"
-        )
-        os.makedirs(fluxsite_analysis_dir)
-
-    fluxsite_bitwise_cmp_dir = Path(root_dir, internal.FLUXSITE_BITWISE_CMP_DIR)
-    if not fluxsite_bitwise_cmp_dir.exists():
-        print(
-            f"Creating {fluxsite_bitwise_cmp_dir.relative_to(root_dir)} directory: "
-            f"{fluxsite_bitwise_cmp_dir}"
-        )
-        os.makedirs(fluxsite_bitwise_cmp_dir)
-
-    print("Creating task directories...")
+    # Create the list of directories as sets because the order is not important
+    # and we want to avoid duplications.
+    fluxsite_paths = []
+
+    # Run directory
+    fluxsite_paths.append(Path(root_dir, internal.RUN_DIR))
+
+    # Fluxsite run directory
+    fluxsite_paths.append(Path(root_dir, internal.FLUXSITE_RUN_DIR))
+
+    # Fluxsite log directory
+    fluxsite_paths.append(Path(root_dir, internal.FLUXSITE_LOG_DIR))
+
+    # Fluxsite output directory
+    fluxsite_paths.append(Path(root_dir, internal.FLUXSITE_OUTPUT_DIR))
+
+    # Fluxsite tasks directory
+    fluxsite_paths.append(Path(root_dir, internal.FLUXSITE_TASKS_DIR))
+
+    # Fluxsite analysis directory
+    fluxsite_paths.append(Path(root_dir, internal.FLUXSITE_ANALYSIS_DIR))
+
+    # Fluxsite bit-wise comparison directory
+    fluxsite_paths.append(Path(root_dir, internal.FLUXSITE_BITWISE_CMP_DIR))
+
+    # Fluxsite tasks directories. append all of them as a set().
+    task_paths = []
     for task in fluxsite_tasks:
-        task_dir = Path(root_dir, internal.FLUXSITE_TASKS_DIR, task.get_task_name())
-        if not task_dir.exists():
-            if verbose:
-                print(f"Creating {task_dir.relative_to(root_dir)}: " f"{task_dir}")
-            os.makedirs(task_dir)
+        task_paths.append(
+            Path(root_dir, internal.FLUXSITE_TASKS_DIR, task.get_task_name()))
+    fluxsite_paths.append(task_paths)
+
+    return fluxsite_paths
+
+
+def setup_fluxsite_directory_tree(
+    fluxsite_paths: list[Union[Path, list]], root_dir=internal.CWD, verbose=False
+):
+    """Generate the directory structure used by `benchcab`."""
+
+    for work_path in fluxsite_paths:
+        if isinstance(work_path, list):
+            print("Creating task directories...")
+
+            for task_dir in work_path:
+                if verbose:
+                    print(f"Creating {task_dir.relative_to(root_dir)}: {task_dir}")
+                task_dir.mkdir(parents=True, exist_ok=True)
+        else:
+            print(
+                f"Creating {work_path.relative_to(root_dir)} directory: {work_path}"
+            )
+            work_path.mkdir(parents=True, exist_ok=True)
diff --git a/tests/test_workdir.py b/tests/test_workdir.py
@@ -11,6 +11,7 @@
 from benchcab.fluxsite import Task
 from benchcab.repository import CableRepository
 from benchcab.workdir import (
+    fluxsite_directory_tree_list,
     setup_fluxsite_directory_tree,
     clean_directory_tree,
     setup_src_dir,
@@ -42,89 +43,100 @@ def setup_mock_tasks() -> list[Task]:
     return tasks
 
 
+def setup_mock_directory_tree_list():
+    """Return the list of directories for the work directory we want benchcab to create"""
+
+    full_directory_lists = []
+    full_directory_lists.append(Path(MOCK_CWD, "runs"))
+    full_directory_lists.append(Path(MOCK_CWD, "runs", "fluxsite"))
+    full_directory_lists.append(Path(MOCK_CWD, "runs", "fluxsite", "logs"))
+    full_directory_lists.append(Path(MOCK_CWD, "runs", "fluxsite", "outputs"))
+    full_directory_lists.append(Path(MOCK_CWD, "runs", "fluxsite", "tasks"))
+    full_directory_lists.append(Path(MOCK_CWD, "runs", "fluxsite", "analysis"))
+    full_directory_lists.append(Path(MOCK_CWD, "runs", "fluxsite",
+                                "analysis", "bitwise-comparisons"))
+
+    task_directories = []
+    task_directories.append(Path(MOCK_CWD, "runs", "fluxsite", "tasks", "site_foo_R0_S0"))
+    task_directories.append(Path(MOCK_CWD, "runs", "fluxsite", "tasks", "site_foo_R0_S1"))
+    task_directories.append(Path(MOCK_CWD, "runs", "fluxsite", "tasks", "site_bar_R0_S0"))
+    task_directories.append(Path(MOCK_CWD, "runs", "fluxsite", "tasks", "site_bar_R0_S1"))
+    task_directories.append(Path(MOCK_CWD, "runs", "fluxsite", "tasks", "site_foo_R1_S0"))
+    task_directories.append(Path(MOCK_CWD, "runs", "fluxsite", "tasks", "site_foo_R1_S1"))
+    task_directories.append(Path(MOCK_CWD, "runs", "fluxsite", "tasks", "site_bar_R1_S0"))
+    task_directories.append(Path(MOCK_CWD, "runs", "fluxsite", "tasks", "site_bar_R1_S1"))
+    full_directory_lists.append(task_directories)
+
+    return full_directory_lists
+
+
+def test_fluxsite_directory_tree_list():
+    """Tests for `fluxsite_directory_tree_list()`."""
+
+    # Success case: generate the full lists of directories for the mock tasks
+    tasks = setup_mock_tasks()
+    full_directory_lists = setup_mock_directory_tree_list()
+
+    fluxsite_paths = fluxsite_directory_tree_list(
+        fluxsite_tasks=tasks, root_dir=MOCK_CWD)
+
+    assert fluxsite_paths == full_directory_lists
+
+
 def test_setup_directory_tree():
     """Tests for `setup_fluxsite_directory_tree()`."""
 
-    # Success case: generate fluxsite directory structure
+    mock_directory_list = [Path(MOCK_CWD, "test"), Path(MOCK_CWD, "test", "test1"), [
+        Path(MOCK_CWD, "test", "task1"), Path(MOCK_CWD, "test", "task2")]]
+
+    # Success case: generate given directory structure
     tasks = setup_mock_tasks()
-    setup_fluxsite_directory_tree(fluxsite_tasks=tasks, root_dir=MOCK_CWD)
-
-    assert len(list(MOCK_CWD.glob("*"))) == 1
-    assert Path(MOCK_CWD, "runs").exists()
-    assert Path(MOCK_CWD, "runs", "fluxsite").exists()
-    assert Path(MOCK_CWD, "runs", "fluxsite", "logs").exists()
-    assert Path(MOCK_CWD, "runs", "fluxsite", "outputs").exists()
-    assert Path(MOCK_CWD, "runs", "fluxsite", "tasks").exists()
-    assert Path(
-        MOCK_CWD, "runs", "fluxsite", "analysis", "bitwise-comparisons"
-    ).exists()
-
-    assert Path(MOCK_CWD, "runs", "fluxsite", "tasks", "site_foo_R0_S0").exists()
-    assert Path(MOCK_CWD, "runs", "fluxsite", "tasks", "site_foo_R0_S1").exists()
-    assert Path(MOCK_CWD, "runs", "fluxsite", "tasks", "site_bar_R0_S0").exists()
-    assert Path(MOCK_CWD, "runs", "fluxsite", "tasks", "site_bar_R0_S1").exists()
-    assert Path(MOCK_CWD, "runs", "fluxsite", "tasks", "site_foo_R1_S0").exists()
-    assert Path(MOCK_CWD, "runs", "fluxsite", "tasks", "site_foo_R1_S1").exists()
-    assert Path(MOCK_CWD, "runs", "fluxsite", "tasks", "site_bar_R1_S0").exists()
-    assert Path(MOCK_CWD, "runs", "fluxsite", "tasks", "site_bar_R1_S1").exists()
-
-    shutil.rmtree(MOCK_CWD / "runs")
+    setup_fluxsite_directory_tree(mock_directory_list, root_dir=MOCK_CWD)
+
+    assert Path(MOCK_CWD, "test").exists()
+    assert Path(MOCK_CWD, "test", "test1").exists()
+    assert Path(MOCK_CWD, "test", "task1").exists()
+    assert Path(MOCK_CWD, "test", "task2").exists()
+
+    shutil.rmtree(MOCK_CWD / "test")
 
     # Success case: test non-verbose output
     with contextlib.redirect_stdout(io.StringIO()) as buf:
-        setup_fluxsite_directory_tree(fluxsite_tasks=tasks, root_dir=MOCK_CWD)
+        setup_fluxsite_directory_tree(mock_directory_list, root_dir=MOCK_CWD)
     assert buf.getvalue() == (
-        f"Creating runs/fluxsite/logs directory: {MOCK_CWD}/runs/fluxsite/logs\n"
-        f"Creating runs/fluxsite/outputs directory: {MOCK_CWD}/runs/fluxsite/outputs\n"
-        f"Creating runs/fluxsite/tasks directory: {MOCK_CWD}/runs/fluxsite/tasks\n"
-        f"Creating runs/fluxsite/analysis directory: {MOCK_CWD}/runs/fluxsite/analysis\n"
-        f"Creating runs/fluxsite/analysis/bitwise-comparisons directory: {MOCK_CWD}"
-        "/runs/fluxsite/analysis/bitwise-comparisons\n"
+        f"Creating test directory: {MOCK_CWD}/test\n"
+        f"Creating test/test1 directory: {MOCK_CWD}/test/test1\n"
         f"Creating task directories...\n"
     )
 
-    shutil.rmtree(MOCK_CWD / "runs")
+    shutil.rmtree(MOCK_CWD / "test")
 
     # Success case: test verbose output
     with contextlib.redirect_stdout(io.StringIO()) as buf:
         setup_fluxsite_directory_tree(
-            fluxsite_tasks=tasks, verbose=True, root_dir=MOCK_CWD
+            mock_directory_list, verbose=True, root_dir=MOCK_CWD
         )
     assert buf.getvalue() == (
-        f"Creating runs/fluxsite/logs directory: {MOCK_CWD}/runs/fluxsite/logs\n"
-        f"Creating runs/fluxsite/outputs directory: {MOCK_CWD}/runs/fluxsite/outputs\n"
-        f"Creating runs/fluxsite/tasks directory: {MOCK_CWD}/runs/fluxsite/tasks\n"
-        f"Creating runs/fluxsite/analysis directory: {MOCK_CWD}/runs/fluxsite/analysis\n"
-        f"Creating runs/fluxsite/analysis/bitwise-comparisons directory: {MOCK_CWD}"
-        "/runs/fluxsite/analysis/bitwise-comparisons\n"
+        f"Creating test directory: {MOCK_CWD}/test\n"
+        f"Creating test/test1 directory: {MOCK_CWD}/test/test1\n"
         f"Creating task directories...\n"
-        f"Creating runs/fluxsite/tasks/site_foo_R0_S0: "
-        f"{MOCK_CWD}/runs/fluxsite/tasks/site_foo_R0_S0\n"
-        f"Creating runs/fluxsite/tasks/site_foo_R0_S1: "
-        f"{MOCK_CWD}/runs/fluxsite/tasks/site_foo_R0_S1\n"
-        f"Creating runs/fluxsite/tasks/site_bar_R0_S0: "
-        f"{MOCK_CWD}/runs/fluxsite/tasks/site_bar_R0_S0\n"
-        f"Creating runs/fluxsite/tasks/site_bar_R0_S1: "
-        f"{MOCK_CWD}/runs/fluxsite/tasks/site_bar_R0_S1\n"
-        f"Creating runs/fluxsite/tasks/site_foo_R1_S0: "
-        f"{MOCK_CWD}/runs/fluxsite/tasks/site_foo_R1_S0\n"
-        f"Creating runs/fluxsite/tasks/site_foo_R1_S1: "
-        f"{MOCK_CWD}/runs/fluxsite/tasks/site_foo_R1_S1\n"
-        f"Creating runs/fluxsite/tasks/site_bar_R1_S0: "
-        f"{MOCK_CWD}/runs/fluxsite/tasks/site_bar_R1_S0\n"
-        f"Creating runs/fluxsite/tasks/site_bar_R1_S1: "
-        f"{MOCK_CWD}/runs/fluxsite/tasks/site_bar_R1_S1\n"
+        f"Creating test/task1: "
+        f"{MOCK_CWD}/test/task1\n"
+        f"Creating test/task2: "
+        f"{MOCK_CWD}/test/task2\n"
     )
 
-    shutil.rmtree(MOCK_CWD / "runs")
+    shutil.rmtree(MOCK_CWD / "test")
 
 
 def test_clean_directory_tree():
     """Tests for `clean_directory_tree()`."""
 
     # Success case: directory tree does not exist after clean
     tasks = setup_mock_tasks()
-    setup_fluxsite_directory_tree(fluxsite_tasks=tasks, root_dir=MOCK_CWD)
+    full_directory_lists = setup_mock_directory_tree_list()
+
+    setup_fluxsite_directory_tree(full_directory_lists, root_dir=MOCK_CWD)
 
     clean_directory_tree(root_dir=MOCK_CWD)
     assert not Path(MOCK_CWD, "runs").exists()