Always cleanup the changed environ

saltstack · May 8, 2019 · ff12fef · ff12fef
1 parent bb48867
commit ff12fef
Show file tree

Hide file tree

Showing 5 changed files with 56 additions and 38 deletions.
diff --git a/tests/integration/modules/test_pip.py b/tests/integration/modules/test_pip.py
@@ -28,8 +28,13 @@ class PipModuleTest(ModuleCase):
 
     def setUp(self):
         super(PipModuleTest, self).setUp()
+
         # Restore the environ
-        self.addCleanup(os.environ.update, os.environ.copy())
+        def cleanup_environ(environ):
+            os.environ.clear()
+            os.environ.update(environ)
+
+        self.addCleanup(cleanup_environ, os.environ.copy())
 
         self.venv_test_dir = tempfile.mkdtemp(dir=TMP)
         # Remove the venv test directory

diff --git a/tests/integration/states/test_pip_state.py b/tests/integration/states/test_pip_state.py
@@ -139,37 +139,34 @@ def test_pip_installed_errors(self):
         venv_dir = os.path.join(
             RUNTIME_VARS.TMP, 'pip-installed-errors'
         )
-        orig_shell = os.environ.get('SHELL')
-        try:
-            # Since we don't have the virtualenv created, pip.installed will
-            # throw an error.
-            # Example error strings:
-            #  * "Error installing 'pep8': /tmp/pip-installed-errors: not found"
-            #  * "Error installing 'pep8': /bin/sh: 1: /tmp/pip-installed-errors: not found"
-            #  * "Error installing 'pep8': /bin/bash: /tmp/pip-installed-errors: No such file or directory"
-            os.environ['SHELL'] = '/bin/sh'
-            ret = self.run_function('state.sls', mods='pip-installed-errors')
-            self.assertSaltFalseReturn(ret)
-            self.assertSaltCommentRegexpMatches(
-                ret,
-                'Error installing \'pep8\':'
-            )
 
-            # We now create the missing virtualenv
-            ret = self._create_virtualenv(venv_dir)
-            self.assertEqual(ret['retcode'], 0)
+        def cleanup_environ(environ):
+            os.environ.clear()
+            os.environ.update(environ)
 
-            # The state should not have any issues running now
-            ret = self.run_function('state.sls', mods='pip-installed-errors')
-            self.assertSaltTrueReturn(ret)
-        finally:
-            if orig_shell is None:
-                # Didn't exist before, don't leave it there. This should never
-                # happen, but if it does, we don't want this test to affect
-                # others elsewhere in the suite.
-                os.environ.pop('SHELL')
-            else:
-                os.environ['SHELL'] = orig_shell
+        self.addCleanup(cleanup_environ, os.environ.copy())
+
+        # Since we don't have the virtualenv created, pip.installed will
+        # throw an error.
+        # Example error strings:
+        #  * "Error installing 'pep8': /tmp/pip-installed-errors: not found"
+        #  * "Error installing 'pep8': /bin/sh: 1: /tmp/pip-installed-errors: not found"
+        #  * "Error installing 'pep8': /bin/bash: /tmp/pip-installed-errors: No such file or directory"
+        os.environ['SHELL'] = '/bin/sh'
+        ret = self.run_function('state.sls', mods='pip-installed-errors')
+        self.assertSaltFalseReturn(ret)
+        self.assertSaltCommentRegexpMatches(
+            ret,
+            'Error installing \'pep8\':'
+        )
+
+        # We now create the missing virtualenv
+        ret = self._create_virtualenv(venv_dir)
+        self.assertEqual(ret['retcode'], 0)
+
+        # The state should not have any issues running now
+        ret = self.run_function('state.sls', mods='pip-installed-errors')
+        self.assertSaltTrueReturn(ret)
 
     @skipIf(six.PY3, 'Issue is specific to carbon module, which is PY2-only')
     @skipIf(salt.utils.is_windows(), "Carbon does not install in Windows")

