opencensus-python/contrib/opencensus-ext-threading/tests/test_tracer.py at a099a0252034cef17f1bfea474357ebe77cf42d1 · census-instrumentation/opencensus-python · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
import unittest
from unittest.mock import patch, MagicMock
from opencensus.ext.threading.trace import wrap_submit, wrap_apply_async


class TestTracer(unittest.TestCase):
    """
    Ensures that sampler, exporter, propagator are passed through
    in case global tracer is present.
    """

    @patch("opencensus.trace.propagation.binary_format.BinaryFormatPropagator")
    @patch("opencensus.ext.threading.trace.wrap_task_func")
    @patch("opencensus.trace.execution_context.get_opencensus_tracer")
    def test_apply_async_context_passed(
        self,
        get_opencensus_tracer_mock: MagicMock,
        wrap_task_func_mock: MagicMock,
        binary_format_propagator_mock: MagicMock,
    ):
        mock_tracer = NoNoopTracerMock()
        # ensure that unique object is generated
        mock_tracer.sampler = MagicMock()
        mock_tracer.exporter = MagicMock()
        mock_tracer.propagator = MagicMock()

        get_opencensus_tracer_mock.return_value = mock_tracer

        submission_function_mock = MagicMock()
        original_function_mock = MagicMock()

        wrap_apply_async(submission_function_mock)(None, original_function_mock)

        # check whether invocation of original function _has_ happened
        call = submission_function_mock.call_args_list[0].kwargs

        self.assertEqual(id(call["kwds"]["sampler"]), id(mock_tracer.sampler))
        self.assertEqual(id(call["kwds"]["exporter"]), id(mock_tracer.exporter))
        self.assertEqual(id(call["kwds"]["propagator"]), id(mock_tracer.propagator))

    @patch("opencensus.trace.propagation.binary_format.BinaryFormatPropagator")
    @patch("opencensus.ext.threading.trace.wrap_task_func")
    @patch("opencensus.trace.execution_context.get_opencensus_tracer")
    def test_wrap_submit_context_passed(
        self,
        get_opencensus_tracer_mock: MagicMock,
        wrap_task_func_mock: MagicMock,
        binary_format_propagator_mock: MagicMock,
    ):
        mock_tracer = NoNoopTracerMock()
        # ensure that unique object is generated
        mock_tracer.sampler = MagicMock()
        mock_tracer.exporter = MagicMock()
        mock_tracer.propagator = MagicMock()

        get_opencensus_tracer_mock.return_value = mock_tracer

        submission_function_mock = MagicMock()
        original_function_mock = MagicMock()

        wrap_submit(submission_function_mock)(None, original_function_mock)

        # check whether invocation of original function _has_ happened
        call = submission_function_mock.call_args_list[0].kwargs

        self.assertEqual(id(call["sampler"]), id(mock_tracer.sampler))
        self.assertEqual(id(call["exporter"]), id(mock_tracer.exporter))
        self.assertEqual(id(call["propagator"]), id(mock_tracer.propagator))


class NoNoopTracerMock(MagicMock):
    pass