diff --git a/tests/profiling/collector/pprof_utils.py b/tests/profiling/collector/pprof_utils.py index 0c93538afb3..463a55172f6 100644 --- a/tests/profiling/collector/pprof_utils.py +++ b/tests/profiling/collector/pprof_utils.py @@ -195,7 +195,7 @@ def get_function_with_id(profile: pprof_pb2.Profile, function_id: int) -> pprof_ def assert_lock_events_of_type( profile: pprof_pb2.Profile, - expected_events: List[LockEvent], + expected_events: Sequence[LockEvent], event_type: LockEventType, ): samples = get_samples_with_value_type( @@ -227,8 +227,8 @@ def assert_lock_events_of_type( def assert_lock_events( profile: pprof_pb2.Profile, - expected_acquire_events: Union[List[LockEvent], None] = None, - expected_release_events: Union[List[LockEvent], None] = None, + expected_acquire_events: Union[Sequence[LockEvent], None] = None, + expected_release_events: Union[Sequence[LockEvent], None] = None, ): if expected_acquire_events: assert_lock_events_of_type(profile, expected_acquire_events, LockEventType.ACQUIRE) diff --git a/tests/profiling/test_main.py b/tests/profiling/test_main.py index 17923679950..345d01fd08e 100644 --- a/tests/profiling/test_main.py +++ b/tests/profiling/test_main.py @@ -1,6 +1,7 @@ # -*- encoding: utf-8 -*- import multiprocessing import os +from pathlib import Path import sys import pytest @@ -10,7 +11,7 @@ from tests.utils import call_program -def test_call_script(): +def test_call_script() -> None: env = os.environ.copy() env["DD_PROFILING_ENABLED"] = "1" stdout, stderr, exitcode, _ = call_program( @@ -20,8 +21,10 @@ def test_call_script(): assert exitcode == 0, (stdout, stderr) else: assert exitcode == 42, (stdout, stderr) - hello, pid = list(s.strip() for s in stdout.decode().strip().split("\n")) - assert hello == "hello world", stdout.decode().strip() + + stdout = stdout.decode() if isinstance(stdout, bytes) else stdout + hello, _ = list(s.strip() for s in stdout.strip().split("\n")) + assert hello == "hello world", stdout.strip() @pytest.mark.skipif(not os.getenv("DD_PROFILE_TEST_GEVENT", False), reason="Not testing gevent") @@ -34,7 +37,7 @@ def test_call_script_gevent(): assert exitcode == 0, (stdout, stderr) -def test_call_script_pprof_output(tmp_path): +def test_call_script_pprof_output(tmp_path: Path) -> None: """This checks if the pprof output and atexit register work correctly. The script does not run for one minute, so if the `stop_on_exit` flag is broken, this test will fail. @@ -54,24 +57,27 @@ def test_call_script_pprof_output(tmp_path): assert exitcode == 0, (stdout, stderr) else: assert exitcode == 42, (stdout, stderr) - _, pid = list(s.strip() for s in stdout.decode().strip().split("\n")) - profile = pprof_utils.parse_newest_profile(filename + "." + str(pid)) + + stdout = stdout.decode() if isinstance(stdout, bytes) else stdout + _, pid = list(s.strip() for s in stdout.strip().split("\n")) + profile = pprof_utils.parse_newest_profile(f"{filename}.{pid}") samples = pprof_utils.get_samples_with_value_type(profile, "cpu-time") assert len(samples) > 0 @pytest.mark.skipif(sys.platform == "win32", reason="fork only available on Unix") -def test_fork(tmp_path): +def test_fork(tmp_path: Path) -> None: filename = str(tmp_path / "pprof") env = os.environ.copy() env["DD_PROFILING_OUTPUT_PPROF"] = filename env["DD_PROFILING_CAPTURE_PCT"] = "100" - stdout, stderr, exitcode, pid = call_program( + stdout, _, exitcode, pid = call_program( "python", os.path.join(os.path.dirname(__file__), "simple_program_fork.py"), env=env ) assert exitcode == 0 - child_pid = stdout.decode().strip() - profile = pprof_utils.parse_newest_profile(filename + "." + str(pid)) + stdout = stdout.decode() if isinstance(stdout, bytes) else stdout + child_pid = stdout.strip() + profile = pprof_utils.parse_newest_profile(f"{filename}.{pid}") parent_expected_acquire_events = [ pprof_utils.LockAcquireEvent( caller_name="", @@ -134,11 +140,9 @@ def test_fork(tmp_path): @pytest.mark.skipif(sys.platform == "win32", reason="fork only available on Unix") @pytest.mark.skipif(not os.getenv("DD_PROFILE_TEST_GEVENT", False), reason="Not testing gevent") -def test_fork_gevent(): +def test_fork_gevent() -> None: env = os.environ.copy() - stdout, stderr, exitcode, pid = call_program( - "python", os.path.join(os.path.dirname(__file__), "gevent_fork.py"), env=env - ) + _, _, exitcode, _ = call_program("python", os.path.join(os.path.dirname(__file__), "gevent_fork.py"), env=env) assert exitcode == 0 @@ -149,7 +153,7 @@ def test_fork_gevent(): "method", set(methods) - {"forkserver", "fork"}, ) -def test_multiprocessing(method, tmp_path): +def test_multiprocessing(method: str, tmp_path: Path) -> None: filename = str(tmp_path / "pprof") env = os.environ.copy() env["DD_PROFILING_OUTPUT_PPROF"] = filename @@ -163,8 +167,9 @@ def test_multiprocessing(method, tmp_path): env=env, ) assert exitcode == 0, (stdout, stderr) - pid, child_pid = list(s.strip() for s in stdout.decode().strip().split("\n")) - profile = pprof_utils.parse_newest_profile(filename + "." + str(pid)) + stdout = stdout.decode() if isinstance(stdout, bytes) else stdout + pid, child_pid = list(s.strip() for s in stdout.strip().split("\n")) + profile = pprof_utils.parse_newest_profile(f"{filename}.{pid}") samples = pprof_utils.get_samples_with_value_type(profile, "cpu-time") assert len(samples) > 0 child_profile = pprof_utils.parse_newest_profile(filename + "." + str(child_pid)) @@ -177,7 +182,7 @@ def test_multiprocessing(method, tmp_path): env=dict(DD_PROFILING_ENABLED="1"), err=lambda _: "RuntimeError: the memalloc module is already started" not in _, ) -def test_memalloc_no_init_error_on_fork(): +def test_memalloc_no_init_error_on_fork() -> None: import os pid = os.fork() @@ -195,7 +200,7 @@ def test_memalloc_no_init_error_on_fork(): out="OK\n", err=None, ) -def test_profiler_start_up_with_module_clean_up_in_protobuf_app(): +def test_profiler_start_up_with_module_clean_up_in_protobuf_app() -> None: # This can cause segfaults if we do module clean up with later versions of # protobuf. This is a regression test. from google.protobuf import empty_pb2 # noqa:F401