Hide navigation sidebar
Hide table of contents sidebar
Skip to content
Back to top
Toggle table of contents sidebar

# Source code for pytestshellutils.utils.ports

    # Copyright 2021-2024 VMware, Inc.
    # SPDX-License-Identifier: Apache-2.0
    """
    Ports related utility functions.
    """
    import contextlib
    import logging
    from typing import Iterable
    from typing import Set

    import pytest

    import pytestshellutils.utils.socket as socket

    log = logging.getLogger(__name__)

    def get_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.
        """
        if not isinstance(use_cache, bool):  # pragma: no cover
            raise pytest.UsageError(
                f"The value of 'use_cache' needs to be an boolean, not {type(use_cache)}"
            )

        with contextlib.closing(socket.socket(family=socket.AF_INET, type=socket.SOCK_STREAM)) as usock:
            usock.bind(("127.0.0.1", 0))
            port: int = usock.getsockname()[1]

        if use_cache:
            try:
                cached_ports = get_unused_localhost_port.__cached_ports__  # type: ignore[attr-defined]
            except AttributeError:
                cached_ports = get_unused_localhost_port.__cached_ports__ = set()  # type: ignore[attr-defined]
            if port in cached_ports:
                return get_unused_localhost_port(use_cache=use_cache)
            cached_ports.add(port)

        return port

    def get_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)

        for port in set(check_ports):
            with contextlib.closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as sock:
                conn = sock.connect_ex(("localhost", port))
                try:
                    if conn == 0:
                        log.debug("Port %s is connectable!", port)
                        connectable_ports.add(port)
                        sock.shutdown(socket.SHUT_RDWR)
                except OSError:  # pragma: no cover
                    continue
        return connectable_ports