Issue #12138: fix use of transient_internet() in test_urllibnet

1 files changed, 62 insertions, 77 deletions

diff --git a/Lib/test/test_urllibnet.py b/Lib/test/test_urllibnet.py
index 32efb2b..03d1708 100644
--- a/Lib/test/test_urllibnet.py
+++ b/Lib/test/test_urllibnet.py
@@ -3,6 +3,7 @@
 import unittest
 from test import support
 
+import contextlib
 import socket
 import urllib.request
 import sys
@@ -27,6 +28,7 @@ class URLTimeoutTest(unittest.TestCase):
             f = urllib.request.urlopen("http://www.python.org/")
             x = f.read()
 
+
 class urlopenNetworkTests(unittest.TestCase):
     """Tests urllib.reqest.urlopen using the network.
 
@@ -42,43 +44,37 @@ class urlopenNetworkTests(unittest.TestCase):
 
     """
 
+    @contextlib.contextmanager
     def urlopen(self, *args, **kwargs):
         resource = args[0]
-        cm = support.transient_internet(resource)
-        cm.__enter__()
-        self.addCleanup(cm.__exit__, None, None, None)
-        return urllib.request.urlopen(*args, **kwargs)
+        with support.transient_internet(resource):
+            r = urllib.request.urlopen(*args, **kwargs)
+            try:
+                yield r
+            finally:
+                r.close()
 
     def test_basic(self):
         # Simple test expected to pass.
-        open_url = self.urlopen("http://www.python.org/")
-        for attr in ("read", "readline", "readlines", "fileno", "close",
-                     "info", "geturl"):
-            self.assertTrue(hasattr(open_url, attr), "object returned from "
-                            "urlopen lacks the %s attribute" % attr)
-        try:
+        with self.urlopen("http://www.python.org/") as open_url:
+            for attr in ("read", "readline", "readlines", "fileno", "close",
+                         "info", "geturl"):
+                self.assertTrue(hasattr(open_url, attr), "object returned from "
+                                "urlopen lacks the %s attribute" % attr)
             self.assertTrue(open_url.read(), "calling 'read' failed")
-        finally:
-            open_url.close()
 
     def test_readlines(self):
         # Test both readline and readlines.
-        open_url = self.urlopen("http://www.python.org/")
-        try:
+        with self.urlopen("http://www.python.org/") as open_url:
             self.assertIsInstance(open_url.readline(), bytes,
                                   "readline did not return a string")
             self.assertIsInstance(open_url.readlines(), list,
                                   "readlines did not return a list")
-        finally:
-            open_url.close()
 
     def test_info(self):
         # Test 'info'.
-        open_url = self.urlopen("http://www.python.org/")
-        try:
+        with self.urlopen("http://www.python.org/") as open_url:
             info_obj = open_url.info()
-        finally:
-            open_url.close()
             self.assertIsInstance(info_obj, email.message.Message,
                                   "object returned by 'info' is not an "
                                   "instance of email.message.Message")
@@ -87,22 +83,20 @@ class urlopenNetworkTests(unittest.TestCase):
     def test_geturl(self):
         # Make sure same URL as opened is returned by geturl.
         URL = "http://www.python.org/"
-        open_url = self.urlopen(URL)
-        try:
+        with self.urlopen(URL) as open_url:
             gotten_url = open_url.geturl()
-        finally:
-            open_url.close()
-        self.assertEqual(gotten_url, URL)
+            self.assertEqual(gotten_url, URL)
 
     def test_getcode(self):
         # test getcode() with the fancy opener to get 404 error codes
         URL = "http://www.python.org/XXXinvalidXXX"
-        open_url = urllib.request.FancyURLopener().open(URL)
-        try:
-            code = open_url.getcode()
-        finally:
-            open_url.close()
-        self.assertEqual(code, 404)
+        with support.transient_internet(URL):
+            open_url = urllib.request.FancyURLopener().open(URL)
+            try:
+                code = open_url.getcode()
+            finally:
+                open_url.close()
+            self.assertEqual(code, 404)
 
     def test_fileno(self):
         if sys.platform in ('win32',):
@@ -110,14 +104,11 @@ class urlopenNetworkTests(unittest.TestCase):
             # test can't pass on Windows.
             return
         # Make sure fd returned by fileno is valid.
