Function test_subprocess

Lib/test/test_profiling/test_sampling_profiler/helpers.py:115–170

Context manager to create a test subprocess with socket synchronization. Args: script: Python code to execute in the subprocess. If wait_for_working is True, script should send b"working" after starting work. wait_for_working: If True, wait for both "ready" and "

(script, wait_for_working=False)

@contextlib.contextmanager
def test_subprocess(script, wait_for_working=False):
    """Context manager to create a test subprocess with socket synchronization.

    Args:
        script: Python code to execute in the subprocess. If wait_for_working
                is True, script should send b"working" after starting work.
        wait_for_working: If True, wait for both "ready" and "working" signals.
                         Default False for backward compatibility.

    Yields:
        SubprocessInfo: Named tuple with process and socket objects
    """
    # Find an unused port for socket communication
    port = find_unused_port()

    # Inject socket connection code at the beginning of the script
    socket_code = f"""
import socket
_test_sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
_test_sock.connect(('localhost', {port}))
_test_sock.sendall(b"ready")
"""

    # Combine socket code with user script
    full_script = socket_code + script

    # Create server socket to wait for process to be ready
    server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    server_socket.bind(("localhost", port))
    server_socket.settimeout(SHORT_TIMEOUT)
    server_socket.listen(1)

    proc = subprocess.Popen(
        [sys.executable, "-c", full_script],
        stdout=subprocess.DEVNULL,
        stderr=subprocess.DEVNULL,
    )

    client_socket = None
    try:
        # Wait for process to connect and send ready signal
        client_socket, _ = server_socket.accept()
        server_socket.close()
        server_socket = None

        # Wait for ready signal, and optionally working signal
        if wait_for_working:
            _wait_for_signal(client_socket, [b"ready", b"working"])
        else:
            _wait_for_signal(client_socket, b"ready")

        yield SubprocessInfo(proc, client_socket)
    finally:
        _cleanup_sockets(client_socket, server_socket)
        _cleanup_process(proc)