more tests, more cleanup

Author: vlad-scherbich

tests/profiling_v2/collector/test_threading.py

@@ -13,29 +13,18 @@
 
 from ddtrace import ext
 from ddtrace.internal.datadog.profiling import ddup
-from ddtrace.profiling.collector import threading as collector_threading
+from ddtrace.profiling.collector.threading import ThreadingLockCollector
+from ddtrace.profiling.collector.threading import ThreadingRLockCollector
 from tests.profiling.collector import pprof_utils
 from tests.profiling.collector import test_collector
 from tests.profiling.collector.lock_utils import get_lock_linenos
 from tests.profiling.collector.lock_utils import init_linenos
 
 
-# Module-level globals for testing global lock profiling
-_test_global_lock = None
-_test_global_bar_instance = None
-
-
-# Type aliases for supported classes
-LockClass = Union[Type[threading.Lock], Type[threading.RLock]]
-CollectorClass = Union[
-    Type[collector_threading.ThreadingLockCollector], Type[collector_threading.ThreadingRLockCollector]
-]
-
-
 # Type aliases for supported classes
 LockClass = Union[Type[threading.Lock], Type[threading.RLock]]
 CollectorClass = Union[
-    Type[collector_threading.ThreadingLockCollector], Type[collector_threading.ThreadingRLockCollector]
+    Type[ThreadingLockCollector], Type[ThreadingRLockCollector]
 ]
 
 # Module-level globals for testing global lock profiling
@@ -65,25 +54,50 @@ def bar(self) -> None:
         self.foo.foo()
 
 
-def test_repr() -> None:
-    test_collector._test_repr(
-        collector_threading.ThreadingLockCollector,
-        "ThreadingLockCollector(status=<ServiceStatus.STOPPED: 'stopped'>, "
-        "capture_pct=1.0, nframes=64, "
-        "endpoint_collection_enabled=True, tracer=None)",
-    )
-
-
-def test_patch() -> None:
-    lock = threading.Lock
-    collector = collector_threading.ThreadingLockCollector()
+@pytest.mark.parametrize(
+    "collector_class,expected_repr",
+    [
+        (
+            ThreadingLockCollector,
+            "ThreadingLockCollector(status=<ServiceStatus.STOPPED: 'stopped'>, "
+            "capture_pct=1.0, nframes=64, "
+            "endpoint_collection_enabled=True, tracer=None)",
+        ),
+        (
+            ThreadingRLockCollector,
+            "ThreadingRLockCollector(status=<ServiceStatus.STOPPED: 'stopped'>, "
+            "capture_pct=1.0, nframes=64, "
+            "endpoint_collection_enabled=True, tracer=None)",
+        ),
+    ],
+)
+def test_repr(
+    collector_class: CollectorClass,
+    expected_repr: str,
+) -> None:
+    test_collector._test_repr(collector_class, expected_repr)
+
+
+@pytest.mark.parametrize(
+    "lock_class,collector_class",
+    [
+        (threading.Lock, ThreadingLockCollector),
+        (threading.RLock, ThreadingRLockCollector),
+    ],
+)
+def test_patch(
+    lock_class: LockClass,
+    collector_class: CollectorClass,
+) -> None:
+    lock = lock_class
+    collector = collector_class()
     collector.start()
     assert lock == collector._original
     # wrapt makes this true
-    assert lock == threading.Lock
+    assert lock == lock_class
     collector.stop()
-    assert lock == threading.Lock
-    assert collector._original == threading.Lock
+    assert lock == lock_class
+    assert collector._original == lock_class
 
 
 @pytest.mark.skipif(not sys.platform.startswith("linux"), reason="only works on linux")
@@ -143,7 +157,7 @@ def test_wrapt_disable_extensions() -> None:
 
     from ddtrace.internal.datadog.profiling import ddup
     from ddtrace.profiling.collector import _lock
-    from ddtrace.profiling.collector import threading as collector_threading
+    from ddtrace.profiling.collector.threading import ThreadingLockCollector
     from tests.profiling.collector import pprof_utils
     from tests.profiling.collector.lock_utils import get_lock_linenos
     from tests.profiling.collector.lock_utils import init_linenos
@@ -165,7 +179,7 @@ def test_wrapt_disable_extensions() -> None:
     assert os.environ.get("WRAPT_DISABLE_EXTENSIONS")
     assert _lock.WRAPT_C_EXT is False
 
-    with collector_threading.ThreadingLockCollector(capture_pct=100):
+    with ThreadingLockCollector(capture_pct=100):
         th_lock = threading.Lock()  # !CREATE! test_wrapt_disable_extensions
         with th_lock:  # !ACQUIRE! !RELEASE! test_wrapt_disable_extensions
             pass
@@ -214,7 +228,7 @@ def test_lock_gevent_tasks() -> None:
     import threading
 
     from ddtrace.internal.datadog.profiling import ddup
-    from ddtrace.profiling.collector import threading as collector_threading
+    from ddtrace.profiling.collector.threading import ThreadingLockCollector
     from tests.profiling.collector import pprof_utils
     from tests.profiling.collector.lock_utils import get_lock_linenos
     from tests.profiling.collector.lock_utils import init_linenos
@@ -235,7 +249,7 @@ def play_with_lock() -> None:
         lock.acquire()  # !ACQUIRE! test_lock_gevent_tasks
         lock.release()  # !RELEASE! test_lock_gevent_tasks
 
-    with collector_threading.ThreadingLockCollector(capture_pct=100):
+    with ThreadingLockCollector(capture_pct=100):
         t = threading.Thread(name="foobar", target=play_with_lock)
         t.start()
         t.join()
@@ -324,7 +338,7 @@ def teardown_method(self, method: Any) -> None:
 
     def test_wrapper(self) -> None:
         # TODO: change to collector_class
-        collector = collector_threading.ThreadingLockCollector()
+        collector = ThreadingLockCollector()
         with collector:
 
             class Foobar(object):
@@ -919,8 +933,8 @@ class TestThreadingLockCollector(BaseThreadingLockCollectorTest):
     """Test threading.Lock profiling"""
 
     @property
-    def collector_class(self) -> Type[collector_threading.ThreadingLockCollector]:
-        return collector_threading.ThreadingLockCollector
+    def collector_class(self) -> Type[ThreadingLockCollector]:
+        return ThreadingLockCollector
 
     @property
     def lock_class(self) -> Type[threading.Lock]:
@@ -931,8 +945,8 @@ class TestThreadingRLockCollector(BaseThreadingLockCollectorTest):
     """Test threading.RLock profiling"""
 
     @property
-    def collector_class(self) -> Type[collector_threading.ThreadingRLockCollector]:
-        return collector_threading.ThreadingRLockCollector
+    def collector_class(self) -> Type[ThreadingRLockCollector]:
+        return ThreadingRLockCollector
 
     @property
     def lock_class(self) -> Type[threading.RLock]: