PR comments; rebased

Author: vlad-scherbich

tests/profiling/collector/test_asyncio.py

@@ -3,13 +3,15 @@
 import glob
 import os
 import sys
+from typing import Optional
 import uuid
 
 import pytest
 
 from ddtrace import ext
 from ddtrace.internal.datadog.profiling import ddup
 from ddtrace.profiling.collector import asyncio as collector_asyncio
+from ddtrace.profiling.collector._lock import _LockAllocatorWrapper as LockAllocatorWrapper
 from tests.profiling.collector import pprof_utils
 from tests.profiling.collector import test_collector
 from tests.profiling.collector.lock_utils import get_lock_linenos
@@ -53,28 +55,18 @@ def teardown_method(self):
 
     async def test_subclassing_wrapped_lock(self) -> None:
         """Test that subclassing of a wrapped lock type when profiling is active."""
-        from typing import Optional
-
-        from ddtrace.profiling.collector._lock import _LockAllocatorWrapper
-
         with collector_asyncio.AsyncioLockCollector(capture_pct=100):
-            assert isinstance(asyncio.Lock, _LockAllocatorWrapper)
+            assert isinstance(asyncio.Lock, LockAllocatorWrapper)
 
             # This should NOT raise TypeError
             class CustomLock(asyncio.Lock):  # type: ignore[misc]
                 def __init__(self) -> None:
                     super().__init__()
-                    self._owner: Optional[int] = None
-                    self._count: int = 0
 
             # Verify subclassing and functionality
             custom_lock: CustomLock = CustomLock()
-            assert hasattr(custom_lock, "_owner")
-            assert hasattr(custom_lock, "_count")
-            assert custom_lock._owner is None
-            assert custom_lock._count == 0
 
-            # Test async acquire/release
+            # Verify subclassing and functionality
             await custom_lock.acquire()
             assert custom_lock.locked()
             custom_lock.release()

tests/profiling/collector/test_threading.py

@@ -173,16 +173,14 @@ def test_lock_repr(
 
 
 def test_patch():
-    from ddtrace.profiling.collector._lock import _LockAllocatorWrapper
-
     lock = threading.Lock
     collector = ThreadingLockCollector()
     collector.start()
     assert lock == collector._original_lock
     # After patching, threading.Lock is replaced with our wrapper
     # The old reference (lock) points to the original builtin Lock class
     assert lock != threading.Lock  # They're different after patching
-    assert isinstance(threading.Lock, _LockAllocatorWrapper)  # threading.Lock is now wrapped
+    assert isinstance(threading.Lock, LockAllocatorWrapper)  # threading.Lock is now wrapped
     assert callable(threading.Lock)  # and it's callable
     collector.stop()
     # After stopping, everything is restored
@@ -1396,22 +1394,16 @@ def test_subclassing_wrapped_lock(self) -> None:
         threading.Lock and threading.RLock are C types that don't support subclassing
         through __mro_entries__.
         """
-        from ddtrace.profiling.collector._lock import _LockAllocatorWrapper
-
         with self.collector_class():
-            assert isinstance(self.lock_class, _LockAllocatorWrapper)
+            assert isinstance(self.lock_class, LockAllocatorWrapper)
 
             # This should NOT raise TypeError
             class CustomLock(self.lock_class):  # type: ignore[misc]
                 def __init__(self) -> None:
                     super().__init__()
-                    self.custom_attr: str = "test"
 
             # Verify subclassing and functionality
             custom_lock: CustomLock = CustomLock()
-            assert hasattr(custom_lock, "custom_attr")
-            assert custom_lock.custom_attr == "test"
-
             assert custom_lock.acquire()
             custom_lock.release()