chore(profiling): improve typing in tests

Author: KowalskiThomas

tests/profiling/test_main.py

@@ -2,6 +2,7 @@
 import multiprocessing
 import os
 import sys
+from typing import cast
 
 import pytest
 
@@ -21,9 +22,9 @@ def test_call_script(monkeypatch):
         assert exitcode == 0, (stdout, stderr)
     else:
         assert exitcode == 42, (stdout, stderr)
-    hello, interval, _ = list(s.strip() for s in stdout.decode().strip().split("\n"))
-    assert hello == "hello world", stdout.decode().strip()
-    assert float(interval) >= 0.01, stdout.decode().strip()
+    hello, interval, _ = list(s.strip() for s in cast(bytes, stdout).decode().strip().split("\n"))
+    assert hello == "hello world", cast(bytes, stdout).decode().strip()
+    assert float(interval) >= 0.01, cast(bytes, stdout).decode().strip()
 
 
 @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):
         assert exitcode == 0, (stdout, stderr)
     else:
         assert exitcode == 42, (stdout, stderr)
-    hello, interval, pid = list(s.strip() for s in stdout.decode().strip().split("\n"))
+    hello, interval, pid = list(s.strip() for s in cast(bytes, stdout).decode().strip().split("\n"))
     utils.check_pprof_file(filename + "." + str(pid))
 
 
@@ -61,11 +62,9 @@ def test_fork(tmp_path, monkeypatch):
     monkeypatch.setenv("DD_PROFILING_API_TIMEOUT_MS", "100")
     monkeypatch.setenv("DD_PROFILING_OUTPUT_PPROF", filename)
     monkeypatch.setenv("DD_PROFILING_CAPTURE_PCT", "100")
-    stdout, stderr, exitcode, pid = call_program(
-        "python", os.path.join(os.path.dirname(__file__), "simple_program_fork.py")
-    )
+    stdout, _, exitcode, pid = call_program("python", os.path.join(os.path.dirname(__file__), "simple_program_fork.py"))
     assert exitcode == 0
-    child_pid = stdout.decode().strip()
+    child_pid = cast(bytes, stdout).decode().strip()
     utils.check_pprof_file(filename + "." + str(pid))
     utils.check_pprof_file(filename + "." + str(child_pid), sample_type="lock-release")
 
@@ -74,7 +73,7 @@ def test_fork(tmp_path, monkeypatch):
 @pytest.mark.skipif(not os.getenv("DD_PROFILE_TEST_GEVENT", False), reason="Not testing gevent")
 def test_fork_gevent(monkeypatch):
     monkeypatch.setenv("DD_PROFILING_API_TIMEOUT_MS", "100")
-    stdout, stderr, exitcode, pid = call_program("python", os.path.join(os.path.dirname(__file__), "gevent_fork.py"))
+    _, _, exitcode, _ = call_program("python", os.path.join(os.path.dirname(__file__), "gevent_fork.py"))
     assert exitcode == 0
 
 
@@ -97,7 +96,7 @@ def test_multiprocessing(method, tmp_path, monkeypatch):
         method,
     )
     assert exitcode == 0, (stdout, stderr)
-    pid, child_pid = list(s.strip() for s in stdout.decode().strip().split("\n"))
+    pid, child_pid = list(s.strip() for s in cast(bytes, stdout).decode().strip().split("\n"))
     utils.check_pprof_file(filename + "." + str(pid))
     utils.check_pprof_file(filename + "." + str(child_pid), sample_type="wall-time")
 

tests/profiling/test_profiler.py

@@ -1,7 +1,7 @@
 import logging
 import time
+from unittest import mock
 
-import mock
 import pytest
 
 import ddtrace
@@ -105,7 +105,7 @@ def snapshot():
 
     class TestProfiler(profiler._ProfilerInstance):
         def _build_default_exporters(self, *args, **kargs):
-            return []
+            return None
 
     p = TestProfiler()
     err_collector = mock.MagicMock(wraps=ErrCollect())
