diff --git a/tests/contrib/aiohttp/test_aiohttp_client.py b/tests/contrib/aiohttp/test_aiohttp_client.py index f272c348324..73ff426e7f5 100644 --- 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 index 72fe1806979..425d205c361 100644 --- 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)