# Copyright 2021-2024 VMware, Inc.# SPDX-License-Identifier: Apache-2.0"""Ports related utility functions."""importcontextlibimportloggingfromtypingimportIterablefromtypingimportSetimportpytestimportpytestshellutils.utils.socketassocketlog=logging.getLogger(__name__)
[docs]defget_unused_localhost_port(use_cache:bool=False)->int:""" Return a random unused port on localhost. Keyword Arguments: use_cache: If ``use_cache`` is ``True``, consecutive calls to this function will never return the cached port. """ifnotisinstance(use_cache,bool):# pragma: no coverraisepytest.UsageError(f"The value of 'use_cache' needs to be an boolean, not {type(use_cache)}")withcontextlib.closing(socket.socket(family=socket.AF_INET,type=socket.SOCK_STREAM))asusock:usock.bind(("127.0.0.1",0))port:int=usock.getsockname()[1]ifuse_cache:try:cached_ports=get_unused_localhost_port.__cached_ports__# type: ignore[attr-defined]exceptAttributeError:cached_ports=get_unused_localhost_port.__cached_ports__=set()# type: ignore[attr-defined]ifportincached_ports:returnget_unused_localhost_port(use_cache=use_cache)cached_ports.add(port)returnport
[docs]defget_connectable_ports(ports:Iterable[int])->Set[int]:""" Given a list of ports, returns those that we can connect to. Arguments: ports: An iterable of ports to try and connect to Returns: set: Returns a set of the ports where connection was successful """connectable_ports=set()check_ports=set(ports)forportinset(check_ports):withcontextlib.closing(socket.socket(socket.AF_INET,socket.SOCK_STREAM))assock:conn=sock.connect_ex(("localhost",port))try:ifconn==0:log.debug("Port %s is connectable!",port)connectable_ports.add(port)sock.shutdown(socket.SHUT_RDWR)exceptOSError:# pragma: no covercontinuereturnconnectable_ports