From 8c8c691f101037c6dd9b3843ae681dc3dd87885b Mon Sep 17 00:00:00 2001
From: Jo Basevi <johanna.basevi@edu.anu.au>
Date: Wed, 2 Aug 2023 16:10:40 +1000
Subject: [PATCH] Modify config.yaml modules option to add module dirs to
 MODULEPATH

---
 payu/experiment.py                        |  6 ++++-
 payu/fsops.py                             |  6 +++++
 test/resources/config_legacy_modules.yaml |  3 +++
 test/resources/config_modules.yaml        |  7 ++++++
 test/{ => resources}/config_mom5.yaml     |  0
 test/test_payu.py                         | 27 +++++++++++++++++++++--
 6 files changed, 46 insertions(+), 3 deletions(-)
 create mode 100644 test/resources/config_legacy_modules.yaml
 create mode 100644 test/resources/config_modules.yaml
 rename test/{ => resources}/config_mom5.yaml (100%)

diff --git a/payu/experiment.py b/payu/experiment.py
index a8bd6e00..544a11fc 100644
--- a/payu/experiment.py
+++ b/payu/experiment.py
@@ -258,7 +258,7 @@ def load_modules(self):
             envmod.module('load', mod)
 
         # User-defined modules
-        user_modules = self.config.get('modules', [])
+        user_modules = self.config.get('modules', {}).get('load', [])
         for mod in user_modules:
             envmod.module('load', mod)
 
@@ -457,6 +457,10 @@ def run(self, *user_flags):
 
         # XXX: This was previously done in reversion
         envmod.setup()
+        
+        # Add any user-defined module dir(s) to MODULEPATH
+        for module_dir in self.config.get('modules', {}).get('use', []):
+            envmod.module('use', module_dir) 
 
         self.load_modules()
 
diff --git a/payu/fsops.py b/payu/fsops.py
index 66317070..7cb9b172 100644
--- a/payu/fsops.py
+++ b/payu/fsops.py
@@ -96,6 +96,12 @@ def read_config(config_fname=None):
 
     config['collate'] = collate_config
 
+    # Transform legacy modules config options
+    modules_config = config.pop('modules', {})
+    if type(modules_config) is list:
+        modules_config = {'load': modules_config}
+    config['modules'] = modules_config
+     
     # Local "control" path. Must be set here so it can be
     # scanned for storage points
     config["control_path"] = config.get('control',
diff --git a/test/resources/config_legacy_modules.yaml b/test/resources/config_legacy_modules.yaml
new file mode 100644
index 00000000..6c99166d
--- /dev/null
+++ b/test/resources/config_legacy_modules.yaml
@@ -0,0 +1,3 @@
+modules:
+  - module_1
+  - module_2
\ No newline at end of file
diff --git a/test/resources/config_modules.yaml b/test/resources/config_modules.yaml
new file mode 100644
index 00000000..fb6af056
--- /dev/null
+++ b/test/resources/config_modules.yaml
@@ -0,0 +1,7 @@
+modules: 
+  use:
+    - path/to/module/dir/1
+    - path/to/module/dir/2
+  load:
+    - module_1
+    - module_2
\ No newline at end of file
diff --git a/test/config_mom5.yaml b/test/resources/config_mom5.yaml
similarity index 100%
rename from test/config_mom5.yaml
rename to test/resources/config_mom5.yaml
diff --git a/test/test_payu.py b/test/test_payu.py
index 7c188dbc..56835be0 100644
--- a/test/test_payu.py
+++ b/test/test_payu.py
@@ -125,7 +125,7 @@ def test_movetree():
 
 
 def test_read_config():
-    config_path = os.path.join('test', 'config_mom5.yaml')
+    config_path = os.path.join('test', 'resources', 'config_mom5.yaml')
     config = payu.fsops.read_config(config_path)
 
     # Test control_path is not set in read_config
@@ -148,11 +148,34 @@ def test_read_config():
 
     assert(config.pop('collate') == {})
     assert(config.pop('control_path') == os.getcwd())
+    assert(config.pop('modules') == {})
     assert(config == {})
 
     os.remove(config_tmp)
 
 
+def test_read_config_modules_legacy_option():
+    # Test transform legacy modules option
+    config_path = os.path.join('test', 'resources', 'config_legacy_modules.yaml')
+
+    config = payu.fsops.read_config(config_path)
+    modules_config = config.get('modules', {})
+
+    assert(modules_config.get('load', []) == ['module_1', 'module_2'])
+    assert(modules_config.get('use', []) == [])
+
+
+def test_read_config_modules_option():
+    # Test modules with load/use options is unchanged
+    config_path = os.path.join('test', 'resources', 'config_modules.yaml')
+
+    config = payu.fsops.read_config(config_path)
+    modules_config = config.get('modules', {})
+
+    assert(modules_config.get('load', []) == ['module_1', 'module_2'])
+    assert(modules_config.get('use', []) == ['path/to/module/dir/1', 'path/to/module/dir/2'])
+
+
 def test_make_symlink():
     tmp_path = 'tmp_file'
     tmp_sym = 'tmp_sym'
@@ -222,7 +245,7 @@ def test_parse_ldd_output():
     with open(ldd_output_path, 'r') as f:
         ldd_output = f.read()
     required_libs = payu.fsops.parse_ldd_output(ldd_output)
-    assert(len(required_libs), 4) 
+    assert(len(required_libs), 4)
     assert(required_libs['libmpi.so.40'], '/apps/openmpi/4.0.2/lib/libmpi.so.40')