Merge pull request #5015 from Flamefire/restore-test-env

Introduce `saved_env` context manager for restoring environment in tests

Author: boegel

easybuild/base/testing.py

@@ -221,6 +221,16 @@ def mocked_stdout_stderr(self, mock_stdout=True, mock_stderr=True):
             if mock_stderr:
                 self.mock_stderr(False)
 
+    @contextmanager
+    def saved_env(self):
+        """Context manager to reset environment to state when it was entered"""
+        orig_env = os.environ.copy()
+        try:
+            yield
+        finally:
+            os.environ.clear()
+            os.environ.update(orig_env)
+
     def tearDown(self):
         """Cleanup after running a test."""
         self.mock_stdout(False)

easybuild/tools/modules.py

@@ -1095,6 +1095,8 @@ def load(self, modules, mod_paths=None, purge=False, init_env=None, allow_reload
         :param init_env: original environment to restore after running 'module purge'
         :param allow_reload: allow reloading an already loaded module
         """
+        if not any((modules, mod_paths, purge)):
+            return  # Avoid costly module paths if nothing to do
         if mod_paths is None:
             mod_paths = []
 

test/framework/easyblock.py

@@ -30,7 +30,6 @@
 @author: Maxime Boissonneault (Compute Canada)
 @author: Jan Andre Reuter (Juelich Supercomputing Centre)
 """
-import copy
 import fileinput
 import os
 import re
@@ -53,7 +52,6 @@
 from easybuild.tools import LooseVersion, config
 from easybuild.tools.build_log import EasyBuildError
 from easybuild.tools.config import get_module_syntax, update_build_option
-from easybuild.tools.environment import modify_env
 from easybuild.tools.filetools import change_dir, copy_dir, copy_file, mkdir, read_file, remove_dir, remove_file
 from easybuild.tools.filetools import symlink, verify_checksum, write_file
 from easybuild.tools.module_generator import module_generator
@@ -1290,14 +1288,14 @@ def test_post_processing_step(self):
         eb.silent = True
         depr_msg = r"EasyBlock.post_install_step\(\) is deprecated, use EasyBlock.post_processing_step\(\) instead"
         expected_error = r"DEPRECATED \(since v6.0\).*" + depr_msg
-        with self.mocked_stdout_stderr():
+        with self.mocked_stdout_stderr(), self.saved_env():
             self.assertErrorRegex(EasyBuildError, expected_error, eb.run_all_steps, True)
 
         change_dir(cwd)
         toy_ec = EasyConfig(toy_ec_fn)
         eb = EB_toy(toy_ec)
         eb.silent = True
-        with self.mocked_stdout_stderr() as (_, stderr):
+        with self.mocked_stdout_stderr() as (_, stderr), self.saved_env():
             eb.run_all_steps(True)
             # no deprecation warning
             stderr = stderr.getvalue()
@@ -1309,14 +1307,13 @@ def test_post_processing_step(self):
         # check again with toy easyblock that still uses post_install_step,
         # to verify that the expected file is being created when deprecated functionality is allow
         remove_file(libtoy_post_a)
-        modify_env(os.environ, self.orig_environ, verbose=False)
         change_dir(cwd)
 
         self.allow_deprecated_behaviour()
         toy_ec = EasyConfig(toy_ec_fn)
         eb = EB_toy_deprecated(toy_ec)
         eb.silent = True
-        with self.mocked_stdout_stderr() as (stdout, stderr):
+        with self.mocked_stdout_stderr() as (stdout, stderr), self.saved_env():
             eb.run_all_steps(True)
 
             regex = re.compile(depr_msg, re.M)
@@ -2606,7 +2603,7 @@ def test_extensions_sanity_check(self):
         error_pattern = r"Sanity check failed: extensions sanity check failed for 1 extensions: toy\n"
         error_pattern += r"failing sanity check for 'toy' extension: "
         error_pattern += r'command "thisshouldfail" failed; output:\n.* thisshouldfail: command not found'
-        with self.mocked_stdout_stderr():
+        with self.mocked_stdout_stderr(), self.saved_env():
             self.assertErrorRegex(EasyBuildError, error_pattern, eb.run_all_steps, True)
 
         # purposely put sanity check command in place that breaks the build,
@@ -2617,7 +2614,7 @@ def test_extensions_sanity_check(self):
         toy_ec['exts_defaultclass'] = 'DummyExtension'
         eb = EB_toy(toy_ec)
         eb.silent = True
-        with self.mocked_stdout_stderr():
+        with self.mocked_stdout_stderr(), self.saved_env():
             eb.run_all_steps(True)
 
     def test_parallel(self):
@@ -2976,23 +2973,19 @@ def test_extension_patch_step(self):
 
         cwd = os.getcwd()
         self.assertExists(cwd)
-        # Take environment with test-specific variable set up
-        orig_environ = copy.deepcopy(os.environ)
 
         def run_extension_step():
-            try:
-                change_dir(cwd)
-                eb = EasyBlock(ec)
-                # Cleanup build directory
-                if os.path.exists(eb.builddir):
-                    remove_dir(eb.builddir)
+            change_dir(cwd)
+            eb = EasyBlock(ec)
+            # Cleanup build directory
+            if os.path.exists(eb.builddir):
+                remove_dir(eb.builddir)
+            # restore original environment to continue testing with a clean slate
+            with self.saved_env():
                 eb.make_builddir()
                 eb.update_config_template_run_step()
                 eb.extensions_step(fetch=True, install=True)
-                return os.path.join(eb.builddir)
-            finally:
-                # restore original environment to continue testing with a clean slate
-                modify_env(os.environ, orig_environ, verbose=False)
+            return os.path.join(eb.builddir)
 
         ec['exts_defaultclass'] = 'DummyExtension'
         ec['exts_list'] = [('toy', '0.0', {'easyblock': 'DummyExtension'})]

test/framework/modules.py

@@ -44,7 +44,6 @@
 from easybuild.framework.easyconfig.easyconfig import EasyConfig
 from easybuild.tools import LooseVersion
 from easybuild.tools.build_log import EasyBuildError
-from easybuild.tools.environment import modify_env
 from easybuild.tools.filetools import adjust_permissions, copy_file, copy_dir, mkdir
 from easybuild.tools.filetools import read_file, remove_dir, remove_file, symlink, write_file
 from easybuild.tools.modules import EnvironmentModules, EnvironmentModulesC, EnvironmentModulesTcl, Lmod, NoModulesTool
@@ -103,11 +102,10 @@ def test_run_module(self):
             os.environ.pop(key, None)
 
         # arguments can be passed in two ways: multiple arguments, or just 1 list argument
-        self.modtool.run_module('load', 'GCC/6.4.0-2.28')
-        self.assertEqual(os.environ['EBROOTGCC'], '/prefix/software/GCC/6.4.0-2.28')
+        with self.saved_env():
+            self.modtool.run_module('load', 'GCC/6.4.0-2.28')
+            self.assertEqual(os.environ['EBROOTGCC'], '/prefix/software/GCC/6.4.0-2.28')
 
-        # restore original environment
-        modify_env(os.environ, self.orig_environ, verbose=False)
         self.reset_modulepath([os.path.join(testdir, 'modules')])
 
         self.assertNotIn('EBROOTGCC', os.environ)

test/framework/options.py

@@ -53,7 +53,6 @@
 from easybuild.tools.build_log import EasyBuildError, EasyBuildLog
 from easybuild.tools.config import DEFAULT_MODULECLASSES, BuildOptions, ConfigurationVariables
 from easybuild.tools.config import build_option, find_last_log, get_build_log_path, get_module_syntax, module_classes
-from easybuild.tools.environment import modify_env
 from easybuild.tools.filetools import adjust_permissions, change_dir, copy_dir, copy_file, download_file
 from easybuild.tools.filetools import is_patch_file, mkdir, move_file, parse_http_header_fields_urlpat
 from easybuild.tools.filetools import read_file, remove_dir, remove_file, which, write_file
@@ -4324,7 +4323,6 @@ def check_tmpdir(tmpdir):
             # cleanup
             os.close(fd)
             shutil.rmtree(mytmpdir)
-            modify_env(os.environ, self.orig_environ)
             tempfile.tempdir = None
 
         orig_tmpdir = tempfile.gettempdir()
@@ -4335,7 +4333,8 @@ def check_tmpdir(tmpdir):
             os.path.join(orig_tmpdir, '[ab @cd]%/#*'),
         ]
         for tmpdir in cand_tmpdirs:
