import asyncio

from ddtrace.provider import DefaultContextProvider
from .utils import AsyncioTestCase, mark_asyncio


class TestAsyncioSafety(AsyncioTestCase):
    """
    Ensure that if the ``AsyncioTracer`` is not properly configured,
    bad traces are produced but the ``Context`` object will not
    leak memory.
    """

    def setUp(self):
        # Asyncio TestCase with the wrong context provider
        super(TestAsyncioSafety, self).setUp()
        self.tracer.configure(context_provider=DefaultContextProvider())

    @mark_asyncio
    def test_get_call_context(self):
        # it should return a context even if not attached to the Task
        ctx = self.tracer.get_call_context()
        assert ctx is not None
        # test that it behaves the wrong way
        task = asyncio.Task.current_task()
        task_ctx = getattr(task, "__datadog_context", None)
        assert task_ctx is None

    @mark_asyncio
    def test_trace_coroutine(self):
        # it should use the task context when invoked in a coroutine
        with self.tracer.trace("coroutine") as span:
            span.resource = "base"

        traces = self.tracer.writer.pop_traces()
        assert 1 == len(traces)
        assert 1 == len(traces[0])
        assert "coroutine" == traces[0][0].name
        assert "base" == traces[0][0].resource

    @mark_asyncio
    def test_trace_multiple_calls(self):
        @asyncio.coroutine
        def coro():
            # another traced coroutine
            with self.tracer.trace("coroutine"):
                yield from asyncio.sleep(0.01)

        ctx = self.tracer.get_call_context()
        futures = [asyncio.ensure_future(coro()) for x in range(1000)]
        for future in futures:
            yield from future

        # the trace is wrong but the Context is finished
        traces = self.tracer.writer.pop_traces()
        assert 1 == len(traces)
        assert 1000 == len(traces[0])
        assert 0 == len(ctx._trace)
