Introduce context manager for restoring environment in tests

Many tests modify the environment in some steps and need to restore it
afterwards.
To avoid missing this the contextmanager does that automatically.

Author: Flamefire

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)

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)
@@ -2605,7 +2602,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,
@@ -2615,7 +2612,7 @@ def test_extensions_sanity_check(self):
         toy_ec.update('sanity_check_commands', [("%(installdir)s/bin/toy && rm %(installdir)s/bin/toy", '')])
         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)
 
         # Verify custom paths and commands of extensions are checked
@@ -2643,7 +2640,7 @@ def test_extensions_sanity_check(self):
         eb = EB_toy(EasyConfig(self.eb_file))
         eb.silent = True
         write_file(os.path.join(eb.installdir, 'any_file'), '')
-        with self.mocked_stdout_stderr():
+        with self.mocked_stdout_stderr(), self.saved_env():
             eb.sanity_check_step()
         logtxt = read_file(eb.logfile)
         self.assertRegex(logtxt, 'Running .*command.*echo "toy output"')
@@ -3021,23 +3018,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
@@ -4322,7 +4321,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()
@@ -4333,7 +4331,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

@@ -54,7 +54,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
@@ -1751,18 +1751,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)]"
@@ -3136,10 +3132,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')
@@ -3529,9 +3521,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__))
@@ -3940,7 +3929,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")
@@ -3952,8 +3942,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)
@@ -3976,10 +3967,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)