|
| 1 | +import asyncio |
| 2 | + |
| 3 | +import pytest |
| 4 | + |
| 5 | +from ddtrace.contrib.asyncio.compat import asyncio_current_task |
| 6 | +from ddtrace.internal.compat import CONTEXTVARS_IS_AVAILABLE |
| 7 | +from ddtrace.provider import DefaultContextProvider |
| 8 | +from tests.utils import flaky |
| 9 | + |
| 10 | + |
| 11 | +pytestmark = pytest.mark.skipif( |
| 12 | + CONTEXTVARS_IS_AVAILABLE, reason="No configuration is necessary when contextvars available." |
| 13 | +) |
| 14 | + |
| 15 | + |
| 16 | +@pytest.mark.asyncio |
| 17 | +async def test_get_call_context(tracer): |
| 18 | + tracer.configure(context_provider=DefaultContextProvider()) |
| 19 | + ctx = tracer.current_trace_context() |
| 20 | + assert ctx is None |
| 21 | + # test that it behaves the wrong way |
| 22 | + task = asyncio_current_task() |
| 23 | + assert task |
| 24 | + task_ctx = getattr(task, "__datadog_context", None) |
| 25 | + assert task_ctx is None |
| 26 | + |
| 27 | + |
| 28 | +def test_trace_coroutine(tracer): |
| 29 | + # it should use the task context when invoked in a coroutine |
| 30 | + with tracer.trace("coroutine") as span: |
| 31 | + span.resource = "base" |
| 32 | + |
| 33 | + traces = tracer.pop_traces() |
| 34 | + assert 1 == len(traces) |
| 35 | + assert 1 == len(traces[0]) |
| 36 | + assert "coroutine" == traces[0][0].name |
| 37 | + assert "base" == traces[0][0].resource |
| 38 | + |
| 39 | + |
| 40 | +@flaky(until=1706677200) |
| 41 | +@pytest.mark.asyncio |
| 42 | +async def test_trace_multiple_calls(tracer): |
| 43 | + tracer.configure(context_provider=DefaultContextProvider()) |
| 44 | + |
| 45 | + async def coro(): |
| 46 | + # another traced coroutine |
| 47 | + with tracer.trace("coroutine"): |
| 48 | + await asyncio.sleep(0.01) |
| 49 | + |
| 50 | + # partial flushing is enabled, ensure the number of spans generated is less than 500 |
| 51 | + futures = [asyncio.ensure_future(coro()) for x in range(400)] |
| 52 | + for future in futures: |
| 53 | + await future |
| 54 | + |
| 55 | + # the trace is wrong but the Context is finished |
| 56 | + traces = tracer.pop_traces() |
| 57 | + assert 1 == len(traces) |
| 58 | + assert 400 == len(traces[0]) |
0 commit comments