Add EnvironmentVarGuard for test_urllib, test_urllib2.py, urllib2_localnet.py

Author: WolframAlph

Lib/test/test_urllib.py

@@ -258,34 +258,31 @@ class ProxyTests_withOrderedEnv(unittest.TestCase):
 
     def setUp(self):
         # We need to test conditions, where variable order _is_ significant
-        self._saved_env = os.environ
-        # Monkey patch os.environ, start with empty fake environment
-        os.environ = collections.OrderedDict()
-
-    def tearDown(self):
-        os.environ = self._saved_env
+        self.env = self.enterContext(os_helper.EnvironmentVarGuard())
+        # Start with empty fake environment
+        self.env.clear()
 
     def test_getproxies_environment_prefer_lowercase(self):
         # Test lowercase preference with removal
-        os.environ['no_proxy'] = ''
-        os.environ['No_Proxy'] = 'localhost'
+        self.env['no_proxy'] = ''
+        self.env['No_Proxy'] = 'localhost'
         self.assertFalse(urllib.request.proxy_bypass_environment('localhost'))
         self.assertFalse(urllib.request.proxy_bypass_environment('arbitrary'))
-        os.environ['http_proxy'] = ''
-        os.environ['HTTP_PROXY'] = 'http://somewhere:3128'
+        self.env['http_proxy'] = ''
+        self.env['HTTP_PROXY'] = 'http://somewhere:3128'
         proxies = urllib.request.getproxies_environment()
         self.assertEqual({}, proxies)
         # Test lowercase preference of proxy bypass and correct matching including ports
-        os.environ['no_proxy'] = 'localhost, noproxy.com, my.proxy:1234'
-        os.environ['No_Proxy'] = 'xyz.com'
+        self.env['no_proxy'] = 'localhost, noproxy.com, my.proxy:1234'
+        self.env['No_Proxy'] = 'xyz.com'
         self.assertTrue(urllib.request.proxy_bypass_environment('localhost'))
         self.assertTrue(urllib.request.proxy_bypass_environment('noproxy.com:5678'))
         self.assertTrue(urllib.request.proxy_bypass_environment('my.proxy:1234'))
         self.assertFalse(urllib.request.proxy_bypass_environment('my.proxy'))
         self.assertFalse(urllib.request.proxy_bypass_environment('arbitrary'))
         # Test lowercase preference with replacement
-        os.environ['http_proxy'] = 'http://somewhere:3128'
-        os.environ['Http_Proxy'] = 'http://somewhereelse:3128'
+        self.env['http_proxy'] = 'http://somewhere:3128'
+        self.env['Http_Proxy'] = 'http://somewhereelse:3128'
         proxies = urllib.request.getproxies_environment()
         self.assertEqual('http://somewhere:3128', proxies['http'])
 

Lib/test/test_urllib2.py

@@ -1444,30 +1444,32 @@ def test_proxy(self):
                              [tup[0:2] for tup in o.calls])
 
     def test_proxy_no_proxy(self):
-        os.environ['no_proxy'] = 'python.org'
-        o = OpenerDirector()
-        ph = urllib.request.ProxyHandler(dict(http="proxy.example.com"))
-        o.add_handler(ph)
-        req = Request("http://www.perl.org/")
-        self.assertEqual(req.host, "www.perl.org")
-        o.open(req)
-        self.assertEqual(req.host, "proxy.example.com")
-        req = Request("http://www.python.org")
-        self.assertEqual(req.host, "www.python.org")
-        o.open(req)
-        self.assertEqual(req.host, "www.python.org")
-        del os.environ['no_proxy']
+        with os_helper.EnvironmentVarGuard() as env:
+            env['no_proxy'] = 'python.org'
+            o = OpenerDirector()
+            ph = urllib.request.ProxyHandler(dict(http="proxy.example.com"))
+            o.add_handler(ph)
+            req = Request("http://www.perl.org/")
+            self.assertEqual(req.host, "www.perl.org")
+            o.open(req)
+            self.assertEqual(req.host, "proxy.example.com")
+            req = Request("http://www.python.org")
+            self.assertEqual(req.host, "www.python.org")
+            o.open(req)
+            self.assertEqual(req.host, "www.python.org")
+            del env['no_proxy']
 
     def test_proxy_no_proxy_all(self):
-        os.environ['no_proxy'] = '*'
-        o = OpenerDirector()
-        ph = urllib.request.ProxyHandler(dict(http="proxy.example.com"))
-        o.add_handler(ph)
-        req = Request("http://www.python.org")
-        self.assertEqual(req.host, "www.python.org")
-        o.open(req)
-        self.assertEqual(req.host, "www.python.org")
-        del os.environ['no_proxy']
+        with os_helper.EnvironmentVarGuard() as env:
+            env['no_proxy'] = '*'
+            o = OpenerDirector()
+            ph = urllib.request.ProxyHandler(dict(http="proxy.example.com"))
+            o.add_handler(ph)
+            req = Request("http://www.python.org")
+            self.assertEqual(req.host, "www.python.org")
+            o.open(req)
+            self.assertEqual(req.host, "www.python.org")
+            del env['no_proxy']
 
     def test_proxy_https(self):
         o = OpenerDirector()

Lib/test/test_urllib2_localnet.py

@@ -11,6 +11,7 @@
 from test import support
 from test.support import hashlib_helper
 from test.support import threading_helper
+from test.support import os_helper
 
 try:
     import ssl
@@ -330,12 +331,9 @@ class ProxyAuthTests(unittest.TestCase):
     def setUp(self):
         super(ProxyAuthTests, self).setUp()
         # Ignore proxy bypass settings in the environment.
-        def restore_environ(old_environ):
-            os.environ.clear()
-            os.environ.update(old_environ)
-        self.addCleanup(restore_environ, os.environ.copy())
-        os.environ['NO_PROXY'] = ''
-        os.environ['no_proxy'] = ''
+        env = self.enterContext(os_helper.EnvironmentVarGuard())
+        env['NO_PROXY'] = ''
+        env['no_proxy'] = ''
 
         self.digest_auth_handler = DigestAuthHandler()
         self.digest_auth_handler.set_users({self.USER: self.PASSWD})
@@ -456,12 +454,9 @@ def setUp(self):
         self.addCleanup(urllib.request.urlcleanup)
 
         # Ignore proxies for localhost tests.
-        def restore_environ(old_environ):
-            os.environ.clear()
-            os.environ.update(old_environ)
-        self.addCleanup(restore_environ, os.environ.copy())
-        os.environ['NO_PROXY'] = '*'
-        os.environ['no_proxy'] = '*'
+        env = self.enterContext(os_helper.EnvironmentVarGuard())
+        env['NO_PROXY'] = '*'
+        env['no_proxy'] = '*'
 
     def urlopen(self, url, data=None, **kwargs):
         l = []