src/python/grpcio_tests/tests/unit/framework/common/__init__.py

   1 # Copyright 2019 The gRPC authors.
   2 #
   3 # Licensed under the Apache License, Version 2.0 (the "License");
   4 # you may not use this file except in compliance with the License.
   5 # You may obtain a copy of the License at
   6 #
   7 #     http://www.apache.org/licenses/LICENSE-2.0
   8 #
   9 # Unless required by applicable law or agreed to in writing, software
  10 # distributed under the License is distributed on an "AS IS" BASIS,
  11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12 # See the License for the specific language governing permissions and
  13 # limitations under the License.
  14
  15 import contextlib
  16 import os
  17 import socket
  18 import errno
  19
  20 _DEFAULT_SOCK_OPTIONS = (socket.SO_REUSEADDR,
  21                          socket.SO_REUSEPORT) if os.name != 'nt' else (
  22                              socket.SO_REUSEADDR,)
  23 _UNRECOVERABLE_ERRNOS = (errno.EADDRINUSE, errno.ENOSR)
  24
  25
  26 def get_socket(bind_address='localhost',
  27                port=0,
  28                listen=True,
  29                sock_options=_DEFAULT_SOCK_OPTIONS):
  30     """Opens a socket.
  31
  32     Useful for reserving a port for a system-under-test.
  33
  34     Args:
  35       bind_address: The host to which to bind.
  36       port: The port to which to bind.
  37       listen: A boolean value indicating whether or not to listen on the socket.
  38       sock_options: A sequence of socket options to apply to the socket.
  39
  40     Returns:
  41       A tuple containing:
  42         - the address to which the socket is bound
  43         - the port to which the socket is bound
  44         - the socket object itself
  45     """
  46     _sock_options = sock_options if sock_options else []
  47     if socket.has_ipv6:
  48         address_families = (socket.AF_INET6, socket.AF_INET)
  49     else:
  50         address_families = (socket.AF_INET)
  51     for address_family in address_families:
  52         try:
  53             sock = socket.socket(address_family, socket.SOCK_STREAM)
  54             for sock_option in _sock_options:
  55                 sock.setsockopt(socket.SOL_SOCKET, sock_option, 1)
  56             sock.bind((bind_address, port))
  57             if listen:
  58                 sock.listen(1)
  59             return bind_address, sock.getsockname()[1], sock
  60         except OSError as os_error:
  61             sock.close()
  62             if os_error.errno in _UNRECOVERABLE_ERRNOS:
  63                 raise
  64             else:
  65                 continue
  66         # For PY2, socket.error is a child class of IOError; for PY3, it is
  67         # pointing to OSError. We need this catch to make it 2/3 agnostic.
  68         except socket.error:  # pylint: disable=duplicate-except
  69             sock.close()
  70             continue
  71     raise RuntimeError("Failed to bind to {} with sock_options {}".format(
  72         bind_address, sock_options))
  73
  74
  75 @contextlib.contextmanager
  76 def bound_socket(bind_address='localhost',
  77                  port=0,
  78                  listen=True,
  79                  sock_options=_DEFAULT_SOCK_OPTIONS):
  80     """Opens a socket bound to an arbitrary port.
  81
  82     Useful for reserving a port for a system-under-test.
  83
  84     Args:
  85       bind_address: The host to which to bind.
  86       port: The port to which to bind.
  87       listen: A boolean value indicating whether or not to listen on the socket.
  88       sock_options: A sequence of socket options to apply to the socket.
  89
  90     Yields:
  91       A tuple containing:
  92         - the address to which the socket is bound
  93         - the port to which the socket is bound
  94     """
  95     host, port, sock = get_socket(bind_address=bind_address,
  96                                   port=port,
  97                                   listen=listen,
  98                                   sock_options=sock_options)
  99     try:
 100         yield host, port
 101     finally:
 102         sock.close()