diff --git a/tests/support/gitfs.py b/tests/support/gitfs.py
@@ -69,7 +69,6 @@
     'git_pillar_includes': True,
 }
 PROC_TIMEOUT = 10
-NOTSET = object()
 
 
 class ProcessManager(object):
@@ -543,14 +542,15 @@ def get_pillar(self, ext_pillar_conf):
         passphraselsess key is used to auth without needing to modify the root
         user's ssh config file.
         '''
-        orig_git_ssh = os.environ.pop('GIT_SSH', NOTSET)
+
+        def cleanup_environ(environ):
+            os.environ.clear()
+            os.environ.update(environ)
+
+        self.addCleanup(cleanup_environ, os.environ.copy())
+
         os.environ['GIT_SSH'] = self.git_ssh
-        try:
-            return super(GitPillarSSHTestBase, self).get_pillar(ext_pillar_conf)
-        finally:
-            os.environ.pop('GIT_SSH', None)
-            if orig_git_ssh is not NOTSET:
-                os.environ['GIT_SSH'] = orig_git_ssh
+        return super(GitPillarSSHTestBase, self).get_pillar(ext_pillar_conf)
 
 
 class GitPillarHTTPTestBase(GitPillarTestBase, WebserverMixin):

diff --git a/tests/unit/config/test_config.py b/tests/unit/config/test_config.py
@@ -76,6 +76,11 @@ def _salt_configuration_error(filename):
 
 class ConfigTestCase(TestCase, AdaptedConfigurationTestCaseMixin):
 
+
+    def _cleanup_environ(self, environ):
+        os.environ.clear()
+        os.environ.update(environ)
+
     def test_sha256_is_default_for_master(self):
         fpath = tempfile.mktemp()
         try:
@@ -151,6 +156,7 @@ def test_common_prefix_stripping(self):
         salt.utils.is_windows(),
         'You can\'t set an environment dynamically in Windows')
     def test_load_master_config_from_environ_var(self):
+        self.addCleanup(self._cleanup_environ, os.environ.copy())
         original_environ = os.environ.copy()
 
         tempdir = tempfile.mkdtemp(dir=TMP)
@@ -198,6 +204,7 @@ def test_load_master_config_from_environ_var(self):
         salt.utils.is_windows(),
         'You can\'t set an environment dynamically in Windows')
     def test_load_minion_config_from_environ_var(self):
+        self.addCleanup(self._cleanup_environ, os.environ.copy())
         original_environ = os.environ.copy()
 
         tempdir = tempfile.mkdtemp(dir=TMP)
@@ -241,6 +248,7 @@ def test_load_minion_config_from_environ_var(self):
                 shutil.rmtree(tempdir)
 
     def test_load_client_config_from_environ_var(self):
+        self.addCleanup(self._cleanup_environ, os.environ.copy())
         original_environ = os.environ.copy()
         try:
             tempdir = tempfile.mkdtemp(dir=TMP)
@@ -1010,6 +1018,7 @@ def test_is_provider_configured_multiple_success(self):
         salt.utils.is_windows(),
         'You can\'t set an environment dynamically in Windows')
     def test_load_cloud_config_from_environ_var(self):
+        self.addCleanup(self._cleanup_environ, os.environ.copy())
         original_environ = os.environ.copy()
 
         tempdir = tempfile.mkdtemp(dir=TMP)

diff --git a/tests/unit/fileserver/test_gitfs.py b/tests/unit/fileserver/test_gitfs.py
@@ -206,6 +206,13 @@ def setUp(self):
             repo = git.Repo.init(self.tmp_repo_dir)
 
         if 'USERNAME' not in os.environ:
+
+            def cleanup_environ(environ):
+                os.environ.clear()
+                os.environ.update(environ)
+
+            self.addCleanup(cleanup_environ, os.environ.copy())
+
             try:
                 import salt.utils
                 os.environ['USERNAME'] = this_user()