From 6733bed57e780008f8c78422d2a9676b9a2710cf Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Walter=20D=C3=B6rwald?= Date: Fri, 1 May 2009 17:35:37 +0000 Subject: Make test.test_support.EnvironmentVarGuard behave like a dictionary. All changes are mirrored to the underlying os.environ dict, but rolled back on exit from the with block. --- Doc/library/test.rst | 8 +++++++- Lib/test/test_getopt.py | 16 +++++++--------- Lib/test/test_gettext.py | 10 ++++------ Lib/test/test_ntpath.py | 14 +++++--------- Lib/test/test_optparse.py | 4 ++-- Lib/test/test_posixpath.py | 16 ++++++---------- Lib/test/test_support.py | 36 +++++++++++++++++++++++++----------- Lib/test/test_tcl.py | 25 +++++++++++-------------- Lib/test/test_tempfile.py | 2 +- Lib/test/test_urllib.py | 2 +- Lib/test/test_xmlrpc.py | 4 ++-- 11 files changed, 71 insertions(+), 66 deletions(-) diff --git a/Doc/library/test.rst b/Doc/library/test.rst index ba2c3b8..b4a6c7a 100644 --- a/Doc/library/test.rst +++ b/Doc/library/test.rst @@ -401,9 +401,14 @@ The :mod:`test.test_support` module defines the following classes: .. class:: EnvironmentVarGuard() Class used to temporarily set or unset environment variables. Instances can be - used as a context manager. + used as a context manager and have a complete dictionary interface for + querying/modifying the underlying ``os.environ``. After exit from the context + manager all changes to environment variables done through this instance will + be rolled back. .. versionadded:: 2.6 + .. versionchanged:: 2.7 + Added dictionary interface. .. method:: EnvironmentVarGuard.set(envvar, value) @@ -415,6 +420,7 @@ The :mod:`test.test_support` module defines the following classes: Temporarily unset the environment variable ``envvar``. + .. class:: WarningsRecorder() Class used to record warnings for unit tests. See documentation of diff --git a/Lib/test/test_getopt.py b/Lib/test/test_getopt.py index aac720d..9031058 100644 --- a/Lib/test/test_getopt.py +++ b/Lib/test/test_getopt.py @@ -1,7 +1,7 @@ # test_getopt.py # David Goodger 2000-08-19 -from test.test_support import verbose, run_doctest, run_unittest +from test.test_support import verbose, run_doctest, run_unittest, EnvironmentVarGuard import unittest import getopt @@ -11,15 +11,13 @@ sentinel = object() class GetoptTests(unittest.TestCase): def setUp(self): - self.old_posixly_correct = os.environ.get("POSIXLY_CORRECT", sentinel) - if self.old_posixly_correct is not sentinel: - del os.environ["POSIXLY_CORRECT"] + self.env = EnvironmentVarGuard() + if "POSIXLY_CORRECT" in self.env: + del self.env["POSIXLY_CORRECT"] def tearDown(self): - if self.old_posixly_correct is sentinel: - os.environ.pop("POSIXLY_CORRECT", None) - else: - os.environ["POSIXLY_CORRECT"] = self.old_posixly_correct + self.env.__exit__() + del self.env def assertError(self, *args, **kwargs): self.assertRaises(getopt.GetoptError, *args, **kwargs) @@ -135,7 +133,7 @@ class GetoptTests(unittest.TestCase): self.assertEqual(args, ['arg1', '-b', '1', '--alpha', '--beta=2']) # Posix style via POSIXLY_CORRECT - os.environ["POSIXLY_CORRECT"] = "1" + self.env["POSIXLY_CORRECT"] = "1" opts, args = getopt.gnu_getopt(cmdline, 'ab:', ['alpha', 'beta=']) self.assertEqual(opts, [('-a', '')]) self.assertEqual(args, ['arg1', '-b', '1', '--alpha', '--beta=2']) diff --git a/Lib/test/test_gettext.py b/Lib/test/test_gettext.py index ab6bc9a..dca99cb 100644 --- a/Lib/test/test_gettext.py +++ b/Lib/test/test_gettext.py @@ -58,10 +58,6 @@ LOCALEDIR = os.path.join('xx', 'LC_MESSAGES') MOFILE = os.path.join(LOCALEDIR, 'gettext.mo') UMOFILE = os.path.join(LOCALEDIR, 'ugettext.mo') MMOFILE = os.path.join(LOCALEDIR, 'metadata.mo') -try: - LANG = os.environ['LANGUAGE'] -except: - LANG = 'en' class GettextBaseTest(unittest.TestCase): @@ -77,10 +73,12 @@ class GettextBaseTest(unittest.TestCase): fp = open(MMOFILE, 'wb') fp.write(base64.decodestring(MMO_DATA)) fp.close() - os.environ['LANGUAGE'] = 'xx' + self.env = test_support.EnvironmentVarGuard() + self.env['LANGUAGE'] = 'xx' def tearDown(self): - os.environ['LANGUAGE'] = LANG + self.env.__exit__() + del self.env shutil.rmtree(os.path.split(LOCALEDIR)[0]) diff --git a/Lib/test/test_ntpath.py b/Lib/test/test_ntpath.py index a077c78..ade0790 100644 --- a/Lib/test/test_ntpath.py +++ b/Lib/test/test_ntpath.py @@ -124,12 +124,11 @@ class TestNtpath(unittest.TestCase): tester("ntpath.normpath('//machine/share//a/b')", r'\\machine\share\a\b') def test_expandvars(self): - oldenv = os.environ.copy() - try: - os.environ.clear() - os.environ["foo"] = "bar" - os.environ["{foo"] = "baz1" - os.environ["{foo}"] = "baz2" + with test_support.EnvironmentVarGuard() as env: + env.clear() + env["foo"] = "bar" + env["{foo"] = "baz1" + env["{foo}"] = "baz2" tester('ntpath.expandvars("foo")', "foo") tester('ntpath.expandvars("$foo bar")', "bar bar") tester('ntpath.expandvars("${foo}bar")', "barbar") @@ -149,9 +148,6 @@ class TestNtpath(unittest.TestCase): tester('ntpath.expandvars("%?bar%")', "%?bar%") tester('ntpath.expandvars("%foo%%bar")', "bar%bar") tester('ntpath.expandvars("\'%foo%\'%bar")', "\'%foo%\'%bar") - finally: - os.environ.clear() - os.environ.update(oldenv) def test_abspath(self): # ntpath.abspath() can only be used on a system with the "nt" module diff --git a/Lib/test/test_optparse.py b/Lib/test/test_optparse.py index f27864a..7483fe7 100644 --- a/Lib/test/test_optparse.py +++ b/Lib/test/test_optparse.py @@ -1465,7 +1465,7 @@ class TestHelp(BaseTest): # screws things up for other tests when it's part of the Python # test suite. with test_support.EnvironmentVarGuard() as env: - env.set('COLUMNS', str(columns)) + env['COLUMNS'] = str(columns) return InterceptingOptionParser(option_list=options) def assertHelpEquals(self, expected_output): @@ -1494,7 +1494,7 @@ class TestHelp(BaseTest): def test_help_title_formatter(self): with test_support.EnvironmentVarGuard() as env: - env.set("COLUMNS", "80") + env["COLUMNS"] = "80" self.parser.formatter = TitledHelpFormatter() self.assertHelpEquals(_expected_help_title_formatter) diff --git a/Lib/test/test_posixpath.py b/Lib/test/test_posixpath.py index b7fbd50..0a0a48b 100644 --- a/Lib/test/test_posixpath.py +++ b/Lib/test/test_posixpath.py @@ -346,18 +346,17 @@ class PosixPathTest(unittest.TestCase): self.assert_(isinstance(posixpath.expanduser("~foo/"), basestring)) with test_support.EnvironmentVarGuard() as env: - env.set('HOME', '/') + env['HOME'] = '/' self.assertEqual(posixpath.expanduser("~"), "/") self.assertRaises(TypeError, posixpath.expanduser) def test_expandvars(self): - oldenv = os.environ.copy() - try: - os.environ.clear() - os.environ["foo"] = "bar" - os.environ["{foo"] = "baz1" - os.environ["{foo}"] = "baz2" + with test_support.EnvironmentVarGuard() as env: + env.clear() + env["foo"] = "bar" + env["{foo"] = "baz1" + env["{foo}"] = "baz2" self.assertEqual(posixpath.expandvars("foo"), "foo") self.assertEqual(posixpath.expandvars("$foo bar"), "bar bar") self.assertEqual(posixpath.expandvars("${foo}bar"), "barbar") @@ -370,9 +369,6 @@ class PosixPathTest(unittest.TestCase): self.assertEqual(posixpath.expandvars("${{foo}}"), "baz1}") self.assertEqual(posixpath.expandvars("$foo$foo"), "barbar") self.assertEqual(posixpath.expandvars("$bar$bar"), "$bar$bar") - finally: - os.environ.clear() - os.environ.update(oldenv) self.assertRaises(TypeError, posixpath.expandvars) diff --git a/Lib/test/test_support.py b/Lib/test/test_support.py index c1b1c69..7d0ec8d 100644 --- a/Lib/test/test_support.py +++ b/Lib/test/test_support.py @@ -13,6 +13,7 @@ import shutil import warnings import unittest import importlib +import UserDict __all__ = ["Error", "TestFailed", "ResourceDenied", "import_module", "verbose", "use_resources", "max_memuse", "record_original_stdout", @@ -526,26 +527,39 @@ class CleanImport(object): sys.modules.update(self.original_modules) -class EnvironmentVarGuard(object): +class EnvironmentVarGuard(UserDict.DictMixin): """Class to help protect the environment variable properly. Can be used as a context manager.""" def __init__(self): + self._environ = os.environ self._changed = {} - def set(self, envvar, value): + def __getitem__(self, envvar): + return self._environ[envvar] + + def __setitem__(self, envvar, value): # Remember the initial value on the first access if envvar not in self._changed: - self._changed[envvar] = os.environ.get(envvar) - os.environ[envvar] = value + self._changed[envvar] = self._environ.get(envvar) + self._environ[envvar] = value - def unset(self, envvar): + def __delitem__(self, envvar): # Remember the initial value on the first access if envvar not in self._changed: - self._changed[envvar] = os.environ.get(envvar) - if envvar in os.environ: - del os.environ[envvar] + self._changed[envvar] = self._environ.get(envvar) + if envvar in self._environ: + del self._environ[envvar] + + def keys(self): + return self._environ.keys() + + def set(self, envvar, value): + self[envvar] = value + + def unset(self, envvar): + del self[envvar] def __enter__(self): return self @@ -553,10 +567,10 @@ class EnvironmentVarGuard(object): def __exit__(self, *ignore_exc): for (k, v) in self._changed.items(): if v is None: - if k in os.environ: - del os.environ[k] + if k in self._environ: + del self._environ[k] else: - os.environ[k] = v + self._environ[k] = v class TransientResource(object): diff --git a/Lib/test/test_tcl.py b/Lib/test/test_tcl.py index 0b05c54..5d16968 100644 --- a/Lib/test/test_tcl.py +++ b/Lib/test/test_tcl.py @@ -144,23 +144,20 @@ class TclTest(unittest.TestCase): import sys if sys.platform.startswith(('win', 'darwin', 'cygwin')): return # no failure possible on windows? - if 'DISPLAY' in os.environ: - old_display = os.environ['DISPLAY'] - del os.environ['DISPLAY'] - # on some platforms, deleting environment variables - # doesn't actually carry through to the process level - # because they don't support unsetenv - # If that's the case, abort. - display = os.popen('echo $DISPLAY').read().strip() - if display: - return - try: + with test_support.EnvironmentVarGuard() as env: + if 'DISPLAY' in os.environ: + del env['DISPLAY'] + # on some platforms, deleting environment variables + # doesn't actually carry through to the process level + # because they don't support unsetenv + # If that's the case, abort. + display = os.popen('echo $DISPLAY').read().strip() + if display: + return + tcl = Tcl() self.assertRaises(TclError, tcl.winfo_geometry) self.assertRaises(TclError, tcl.loadtk) - finally: - if old_display is not None: - os.environ['DISPLAY'] = old_display def test_main(): test_support.run_unittest(TclTest, TkinterTest) diff --git a/Lib/test/test_tempfile.py b/Lib/test/test_tempfile.py index 2b46c46..8b883b0 100644 --- a/Lib/test/test_tempfile.py +++ b/Lib/test/test_tempfile.py @@ -153,7 +153,7 @@ class test__candidate_tempdir_list(TC): for envname in 'TMPDIR', 'TEMP', 'TMP': dirname = os.getenv(envname) if not dirname: - env.set(envname, os.path.abspath(envname)) + env[envname] = os.path.abspath(envname) cand = tempfile._candidate_tempdir_list() diff --git a/Lib/test/test_urllib.py b/Lib/test/test_urllib.py index 8b687cf7..ffc2f31 100644 --- a/Lib/test/test_urllib.py +++ b/Lib/test/test_urllib.py @@ -103,7 +103,7 @@ class ProxyTests(unittest.TestCase): # Delete all proxy related env vars for k, v in os.environ.iteritems(): if 'proxy' in k.lower(): - env.unset(k) + del env[k] def tearDown(self): # Restore all proxy related env vars diff --git a/Lib/test/test_xmlrpc.py b/Lib/test/test_xmlrpc.py index d38b74a..8d527c3 100644 --- a/Lib/test/test_xmlrpc.py +++ b/Lib/test/test_xmlrpc.py @@ -605,7 +605,7 @@ class CGIHandlerTestCase(unittest.TestCase): def test_cgi_get(self): with test_support.EnvironmentVarGuard() as env: - env.set('REQUEST_METHOD', 'GET') + env['REQUEST_METHOD'] = 'GET' # if the method is GET and no request_text is given, it runs handle_get # get sysout output tmp = sys.stdout @@ -646,7 +646,7 @@ class CGIHandlerTestCase(unittest.TestCase): sys.stdout = open(test_support.TESTFN, "w") with test_support.EnvironmentVarGuard() as env: - env.set('CONTENT_LENGTH', str(len(data))) + env['CONTENT_LENGTH'] = str(len(data)) self.cgi.handle_request() sys.stdin.close() -- cgit v0.12