Skip to content

Commit 3fc3fac

Browse files
chore(profiling): improve typing in tests (#15249)
## Description As title says!
1 parent 278abf0 commit 3fc3fac

File tree

5 files changed

+21
-23
lines changed

5 files changed

+21
-23
lines changed

tests/profiling/test_main.py

Lines changed: 9 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22
import multiprocessing
33
import os
44
import sys
5+
from typing import cast
56

67
import pytest
78

@@ -21,9 +22,9 @@ def test_call_script(monkeypatch):
2122
assert exitcode == 0, (stdout, stderr)
2223
else:
2324
assert exitcode == 42, (stdout, stderr)
24-
hello, interval, _ = list(s.strip() for s in stdout.decode().strip().split("\n"))
25-
assert hello == "hello world", stdout.decode().strip()
26-
assert float(interval) >= 0.01, stdout.decode().strip()
25+
hello, interval, _ = list(s.strip() for s in cast(bytes, stdout).decode().strip().split("\n"))
26+
assert hello == "hello world", cast(bytes, stdout).decode().strip()
27+
assert float(interval) >= 0.01, cast(bytes, stdout).decode().strip()
2728

2829

2930
@pytest.mark.skipif(not os.getenv("DD_PROFILE_TEST_GEVENT", False), reason="Not testing gevent")
@@ -51,7 +52,7 @@ def test_call_script_pprof_output(tmp_path, monkeypatch):
5152
assert exitcode == 0, (stdout, stderr)
5253
else:
5354
assert exitcode == 42, (stdout, stderr)
54-
hello, interval, pid = list(s.strip() for s in stdout.decode().strip().split("\n"))
55+
hello, interval, pid = list(s.strip() for s in cast(bytes, stdout).decode().strip().split("\n"))
5556
utils.check_pprof_file(filename + "." + str(pid))
5657

5758

@@ -61,11 +62,9 @@ def test_fork(tmp_path, monkeypatch):
6162
monkeypatch.setenv("DD_PROFILING_API_TIMEOUT_MS", "100")
6263
monkeypatch.setenv("DD_PROFILING_OUTPUT_PPROF", filename)
6364
monkeypatch.setenv("DD_PROFILING_CAPTURE_PCT", "100")
64-
stdout, stderr, exitcode, pid = call_program(
65-
"python", os.path.join(os.path.dirname(__file__), "simple_program_fork.py")
66-
)
65+
stdout, _, exitcode, pid = call_program("python", os.path.join(os.path.dirname(__file__), "simple_program_fork.py"))
6766
assert exitcode == 0
68-
child_pid = stdout.decode().strip()
67+
child_pid = cast(bytes, stdout).decode().strip()
6968
utils.check_pprof_file(filename + "." + str(pid))
7069
utils.check_pprof_file(filename + "." + str(child_pid), sample_type="lock-release")
7170

@@ -74,7 +73,7 @@ def test_fork(tmp_path, monkeypatch):
7473
@pytest.mark.skipif(not os.getenv("DD_PROFILE_TEST_GEVENT", False), reason="Not testing gevent")
7574
def test_fork_gevent(monkeypatch):
7675
monkeypatch.setenv("DD_PROFILING_API_TIMEOUT_MS", "100")
77-
stdout, stderr, exitcode, pid = call_program("python", os.path.join(os.path.dirname(__file__), "gevent_fork.py"))
76+
_, _, exitcode, _ = call_program("python", os.path.join(os.path.dirname(__file__), "gevent_fork.py"))
7877
assert exitcode == 0
7978

8079

@@ -97,7 +96,7 @@ def test_multiprocessing(method, tmp_path, monkeypatch):
9796
method,
9897
)
9998
assert exitcode == 0, (stdout, stderr)
100-
pid, child_pid = list(s.strip() for s in stdout.decode().strip().split("\n"))
99+
pid, child_pid = list(s.strip() for s in cast(bytes, stdout).decode().strip().split("\n"))
101100
utils.check_pprof_file(filename + "." + str(pid))
102101
utils.check_pprof_file(filename + "." + str(child_pid), sample_type="wall-time")
103102

tests/profiling/test_profiler.py

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
import logging
22
import time
3+
from unittest import mock
34

4-
import mock
55
import pytest
66

77
import ddtrace
@@ -105,7 +105,7 @@ def snapshot():
105105

