cpython/Lib/test/test_profiling/test_sampling_profiler/helpers.py

"""Helper utilities for sampling profiler tests."""

import contextlib
import socket
import subprocess
import sys
import unittest
from collections import namedtuple

from test.support import SHORT_TIMEOUT
from test.support.socket_helper import find_unused_port
from test.support.os_helper import unlink


PROCESS_VM_READV_SUPPORTED = False

try:
    from _remote_debugging import PROCESS_VM_READV_SUPPORTED  # noqa: F401
    import _remote_debugging  # noqa: F401
except ImportError:
    raise unittest.SkipTest(
        "Test only runs when _remote_debugging is available"
    )
else:
    import profiling.sampling  # noqa: F401
    from profiling.sampling.sample import SampleProfiler  # noqa: F401


skip_if_not_supported = unittest.skipIf(
    (
        sys.platform != "darwin"
        and sys.platform != "linux"
        and sys.platform != "win32"
    ),
    "Test only runs on Linux, Windows and MacOS",
)

SubprocessInfo = namedtuple("SubprocessInfo", ["process", "socket"])


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

    Args:
        script: Python code to execute in the subprocess

    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()
        response = client_socket.recv(1024)
        if response != b"ready":
            raise RuntimeError(
                f"Unexpected response from subprocess: {response!r}"
            )

        yield SubprocessInfo(proc, client_socket)
    finally:
        if client_socket is not None:
            client_socket.close()
        if proc.poll() is None:
            proc.kill()
        proc.wait()


def close_and_unlink(file):
    """Close a file and unlink it from the filesystem."""
    file.close()
    unlink(file.name)