refactor make/list/del restart dir test code to common

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

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

test/test_prune_restarts.py

@@ -1,5 +1,4 @@
 import copy
-import os
 import shutil
 import re
 
@@ -12,15 +11,14 @@
 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
 
 # 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