diff options
| -rw-r--r-- | setuptools/tests/contexts.py | 30 | ||||
| -rw-r--r-- | setuptools/tests/test_msvc9compiler.py | 31 |
2 files changed, 26 insertions, 35 deletions
diff --git a/setuptools/tests/contexts.py b/setuptools/tests/contexts.py index 10691606..a604cd41 100644 --- a/setuptools/tests/contexts.py +++ b/setuptools/tests/contexts.py @@ -21,15 +21,31 @@ def tempdir(cd=lambda dir:None, **kwargs): @contextlib.contextmanager -def environment(**updates): - old_env = os.environ.copy() - os.environ.update(updates) +def environment(**replacements): + """ + In a context, patch the environment with replacements. Pass None values + to clear the values. + """ + saved = dict( + (key, os.environ['key']) + for key in replacements + if key in os.environ + ) + + # remove values that are null + remove = (key for (key, value) in replacements.items() if value is None) + for key in list(remove): + os.environ.pop(key, None) + replacements.pop(key) + + os.environ.update(replacements) + try: - yield + yield saved finally: - for key in updates: - del os.environ[key] - os.environ.update(old_env) + for key in replacements: + os.environ.pop(key, None) + os.environ.update(saved) @contextlib.contextmanager diff --git a/setuptools/tests/test_msvc9compiler.py b/setuptools/tests/test_msvc9compiler.py index 27747512..2a117dc9 100644 --- a/setuptools/tests/test_msvc9compiler.py +++ b/setuptools/tests/test_msvc9compiler.py @@ -10,10 +10,11 @@ import shutil import tempfile import unittest import distutils.errors -import contextlib import pytest +from . import contexts + # importing only setuptools should apply the patch __import__('setuptools') @@ -62,32 +63,6 @@ class MockReg: distutils.msvc9compiler.Reg.read_keys = self.original_read_keys distutils.msvc9compiler.Reg.read_values = self.original_read_values -@contextlib.contextmanager -def patch_env(**replacements): - """ - In a context, patch the environment with replacements. Pass None values - to clear the values. - """ - saved = dict( - (key, os.environ['key']) - for key in replacements - if key in os.environ - ) - - # remove values that are null - remove = (key for (key, value) in replacements.items() if value is None) - for key in list(remove): - os.environ.pop(key, None) - replacements.pop(key) - - os.environ.update(replacements) - - try: - yield saved - finally: - for key in replacements: - os.environ.pop(key, None) - os.environ.update(saved) class TestMSVC9Compiler(unittest.TestCase): @@ -100,7 +75,7 @@ class TestMSVC9Compiler(unittest.TestCase): # No registry entries or environment variable means we should # not find anything - with patch_env(VS90COMNTOOLS=None): + with contexts.environment(VS90COMNTOOLS=None): with MockReg(): self.assertIsNone(find_vcvarsall(9.0)) |