-            check_tmpdir(tmpdir)
+            with self.saved_env():
+                check_tmpdir(tmpdir)
 
     def test_minimal_toolchains(self):
         """End-to-end test for --minimal-toolchains."""

test/framework/systemtools.py

@@ -28,7 +28,6 @@
 @author: Kenneth hoste (Ghent University)
 @author: Ward Poelmans (Ghent University)
 """
-import copy
 import ctypes
 import logging
 import os
@@ -41,7 +40,7 @@
 
 import easybuild.tools.systemtools as st
 from easybuild.tools.build_log import EasyBuildError
-from easybuild.tools.environment import modify_env, setvar
+from easybuild.tools.environment import setvar
 from easybuild.tools.filetools import adjust_permissions, mkdir, read_file, symlink, which, write_file
 from easybuild.tools.run import RunShellCmdResult, run_shell_cmd
 from easybuild.tools.systemtools import CPU_ARCHITECTURES, AARCH32, AARCH64, POWER, X86_64
@@ -740,9 +739,9 @@ def test_cpu_architecture(self):
             'x86_64': X86_64,
             'some_fancy_arch': UNKNOWN,
         }
-        for name in machine_names:
+        for name, arch in machine_names.items():
             MACHINE_NAME = name
-            self.assertEqual(get_cpu_architecture(), machine_names[name])
+            self.assertEqual(get_cpu_architecture(), arch)
 
     def test_cpu_arch_name_native(self):
         """Test getting CPU arch name."""
@@ -1323,9 +1322,6 @@ def test_find_library_path(self):
 
     def test_get_cuda_object_dump_raw(self):
         """Test get_cuda_object_dump_raw function"""
-        # This test modifies environment, make sure we can revert the changes:
-        start_env = copy.deepcopy(os.environ)
-
         # Mock the shell command for certain known commands
         st.run_shell_cmd = mocked_run_shell_cmd
 
@@ -1373,14 +1369,8 @@ def test_get_cuda_object_dump_raw(self):
         # Test case 7: call on CUDA static lib, which only contains device code
         self.assertEqual(get_cuda_object_dump_raw('mock_cuda_staticlib'), CUOBJDUMP_DEVICE_CODE_ONLY)
 
-        # Restore original environment
-        modify_env(os.environ, start_env, verbose=False)
-
     def test_get_cuda_architectures(self):
         """Test get_cuda_architectures function"""
-        # This test modifies environment, make sure we can revert the changes:
-        start_env = copy.deepcopy(os.environ)
-
         # Mock the shell command for certain known commands
         st.run_shell_cmd = mocked_run_shell_cmd
 
@@ -1450,9 +1440,6 @@ def test_get_cuda_architectures(self):
         self.assertTrue(warning_regex_elf.search(logtxt), fail_msg)
         self.assertIsNone(res_elf)
 
-        # Restore original environment
-        modify_env(os.environ, start_env, verbose=False)
-
     def test_get_linked_libs_raw(self):
         """
         Test get_linked_libs_raw function.

