fix get_from in testpin

DataDog · Feb 4, 2025 · 4765935 · 4765935
1 parent d3624a0
commit 4765935
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 1 deletion.
diff --git a/tests/contrib/aiohttp/test_aiohttp_client.py b/tests/contrib/aiohttp/test_aiohttp_client.py
@@ -184,7 +184,7 @@ def test_configure_service_name_pin(ddtrace_run_python_code_in_subprocess):
 import asyncio
 import sys
 import aiohttp
-from tests.utils import TestPin as Pin
+from ddtrace.trace import Pin
 from tests.contrib.aiohttp.test_aiohttp_client import URL_200
 
 async def test():

diff --git a/tests/utils.py b/tests/utils.py
@@ -1038,6 +1038,15 @@ def clone(
             _config=pin._config,
         )
 
+    @staticmethod
+    def get_from(obj):
+        pin = Pin.get_from(obj)
+        if isinstance(pin, Pin):
+            # replace ddtrace.trace.Pin with tests.utils.TestPin in tests
+            pin.remove_from(obj)
+            return TestPin(service=pin.service, tags=pin.tags, tracer=pin.tracer, _config=pin._config).onto(obj)
+        return pin
+
 
 def assert_dict_issuperset(a, b):
     assert set(a.items()).issuperset(set(b.items())), "{a} is not a superset of {b}".format(a=a, b=b)