106106
class TestProfiler(profiler._ProfilerInstance):
107107
def _build_default_exporters(self, *args, **kargs):
108-
return []
108+
return None
109109

110110
p = TestProfiler()
111111
err_collector = mock.MagicMock(wraps=ErrCollect())
@@ -140,7 +140,6 @@ def test_default_collectors():
140140

141141

142142
def test_profiler_serverless(monkeypatch):
143-
# type: (...) -> None
144143
monkeypatch.setenv("AWS_LAMBDA_FUNCTION_NAME", "foobar")
145144
p = profiler.Profiler()
146145
assert isinstance(p._scheduler, scheduler.ServerlessScheduler)

tests/profiling_v2/test_profiler.py

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,8 @@
11
import logging
22
import sys
33
import time
4+
from unittest import mock
45

5-
import mock
66
import pytest
77

88
import ddtrace
@@ -107,7 +107,7 @@ def snapshot():
107107

108108
class TestProfiler(profiler._ProfilerInstance):
109109
def _build_default_exporters(self, *args, **kargs):
110-
return []
110+
return None
111111

112112
p = TestProfiler()
113113
err_collector = mock.MagicMock(wraps=ErrCollect())
@@ -182,7 +182,7 @@ def test_libdd_failure_telemetry_logging():
182182
2) import ddtrace.profiling.auto
183183
"""
184184

185-
import mock
185+
from unittest import mock
186186

187187
with mock.patch.multiple(
188188
"ddtrace.internal.datadog.profiling.ddup",
@@ -206,7 +206,7 @@ def test_libdd_failure_telemetry_logging():
206206
err=None
207207
)
208208
def test_libdd_failure_telemetry_logging_with_auto():
209-
import mock
209+
from unittest import mock
210210

211211
with mock.patch.multiple(
212212
"ddtrace.internal.datadog.profiling.ddup",
@@ -233,7 +233,7 @@ def test_stack_v2_failure_telemetry_logging():
233233
# mimicking the behavior of ddtrace-run, where the config is imported to
234234
# determine if profiling/stack_v2 is enabled
235235

236-
import mock
236+
from unittest import mock
237237

238238
with mock.patch.multiple(
239239
"ddtrace.internal.datadog.profiling.stack_v2",
@@ -257,7 +257,7 @@ def test_stack_v2_failure_telemetry_logging():
257257
err=None,
258258
)
259259
def test_stack_v2_failure_telemetry_logging_with_auto():
260-
import mock
260+
from unittest import mock
261261

262262
with mock.patch.multiple(
263263
"ddtrace.internal.datadog.profiling.stack_v2",
@@ -305,7 +305,7 @@ def test_user_threads_have_native_id():
305305
for _ in range(10):
306306
try:
307307
# The TID should be higher than the PID, but not too high
308-
assert 0 < t.native_id - getpid() < 100, (t.native_id, getpid())
308+
assert 0 < t.native_id - getpid() < 100, (t.native_id, getpid()) # pyright: ignore[reportOptionalOperand]
309309
except AttributeError:
310310
# The native_id attribute is set by the thread so we might have to
311311
# wait a bit for it to be set.

tests/profiling_v2/test_pytorch.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@ def test_call_script_pytorch_gpu(tmp_path, monkeypatch):
1313
monkeypatch.setenv("DD_PROFILING_OUTPUT_PPROF", filename)
1414
monkeypatch.setenv("DD_PROFILING_ENABLED", "1")
1515
monkeypatch.setenv("DD_PROFILING_PYTORCH_ENABLED", "1")
16-
stdout, stderr, exitcode, pid = call_program(
16+
_, stderr, exitcode, _ = call_program(
1717
"ddtrace-run", sys.executable, os.path.join(os.path.dirname(__file__), "simple_program_pytorch_gpu.py")
1818
)
1919
assert exitcode == 0, f"Profiler exited with code {exitcode}. Stderr: {stderr}"

tests/profiling_v2/test_scheduler.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,15 +1,15 @@
11
# -*- encoding: utf-8 -*-
22
import logging
33
import time
4-
5-
import mock
4+
from unittest import mock
65

76
from ddtrace.profiling import scheduler
87

98

109
def test_thread_name():
1110
s = scheduler.Scheduler()
1211
s.start()
12+
assert s._worker is not None
1313
assert s._worker.name == "ddtrace.profiling.scheduler:Scheduler"
1414
s.stop()
1515

0 commit comments

Comments
 (0)