mirror of
https://github.com/python/cpython.git
synced 2025-12-08 06:10:17 +00:00
101 lines
2.8 KiB
Python
101 lines
2.8 KiB
Python
"""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)
|