1 files changed, 61 insertions, 19 deletions
diff --git a/Lib/test/support.py b/Lib/test/support.py
index 11bb168..8ef87c2 100644
--- a/Lib/test/support.py
+++ b/Lib/test/support.py
@@ -17,19 +17,21 @@ import unittest
 import importlib
 import collections
 
-__all__ = ["Error", "TestFailed", "ResourceDenied", "import_module",
-           "verbose", "use_resources", "max_memuse", "record_original_stdout",
-           "get_original_stdout", "unload", "unlink", "rmtree", "forget",
-           "is_resource_enabled", "requires", "find_unused_port", "bind_port",
-           "fcmp", "is_jython", "TESTFN", "HOST", "FUZZ", "findfile", "verify",
-           "vereq", "sortdict", "check_syntax_error", "open_urlresource",
-           "check_warnings", "CleanImport", "EnvironmentVarGuard",
-           "TransientResource", "captured_output", "captured_stdout",
-           "time_out", "socket_peer_reset", "ioerror_peer_reset",
-           "run_with_locale",
-           "set_memlimit", "bigmemtest", "bigaddrspacetest", "BasicTestRunner",
-           "run_unittest", "run_doctest", "threading_setup", "threading_cleanup",
-           "reap_children", "cpython_only", "check_impl_detail", "get_attribute"]
+__all__ = [
+    "Error", "TestFailed", "ResourceDenied", "import_module",
+    "verbose", "use_resources", "max_memuse", "record_original_stdout",
+    "get_original_stdout", "unload", "unlink", "rmtree", "forget",
+    "is_resource_enabled", "requires", "find_unused_port", "bind_port",
+    "fcmp", "is_jython", "TESTFN", "HOST", "FUZZ", "findfile", "verify",
+    "vereq", "sortdict", "check_syntax_error", "open_urlresource",
+    "check_warnings", "CleanImport", "EnvironmentVarGuard",
+    "TransientResource", "captured_output", "captured_stdout",
+    "time_out", "socket_peer_reset", "ioerror_peer_reset",
+    "run_with_locale", "transient_internet",
+    "set_memlimit", "bigmemtest", "bigaddrspacetest", "BasicTestRunner",
+    "run_unittest", "run_doctest", "threading_setup", "threading_cleanup",
+    "reap_children", "cpython_only", "check_impl_detail", "get_attribute",
+    ]
 
 class Error(Exception):
     """Base class for regression test exceptions."""
@@ -604,23 +606,63 @@ class TransientResource(object):
             else:
                 raise ResourceDenied("an optional resource is not available")
 
-
 # Context managers that raise ResourceDenied when various issues
 # with the Internet connection manifest themselves as exceptions.
+# XXX deprecate these and use transient_internet() instead
 time_out = TransientResource(IOError, errno=errno.ETIMEDOUT)
 socket_peer_reset = TransientResource(socket.error, errno=errno.ECONNRESET)
 ioerror_peer_reset = TransientResource(IOError, errno=errno.ECONNRESET)
 
 
 @contextlib.contextmanager
-def transient_internet():
+def transient_internet(resource_name, *, timeout=30.0, errnos=()):
     """Return a context manager that raises ResourceDenied when various issues
     with the Internet connection manifest themselves as exceptions."""
-    time_out = TransientResource(IOError, errno=errno.ETIMEDOUT)
-    socket_peer_reset = TransientResource(socket.error, errno=errno.ECONNRESET)
-    ioerror_peer_reset = TransientResource(IOError, errno=errno.ECONNRESET)
-    with time_out, socket_peer_reset, ioerror_peer_reset:
+    default_errnos = [
+        ('ECONNREFUSED', 111),
+        ('ECONNRESET', 104),
+        ('ENETUNREACH', 101),
+        ('ETIMEDOUT', 110),
+    ]
+
+    denied = ResourceDenied("Resource '%s' is not available" % resource_name)
+    captured_errnos = errnos
+    if not captured_errnos:
+        captured_errnos = [getattr(errno, name, num)
+                           for (name, num) in default_errnos]
+
+    def filter_error(err):
+        if (isinstance(err, socket.timeout) or
+            getattr(err, 'errno', None) in captured_errnos):
+            if not verbose:
+                sys.stderr.write(denied.args[0] + "\n")
+            raise denied from err
+
+    old_timeout = socket.getdefaulttimeout()
+    try:
+        if timeout is not None:
+            socket.setdefaulttimeout(timeout)
         yield
+    except IOError as err:
+        # urllib can wrap original socket errors multiple times (!), we must
+        # unwrap to get at the original error.
+        while True:
+            a = err.args
+            if len(a) >= 1 and isinstance(a[0], IOError):
+                err = a[0]
+            # The error can also be wrapped as args[1]:
+            #    except socket.error as msg:
+            #        raise IOError('socket error', msg).with_traceback(sys.exc_info()[2])
+            elif len(a) >= 2 and isinstance(a[1], IOError):
+                err = a[1]
+            else:
+                break
+        filter_error(err)
+        raise
+    # XXX should we catch generic exceptions and look for their
+    # __cause__ or __context__?
+    finally:
+        socket.setdefaulttimeout(old_timeout)
 
 
 @contextlib.contextmanager