From 430506275f42fc2b4275a543f817101e6c50a708 Mon Sep 17 00:00:00 2001 From: Claire Carouge Date: Thu, 5 Oct 2023 15:19:32 +1100 Subject: [PATCH] #150- refactor work_dir creation --- benchcab/benchcab.py | 6 +- benchcab/workdir.py | 110 ++++++++++++++++++------------------- tests/test_workdir.py | 124 +++++++++++++++++++++++------------------- 3 files changed, 127 insertions(+), 113 deletions(-) diff --git a/benchcab/benchcab.py b/benchcab/benchcab.py index 25c4bf73..f35e4169 100644 --- 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 index fa7fa32d..f578cdba 100644 --- 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 index 2df618a0..37f3181e 100644 --- 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,81 +43,90 @@ 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(): @@ -124,7 +134,9 @@ def test_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()