test/framework/toy_build.py

@@ -29,7 +29,6 @@
 @author: Kenneth Hoste (Ghent University)
 @author: Damian Alvarez (Forschungszentrum Juelich GmbH)
 """
-import copy
 import glob
 import grp
 import os
@@ -54,7 +53,7 @@
 from easybuild.main import main_with_hooks
 from easybuild.tools.build_log import EasyBuildError
 from easybuild.tools.config import get_module_syntax, get_repositorypath
-from easybuild.tools.environment import modify_env, setvar
+from easybuild.tools.environment import setvar
 from easybuild.tools.filetools import adjust_permissions, change_dir, copy_file, mkdir, move_file
 from easybuild.tools.filetools import read_file, remove_dir, remove_file, which, write_file
 from easybuild.tools.module_generator import ModuleGeneratorTcl
@@ -1759,18 +1758,14 @@ def test_external_dependencies(self):
         installed_test_modules = os.path.join(self.test_installpath, 'modules', 'all')
         self.reset_modulepath([modulepath, installed_test_modules])
 
-        start_env = copy.deepcopy(os.environ)
-
         with self.mocked_stdout_stderr():
             self._test_toy_build(ec_file=toy_ec, versionsuffix='-external-deps', verbose=True, raise_error=True)
 
-        self.modtool.load(['toy/0.0-external-deps'])
-        # note build dependency is not loaded
-        mods = ['intel/2018a', 'GCC/6.4.0-2.28', 'foobar/1.2.3', 'toy/0.0-external-deps']
-        self.assertEqual([x['mod_name'] for x in self.modtool.list()], mods)
-
-        # restore original environment (to undo 'module load' done above)
-        modify_env(os.environ, start_env, verbose=False)
+        with self.saved_env():
+            self.modtool.load(['toy/0.0-external-deps'])
+            # note build dependency is not loaded
+            mods = ['intel/2018a', 'GCC/6.4.0-2.28', 'foobar/1.2.3', 'toy/0.0-external-deps']
+            self.assertEqual([x['mod_name'] for x in self.modtool.list()], mods)
 
         # check behaviour when a non-existing external (build) dependency is included
         extraectxt = "\nbuilddependencies = [('nosuchbuilddep/0.0.0', EXTERNAL_MODULE)]"
@@ -3147,10 +3142,6 @@ def test_toy_filter_rpath_sanity_libs(self):
 
     def test_toy_cuda_sanity_check(self):
         """Test the CUDA sanity check"""
-        # We need to mock a cuobjdump executable and prepend in on the PATH
-        # First, make sure we can restore environment at the end of this test
-        start_env = copy.deepcopy(os.environ)
-
         # Define the toy_ec file we want to use
         topdir = os.path.dirname(os.path.abspath(__file__))
         toy_ec = os.path.join(topdir, 'easyconfigs', 'test_ecs', 't', 'toy', 'toy-0.0.eb')
@@ -3540,9 +3531,6 @@ def assert_cuda_report(missing_cc, additional_cc, missing_ptx, log, stdout=None,
         self.assertTrue(expected_result.search(outtxt), msg)
         assert_cuda_report(missing_cc=0, additional_cc=0, missing_ptx=0, log=outtxt, stdout=stdout, num_checked=0)
 
-        # Restore original environment
-        modify_env(os.environ, start_env, verbose=False)
-
     def test_toy_modaltsoftname(self):
         """Build two dependent toys as in test_toy_toy but using modaltsoftname"""
         topdir = os.path.dirname(os.path.abspath(__file__))
@@ -3953,7 +3941,8 @@ def check_toy_load(depends_on=False):
                 # just undo
                 self.modtool.unload(['toy/0.0', 'GCC/4.6.3'])
 
-        check_toy_load()
+        with self.saved_env():
+            check_toy_load()
 
         # this behaviour can be disabled via "multi_dep_load_defaults = False"
         write_file(test_ec, test_ec_txt + "\nmulti_deps_load_default = False")
@@ -3965,8 +3954,9 @@ def check_toy_load(depends_on=False):
 
         self.assertNotIn(expected, toy_mod_txt)
 
-        self.modtool.load(['toy/0.0'])
-        loaded_mod_names = [x['mod_name'] for x in self.modtool.list()]
+        with self.saved_env():
+            self.modtool.load(['toy/0.0'])
+            loaded_mod_names = [x['mod_name'] for x in self.modtool.list()]
         self.assertIn('toy/0.0', loaded_mod_names)
         self.assertNotIn('GCC/4.6.3', loaded_mod_names)
         self.assertNotIn('GCC/7.3.0-2.30', loaded_mod_names)
@@ -3989,10 +3979,6 @@ def check_toy_load(depends_on=False):
         error_msg_whatis = "Pattern '%s' should be found in: %s" % (expected_whatis_no_default, toy_mod_txt)
         self.assertIn(expected_whatis_no_default, toy_mod_txt, error_msg_whatis)
 
-        # restore original environment to continue testing with a clean slate
-        modify_env(os.environ, self.orig_environ, verbose=False)
-        self.modtool.use(test_mod_path)
-
         # disable showing of progress bars (again), doesn't make sense when running tests
         os.environ['EASYBUILD_DISABLE_SHOW_PROGRESS_BAR'] = '1'
 

test/framework/utilities.py

@@ -501,9 +501,9 @@ def init_config(args=None, build_options=None, with_include=True, clear_caches=T
         'valid_module_classes': module_classes(),
         'valid_stops': [x[0] for x in EasyBlock.get_steps()],
     }
-    for key in default_build_options:
+    for key, def_option in default_build_options.items():
         if key not in build_options:
-            build_options[key] = default_build_options[key]
+            build_options[key] = def_option
 
     config.init_build_options(build_options=build_options)