extend Experiment.prune_restarts test to use access-om2 model

payu-org · Sep 27, 2023 · 874b99a · 874b99a
1 parent 0bf3bd7
commit 874b99a
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 11 deletions.
diff --git a/payu/models/accessom2.py b/payu/models/accessom2.py
@@ -88,9 +88,13 @@ def collate(self):
         pass
 
     def get_restart_datetime(self, restart_path):
+        """Given a restart path, parse the restart files and
+        return a cftime datetime (for date-based restart pruning)"""
         for model in self.expt.models:
             if model.model_type == 'mom':
                 mom_restart_path = os.path.join(restart_path, model.name)
                 return model.get_restart_datetime(mom_restart_path)
-
-        raise NotImplementedError('Cannot find mom sub-model: access-om2 date-based restart pruning requires the mom sub-model to determine restart dates')
+
+        raise NotImplementedError(
+            'Cannot find mom sub-model: access-om2 date-based restart pruning '
+            'requires the mom sub-model to determine restart dates')
diff --git a/test/models/test_mom.py b/test/models/test_mom.py
@@ -151,10 +151,10 @@ 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')]
+    restarts = [d for d in os.listdir(archive_path)
+                if d.startswith('restart')]
 
-        for index, restart in enumerate(restarts):
-            restart_path = os.path.join(archive_path, restart)
-            run_dt = expt.model.get_restart_datetime(restart_path)
-            assert run_dt == expected_cftimes[index]
+    for index, expected_cftime in enumerate(expected_cftimes):
+        restart_path = os.path.join(archive_path, restarts[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
@@ -102,7 +102,6 @@ def test_prune_restarts(restart_freq,
     # (1900/03/01, restart001)
     #  ...
     # (1902/02/01, restart024)
-    init_dt_array = [1900, 1, 1, 0, 0, 0]
     restart_dts = []
     for year in [1900, 1901, 1902]:
         for month in range(1, 13):
@@ -111,11 +110,19 @@ def test_prune_restarts(restart_freq,
                 continue
             restart_dts.append([year, month, 1, 0, 0, 0])
 
-    make_ocean_restart_files(init_dt_array, restart_dts, 4)
+    make_ocean_restart_files(
+        init_dt_array=[1900, 1, 1, 0, 0, 0],
+        run_dt_arrays=restart_dts,
+        calendar=4,
+        additional_restart_path='ocean')
 
     # Set up config
     test_config = config
-    test_config['model'] = 'mom'
+    test_config['model'] = 'access-om2'
+    test_config['submodels'] = [
+        {'name': 'atmosphere', 'model': 'yatm'},
+        {'name': 'ocean', 'model': 'mom'}
+    ]
     test_config['restart_freq'] = restart_freq
     test_config['restart_history'] = restart_history
     write_config(test_config)