@@ -140,7 +140,6 @@ def test_default_collectors():
 
 
 def test_profiler_serverless(monkeypatch):
-    # type: (...) -> None
     monkeypatch.setenv("AWS_LAMBDA_FUNCTION_NAME", "foobar")
     p = profiler.Profiler()
     assert isinstance(p._scheduler, scheduler.ServerlessScheduler)

tests/profiling_v2/test_profiler.py

@@ -1,8 +1,8 @@
 import logging
 import sys
 import time
+from unittest import mock
 
-import mock
 import pytest
 
 import ddtrace
@@ -107,7 +107,7 @@ def snapshot():
 
     class TestProfiler(profiler._ProfilerInstance):
         def _build_default_exporters(self, *args, **kargs):
-            return []
+            return None
 
     p = TestProfiler()
     err_collector = mock.MagicMock(wraps=ErrCollect())
@@ -182,7 +182,7 @@ def test_libdd_failure_telemetry_logging():
     2) import ddtrace.profiling.auto
     """
 
-    import mock
+    from unittest import mock
 
     with mock.patch.multiple(
         "ddtrace.internal.datadog.profiling.ddup",
@@ -206,7 +206,7 @@ def test_libdd_failure_telemetry_logging():
     err=None
 )
 def test_libdd_failure_telemetry_logging_with_auto():
-    import mock
+    from unittest import mock
 
     with mock.patch.multiple(
         "ddtrace.internal.datadog.profiling.ddup",
@@ -233,7 +233,7 @@ def test_stack_v2_failure_telemetry_logging():
     # mimicking the behavior of ddtrace-run, where the config is imported to
     # determine if profiling/stack_v2 is enabled
 
-    import mock
+    from unittest import mock
 
     with mock.patch.multiple(
         "ddtrace.internal.datadog.profiling.stack_v2",
@@ -257,7 +257,7 @@ def test_stack_v2_failure_telemetry_logging():
     err=None,
 )
 def test_stack_v2_failure_telemetry_logging_with_auto():
-    import mock
+    from unittest import mock
 
     with mock.patch.multiple(
         "ddtrace.internal.datadog.profiling.stack_v2",
@@ -305,7 +305,7 @@ def test_user_threads_have_native_id():
     for _ in range(10):
         try:
             # The TID should be higher than the PID, but not too high
-            assert 0 < t.native_id - getpid() < 100, (t.native_id, getpid())
+            assert 0 < t.native_id - getpid() < 100, (t.native_id, getpid())  # pyright: ignore[reportOptionalOperand]
         except AttributeError:
             # The native_id attribute is set by the thread so we might have to
             # wait a bit for it to be set.

tests/profiling_v2/test_pytorch.py

@@ -13,7 +13,7 @@ def test_call_script_pytorch_gpu(tmp_path, monkeypatch):
     monkeypatch.setenv("DD_PROFILING_OUTPUT_PPROF", filename)
     monkeypatch.setenv("DD_PROFILING_ENABLED", "1")
     monkeypatch.setenv("DD_PROFILING_PYTORCH_ENABLED", "1")
-    stdout, stderr, exitcode, pid = call_program(
+    _, stderr, exitcode, _ = call_program(
         "ddtrace-run", sys.executable, os.path.join(os.path.dirname(__file__), "simple_program_pytorch_gpu.py")
     )
     assert exitcode == 0, f"Profiler exited with code {exitcode}. Stderr: {stderr}"

tests/profiling_v2/test_scheduler.py

@@ -1,15 +1,15 @@
 # -*- encoding: utf-8 -*-
 import logging
 import time
-
-import mock
+from unittest import mock
 
 from ddtrace.profiling import scheduler
 
 
 def test_thread_name():
     s = scheduler.Scheduler()
     s.start()
+    assert s._worker is not None
     assert s._worker.name == "ddtrace.profiling.scheduler:Scheduler"
     s.stop()