From 85f5f5e1470e582a3a4640a35934c849cda55f2c Mon Sep 17 00:00:00 2001
From: Xavier Delaruelle <xavier.delaruelle@cea.fr>
Date: Sat, 28 Oct 2023 20:55:49 +0200
Subject: [PATCH] Adapt module function check for EnvironmentModules

Default check_module_function tests that module command is called from
module shell function. With EnvironmentModules v4+, module command is
usually called from the _module_raw shell function.

This commit adds a specific version of check_module_function for
EnvironmentModules class. Module command is first checked within
_module_raw shell function definition. If not found, default test (that
checks module function) is run.

Adapt "test_module_mismatch" unit test to specifically check module
command definition with EnvironmentModules.

Fixes #4368
---
 easybuild/tools/modules.py    | 23 +++++++++++++++++++++++
 easybuild/tools/run.py        |  1 +
 test/framework/modulestool.py | 15 ++++++++++++++-
 3 files changed, 38 insertions(+), 1 deletion(-)

diff --git a/easybuild/tools/modules.py b/easybuild/tools/modules.py
index 386643e706..d8e425a19d 100644
--- a/easybuild/tools/modules.py
+++ b/easybuild/tools/modules.py
@@ -1328,6 +1328,29 @@ class EnvironmentModules(EnvironmentModulesTcl):
     MAX_VERSION = None
     VERSION_REGEXP = r'^Modules\s+Release\s+(?P<version>\d\S*)\s'
 
+    def check_module_function(self, allow_mismatch=False, regex=None):
+        """Check whether selected module tool matches 'module' function definition."""
+        # Modules 5.1.0+: module command is called from _module_raw shell function
+        # Modules 4.2.0..5.0.1: module command is called from _module_raw shell function if it has
+        #   been initialized in an interactive shell session (i.e., a session attached to a tty)
+        if self.testing:
+            if '_module_raw' in os.environ:
+                out, ec = os.environ['_module_raw'], 0
+            else:
+                out, ec = None, 1
+        else:
+            cmd = "type _module_raw"
+            out, ec = run_cmd(cmd, simple=False, log_ok=False, log_all=False, force_in_dry_run=True, trace=False)
+
+        if regex is None:
+            regex = r".*%s" % os.path.basename(self.cmd)
+        mod_cmd_re = re.compile(regex, re.M)
+
+        if ec == 0 and mod_cmd_re.search(out):
+            self.log.debug("Found pattern '%s' in defined 'module' function." % mod_cmd_re.pattern)
+        else:
+            super(EnvironmentModules, self).check_module_function(allow_mismatch, regex)
+
     def check_module_output(self, cmd, stdout, stderr):
         """Check output of 'module' command, see if if is potentially invalid."""
         if "_mlstatus = False" in stdout:
diff --git a/easybuild/tools/run.py b/easybuild/tools/run.py
index 8916d80795..08b109ec15 100644
--- a/easybuild/tools/run.py
+++ b/easybuild/tools/run.py
@@ -71,6 +71,7 @@
     "sysctl -n machdep.cpu.brand_string",  # used in get_cpu_model (OS X)
     "sysctl -n machdep.cpu.vendor",  # used in get_cpu_vendor (OS X)
     "type module",  # used in ModulesTool.check_module_function
+    "type _module_raw",  # used in EnvironmentModules.check_module_function
     "ulimit -u",  # used in det_parallelism
 ]
 
diff --git a/test/framework/modulestool.py b/test/framework/modulestool.py
index 59c6872b93..66cf4ee333 100644
--- a/test/framework/modulestool.py
+++ b/test/framework/modulestool.py
@@ -39,7 +39,7 @@
 from easybuild.tools import modules, StrictVersion
 from easybuild.tools.build_log import EasyBuildError
 from easybuild.tools.filetools import read_file, which, write_file
-from easybuild.tools.modules import Lmod
+from easybuild.tools.modules import EnvironmentModules, Lmod
 from test.framework.utilities import init_config
 
 
@@ -108,6 +108,19 @@ def test_environment_command(self):
 
     def test_module_mismatch(self):
         """Test whether mismatch detection between modules tool and 'module' function works."""
+        # redefine 'module' and '_module_raw' function (deliberate mismatch with used module
+        # command in EnvironmentModules)
+        os.environ['_module_raw'] = "() {  eval `/usr/share/Modules/libexec/foo.tcl' bash $*`;\n}"
+        os.environ['module'] = "() {  _module_raw \"$@\" 2>&1;\n}"
+        error_regex = ".*pattern .* not found in defined 'module' function"
+        self.assertErrorRegex(EasyBuildError, error_regex, EnvironmentModules, testing=True)
+
+        # redefine '_module_raw' function with correct module command
+        os.environ['_module_raw'] = "() {  eval `/usr/share/Modules/libexec/modulecmd.tcl' bash $*`;\n}"
+        mt = EnvironmentModules(testing=True)
+        self.assertIsInstance(mt.loaded_modules(), list)  # dummy usage
+        del os.environ['_module_raw']
+
         # redefine 'module' function (deliberate mismatch with used module command in MockModulesTool)
         os.environ['module'] = "() {  eval `/tmp/Modules/$MODULE_VERSION/bin/modulecmd bash $*`\n}"
         error_regex = ".*pattern .* not found in defined 'module' function"