def test_list_collected_spans(self):
        tracer = asyncio_context_tracer.ContextTracer()
        span1 = mock.Mock()
        span2 = mock.Mock()
        tracer._spans_list.append(span1)
        tracer._spans_list.append(span2)

        spans = tracer.list_collected_spans()

        self.assertEqual(spans, [span1, span2])
    def test_end_span_batch_export(self):
        exporter = mock.Mock()
        tracer = asyncio_context_tracer.ContextTracer(exporter=exporter)
        span = tracer.start_span('test')
        parent_span_id = '6e0c63257de34c92'
        span.parent_span.span_id = parent_span_id
        span.finish = mock.Mock()
        tracer.end_span()

        self.assertTrue(span.finish.called)
        self.assertEqual(tracer.span_context.span_id, parent_span_id)
        self.assertTrue(tracer.exporter.export.called)
    def test_constructor_explicit(self):
        from opencensus.trace import span_context

        span_context = span_context.SpanContext()
        exporter = mock.Mock()
        tracer = asyncio_context_tracer.ContextTracer(
            exporter=exporter, span_context=span_context)

        self.assertIs(tracer.span_context, span_context)
        self.assertIs(tracer.exporter, exporter)
        self.assertEqual(tracer.trace_id, span_context.trace_id)
        self.assertEqual(tracer._spans_list, [])
        self.assertEqual(tracer.root_span_id, span_context.span_id)
    def test_add_attribute_to_current_span(self):
        from opencensus.trace.span import Span
        from qubit.opencensus.trace import asyncio_context

        tracer = asyncio_context_tracer.ContextTracer()
        span1 = mock.Mock(spec=Span)

        span1.attributes = {}
        asyncio_context.set_current_span(span1)

        attribute_key = 'key'
        attribute_value = 'value'

        tracer.add_attribute_to_current_span(attribute_key, attribute_value)

        span1.add_attribute.assert_called_once_with(attribute_key,
                                                    attribute_value)
    def test_start_span(self, current_span_mock):
        from opencensus.trace import span_context

        span_id = '6e0c63257de34c92'
        span_name = 'test_span'
        mock_span = mock.Mock()
        mock_span.span_id = span_id
        span_context = span_context.SpanContext(span_id=span_id)
        tracer = asyncio_context_tracer.ContextTracer(
            span_context=span_context)
        current_span_mock.return_value = mock_span

        span = tracer.start_span(name=span_name)

        self.assertEqual(span.parent_span.span_id, span_id)
        self.assertEqual(span.name, span_name)
        self.assertEqual(len(tracer._spans_list), 1)
        self.assertEqual(span_context.span_id, span.span_id)
    def test_end_span_without_parent(self, mock_current_span):
        from qubit.opencensus.trace.asyncio_context import get_current_span

        tracer = asyncio_context_tracer.ContextTracer()
        mock_span = mock.Mock()
        mock_span.name = 'span'
        mock_span.children = []
        mock_span.status = None
        mock_span.links = None
        mock_span.stack_trace = None
        mock_span.time_events = None
        mock_span.attributes = {}
        mock_span.__iter__ = mock.Mock(return_value=iter([mock_span]))
        mock_current_span.return_value = mock_span
        tracer.end_span()

        cur_span = get_current_span()
        self.assertIsNone(cur_span)
    def test_end_span_active(self, mock_current_span):
        exporter = mock.Mock()
        tracer = asyncio_context_tracer.ContextTracer(exporter=exporter)
        mock_span = mock.Mock()
        mock_span.name = 'span'
        mock_span.children = []
        mock_span.status = None
        mock_span.links = None
        mock_span.stack_trace = None
        mock_span.time_events = None
        mock_span.attributes = {}
        mock_span.__iter__ = mock.Mock(return_value=iter([mock_span]))
        parent_span_id = '6e0c63257de34c92'
        mock_span.parent_span.span_id = parent_span_id
        mock_current_span.return_value = mock_span
        tracer.end_span()

        self.assertTrue(mock_span.finish.called)
        self.assertEqual(tracer.span_context.span_id, parent_span_id)
        self.assertFalse(tracer.exporter.export.called)
    def test_end_leftover_spans(self):
        tracer = asyncio_context_tracer.ContextTracer()
        tracer._spans_list = [span.Span(name='span')]
        tracer.finish()

        self.assertEqual(tracer._spans_list, [])
    def test_finish_with_spans(self):
        tracer = asyncio_context_tracer.ContextTracer()
        tracer.start_span('span')
        tracer.finish()

        self.assertEqual(tracer._spans_list, [])