-        open_url = self.urlopen("http://www.python.org/", timeout=None)
-        fd = open_url.fileno()
-        FILE = os.fdopen(fd, encoding='utf-8')
-        try:
-            self.assertTrue(FILE.read(), "reading from file created using fd "
-                                      "returned by fileno failed")
-        finally:
-            FILE.close()
+        with self.urlopen("http://www.python.org/", timeout=None) as open_url:
+            fd = open_url.fileno()
+            with os.fdopen(fd, encoding='utf-8') as f:
+                self.assertTrue(f.read(), "reading from file created using fd "
+                                          "returned by fileno failed")
 
     def test_bad_address(self):
         # Make sure proper exception is raised when connecting to a bogus
@@ -133,66 +124,60 @@ class urlopenNetworkTests(unittest.TestCase):
                           urllib.request.urlopen,
                           "http://sadflkjsasf.i.nvali.d/")
 
+
 class urlretrieveNetworkTests(unittest.TestCase):
     """Tests urllib.request.urlretrieve using the network."""
 
+    @contextlib.contextmanager
     def urlretrieve(self, *args):
         resource = args[0]
-        cm = support.transient_internet(resource)
-        cm.__enter__()
-        self.addCleanup(cm.__exit__, None, None, None)
-        return urllib.request.urlretrieve(*args)
+        with support.transient_internet(resource):
+            file_location, info = urllib.request.urlretrieve(*args)
+            try:
+                yield file_location, info
+            finally:
+                support.unlink(file_location)
 
     def test_basic(self):
         # Test basic functionality.
-        file_location,info = self.urlretrieve("http://www.python.org/")
-        self.assertTrue(os.path.exists(file_location), "file location returned by"
-                        " urlretrieve is not a valid path")
-        FILE = open(file_location, encoding='utf-8')
-        try:
-            self.assertTrue(FILE.read(), "reading from the file location returned"
-                         " by urlretrieve failed")
-        finally:
-            FILE.close()
-            os.unlink(file_location)
+        with self.urlretrieve("http://www.python.org/") as (file_location, info):
+            self.assertTrue(os.path.exists(file_location), "file location returned by"
+                            " urlretrieve is not a valid path")
+            with open(file_location, encoding='utf-8') as f:
+                self.assertTrue(f.read(), "reading from the file location returned"
+                                " by urlretrieve failed")
 
     def test_specified_path(self):
         # Make sure that specifying the location of the file to write to works.
-        file_location,info = self.urlretrieve("http://www.python.org/",
-                                              support.TESTFN)
-        self.assertEqual(file_location, support.TESTFN)
-        self.assertTrue(os.path.exists(file_location))
-        FILE = open(file_location, encoding='utf-8')
-        try:
-            self.assertTrue(FILE.read(), "reading from temporary file failed")
-        finally:
-            FILE.close()
-            os.unlink(file_location)
+        with self.urlretrieve("http://www.python.org/",
+                              support.TESTFN) as (file_location, info):
+            self.assertEqual(file_location, support.TESTFN)
+            self.assertTrue(os.path.exists(file_location))
+            with open(file_location, encoding='utf-8') as f:
+                self.assertTrue(f.read(), "reading from temporary file failed")
 
     def test_header(self):
         # Make sure header returned as 2nd value from urlretrieve is good.
-        file_location, header = self.urlretrieve("http://www.python.org/")
-        os.unlink(file_location)
-        self.assertIsInstance(header, email.message.Message,
-                              "header is not an instance of email.message.Message")
+        with self.urlretrieve("http://www.python.org/") as (file_location, info):
+            self.assertIsInstance(info, email.message.Message,
+                                  "info is not an instance of email.message.Message")
 
     def test_data_header(self):
         logo = "http://www.python.org/community/logos/python-logo-master-v3-TM.png"
-        file_location, fileheaders = self.urlretrieve(logo)
-        os.unlink(file_location)
-        datevalue = fileheaders.get('Date')
-        dateformat = '%a, %d %b %Y %H:%M:%S GMT'
-        try:
-            time.strptime(datevalue, dateformat)
-        except ValueError:
-            self.fail('Date value not in %r format', dateformat)
+        with self.urlretrieve(logo) as (file_location, fileheaders):
+            datevalue = fileheaders.get('Date')
+            dateformat = '%a, %d %b %Y %H:%M:%S GMT'
+            try:
+                time.strptime(datevalue, dateformat)
+            except ValueError:
+                self.fail('Date value not in %r format', dateformat)
 
 
 def test_main():
     support.requires('network')
     support.run_unittest(URLTimeoutTest,
-                              urlopenNetworkTests,
-                              urlretrieveNetworkTests)
+                         urlopenNetworkTests,
+                         urlretrieveNetworkTests)
 
 if __name__ == "__main__":
     test_main()