mirror of
https://github.com/open-telemetry/opentelemetry-python-contrib.git
synced 2025-08-01 06:33:52 +08:00
92 lines
3.3 KiB
Python
92 lines
3.3 KiB
Python
import asyncio
|
|
import pytest
|
|
|
|
from ddtrace.context import Context
|
|
from ddtrace.internal.context_manager import CONTEXTVARS_IS_AVAILABLE
|
|
from ddtrace.contrib.asyncio import helpers
|
|
from .utils import AsyncioTestCase, mark_asyncio
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
CONTEXTVARS_IS_AVAILABLE,
|
|
reason='only applicable to legacy asyncio integration'
|
|
)
|
|
class TestAsyncioHelpers(AsyncioTestCase):
|
|
"""
|
|
Ensure that helpers set the ``Context`` properly when creating
|
|
new ``Task`` or threads.
|
|
"""
|
|
@mark_asyncio
|
|
def test_set_call_context(self):
|
|
# a different Context is set for the current logical execution
|
|
task = asyncio.Task.current_task()
|
|
ctx = Context()
|
|
helpers.set_call_context(task, ctx)
|
|
assert ctx == self.tracer.get_call_context()
|
|
|
|
@mark_asyncio
|
|
def test_ensure_future(self):
|
|
# the wrapper should create a new Future that has the Context attached
|
|
@asyncio.coroutine
|
|
def future_work():
|
|
# the ctx is available in this task
|
|
ctx = self.tracer.get_call_context()
|
|
assert 1 == len(ctx._trace)
|
|
assert 'coroutine' == ctx._trace[0].name
|
|
return ctx._trace[0].name
|
|
|
|
self.tracer.trace('coroutine')
|
|
# schedule future work and wait for a result
|
|
delayed_task = helpers.ensure_future(future_work(), tracer=self.tracer)
|
|
result = yield from asyncio.wait_for(delayed_task, timeout=1)
|
|
assert 'coroutine' == result
|
|
|
|
@mark_asyncio
|
|
def test_run_in_executor_proxy(self):
|
|
# the wrapper should pass arguments and results properly
|
|
def future_work(number, name):
|
|
assert 42 == number
|
|
assert 'john' == name
|
|
return True
|
|
|
|
future = helpers.run_in_executor(self.loop, None, future_work, 42, 'john', tracer=self.tracer)
|
|
result = yield from future
|
|
assert result
|
|
|
|
@mark_asyncio
|
|
def test_run_in_executor_traces(self):
|
|
# the wrapper should create a different Context when the Thread
|
|
# is started; the new Context creates a new trace
|
|
def future_work():
|
|
# the Context is empty but the reference to the latest
|
|
# span is here to keep the parenting
|
|
ctx = self.tracer.get_call_context()
|
|
assert 0 == len(ctx._trace)
|
|
assert 'coroutine' == ctx._current_span.name
|
|
return True
|
|
|
|
span = self.tracer.trace('coroutine')
|
|
future = helpers.run_in_executor(self.loop, None, future_work, tracer=self.tracer)
|
|
# we close the Context
|
|
span.finish()
|
|
result = yield from future
|
|
assert result
|
|
|
|
@mark_asyncio
|
|
def test_create_task(self):
|
|
# the helper should create a new Task that has the Context attached
|
|
@asyncio.coroutine
|
|
def future_work():
|
|
# the ctx is available in this task
|
|
ctx = self.tracer.get_call_context()
|
|
assert 0 == len(ctx._trace)
|
|
child_span = self.tracer.trace('child_task')
|
|
return child_span
|
|
|
|
root_span = self.tracer.trace('main_task')
|
|
# schedule future work and wait for a result
|
|
task = helpers.create_task(future_work())
|
|
result = yield from task
|
|
assert root_span.trace_id == result.trace_id
|
|
assert root_span.span_id == result.parent_id
|