From 25b12e89278e91da169c62128c3b761cfd01c691 Mon Sep 17 00:00:00 2001 From: Jo Basevi Date: Fri, 29 Sep 2023 09:45:41 +1000 Subject: [PATCH] refactor make/list/del restart dir test code to common --- test/common.py | 49 ++++++++++++++++++++++++++++++++++++- test/models/test_mom.py | 34 +++++++++---------------- test/test_prune_restarts.py | 16 +++--------- 3 files changed, 63 insertions(+), 36 deletions(-) diff --git a/test/common.py b/test/common.py index 7b185c01..e0aa81f3 100644 --- a/test/common.py +++ b/test/common.py @@ -1,6 +1,8 @@ from contextlib import contextmanager import os from pathlib import Path +import re +import shutil import yaml @@ -12,13 +14,18 @@ from payu.subcommands.setup_cmd import runcmd as payu_setup_orignal from payu.subcommands.sweep_cmd import runcmd as payu_sweep +ctrldir_basename = 'ctrl' + testdir = Path().cwd() / Path('test') tmpdir = testdir / 'tmp' -ctrldir = tmpdir / 'ctrl' +ctrldir = tmpdir / ctrldir_basename labdir = tmpdir / 'lab' workdir = ctrldir / 'work' payudir = tmpdir / 'payu' +archive_dir = labdir / 'archive' +expt_archive_dir = archive_dir / ctrldir_basename + print('tmpdir: {}'.format(tmpdir)) config = { @@ -42,6 +49,7 @@ } } + @contextmanager def cd(directory): """ @@ -153,6 +161,45 @@ def make_restarts(fnames=None): make_random_file(restartdir/fname, 5000**2 + i) +def make_expt_archive_dirs(dir_type, num_dirs=5, additional_path=None): + """Make experiment archive directories of given type (i.e. "restart" or + "output")""" + created_dirs = [] + for i in range(num_dirs): + dir_path = os.path.join(expt_archive_dir, f'{dir_type}{i:03d}') + if additional_path: + dir_path = os.path.join(dir_path, additional_path) + + os.makedirs(dir_path) + created_dirs.append(dir_path) + return created_dirs + + +def list_expt_archive_dirs(dir_type='restart', full_path=True): + """Return a list of output/restart paths in experiment archive + path""" + dirs = [] + if os.path.exists(expt_archive_dir): + if os.path.isdir(expt_archive_dir): + naming_pattern = re.compile(fr"^{dir_type}[0-9][0-9][0-9]$") + dirs = [d for d in os.listdir(expt_archive_dir) + if naming_pattern.match(d)] + + if full_path: + dirs = [os.path.join(expt_archive_dir, d) for d in dirs] + return dirs + + +def remove_expt_archive_dirs(dir_type='restart'): + """Remove experiment archive directories of the given type (i.e. "restart" + or "output"). Useful for cleaning up archive between tests""" + for dir_path in list_expt_archive_dirs(dir_type): + try: + shutil.rmtree(dir_path) + except Exception as e: + print(e) + + def make_all_files(): make_inputs() make_exe() diff --git a/test/models/test_mom.py b/test/models/test_mom.py index 03caa839..885fd01a 100644 --- a/test/models/test_mom.py +++ b/test/models/test_mom.py @@ -12,14 +12,15 @@ from test.common import config as config_orig from test.common import write_config from test.common import make_all_files +from test.common import list_expt_archive_dirs +from test.common import make_expt_archive_dirs, remove_expt_archive_dirs + verbose = True # Global config config = copy.deepcopy(config_orig) -archive_path = os.path.join(labdir, 'archive', 'ctrl') - def setup_module(module): """ @@ -39,7 +40,6 @@ def setup_module(module): labdir.mkdir() ctrldir.mkdir() make_all_files() - os.makedirs(archive_path) except Exception as e: print(e) @@ -64,29 +64,20 @@ def teardown(): yield # Remove any created restart files - restart_dirs = [d for d in os.listdir(archive_path) - if d.startswith('restart')] - - for restart in restart_dirs: - try: - shutil.rmtree(os.path.join(archive_path, restart)) - except Exception as e: - print(e) + remove_expt_archive_dirs(dir_type='restart') def make_ocean_restart_files(init_dt_array, run_dt_arrays, calendar, - additional_restart_path=None): - for index, run_dt_array in enumerate(run_dt_arrays): - # Make restart dir - restart_path = os.path.join(archive_path, f'restart{index:03d}') - if additional_restart_path is not None: - restart_path = os.path.join(restart_path, additional_restart_path) - os.makedirs(restart_path) + additional_path=None): + restart_paths = make_expt_archive_dirs(dir_type='restart', + num_dirs=len(run_dt_arrays), + additional_path=additional_path) + for index, run_dt_array in enumerate(run_dt_arrays): # Create ocean_solo.res file - make_ocean_solo_file(restart_path, + make_ocean_solo_file(restart_paths[index], init_dt_array, run_dt_array, calendar) @@ -151,10 +142,9 @@ def test_mom_get_restart_datetime(run_dt_arrays, calendar, expected_cftimes): lab = payu.laboratory.Laboratory(lab_path=str(labdir)) expt = payu.experiment.Experiment(lab, reproduce=False) - restarts = [d for d in os.listdir(archive_path) - if d.startswith('restart')] + restart_paths = list_expt_archive_dirs(dir_type='restart') for index, expected_cftime in enumerate(expected_cftimes): - restart_path = os.path.join(archive_path, restarts[index]) + restart_path = restart_paths[index] run_dt = expt.model.get_restart_datetime(restart_path) assert run_dt == expected_cftime diff --git a/test/test_prune_restarts.py b/test/test_prune_restarts.py index fc943140..2fb1c79f 100644 --- a/test/test_prune_restarts.py +++ b/test/test_prune_restarts.py @@ -1,5 +1,4 @@ import copy -import os import shutil import re @@ -12,6 +11,7 @@ from test.common import config as config_orig from test.common import write_config from test.common import make_all_files +from test.common import remove_expt_archive_dirs from test.models.test_mom import make_ocean_restart_files verbose = True @@ -19,8 +19,6 @@ # Global config config = copy.deepcopy(config_orig) -archive_path = os.path.join(labdir, 'archive', 'ctrl') - def setup_module(module): """ @@ -40,7 +38,6 @@ def setup_module(module): labdir.mkdir() ctrldir.mkdir() make_all_files() - os.makedirs(archive_path) except Exception as e: print(e) @@ -65,14 +62,7 @@ def teardown(): yield # Remove any created restart files - restart_dirs = [d for d in os.listdir(archive_path) - if d.startswith('restart')] - - for restart in restart_dirs: - try: - shutil.rmtree(os.path.join(archive_path, restart)) - except Exception as e: - print(e) + remove_expt_archive_dirs(dir_type='restart') @pytest.mark.parametrize( @@ -114,7 +104,7 @@ def test_prune_restarts(restart_freq, init_dt_array=[1900, 1, 1, 0, 0, 0], run_dt_arrays=restart_dts, calendar=4, - additional_restart_path='ocean') + additional_path='ocean') # Set up config test_config = config