13
13
# limitations under the License.
14
14
from unittest import TestCase , mock
15
15
16
+ from pika .channel import Channel
17
+ from pika .spec import Basic , BasicProperties
18
+
16
19
from opentelemetry .instrumentation .pika import utils
17
- from opentelemetry .semconv .trace import SpanAttributes
18
- from opentelemetry .trace import Span , Tracer
20
+ from opentelemetry .semconv .trace import (
21
+ MessagingOperationValues ,
22
+ SpanAttributes ,
23
+ )
24
+ from opentelemetry .trace import Span , SpanKind , Tracer
19
25
20
26
21
27
class TestUtils (TestCase ):
@@ -32,12 +38,15 @@ def test_get_span(
32
38
channel = mock .MagicMock ()
33
39
properties = mock .MagicMock ()
34
40
task_name = "test.test"
41
+ span_kind = mock .MagicMock (spec = SpanKind )
35
42
get_value .return_value = None
36
43
ctx = mock .MagicMock ()
37
- _ = utils ._get_span (tracer , channel , properties , task_name , ctx )
44
+ _ = utils ._get_span (
45
+ tracer , channel , properties , task_name , span_kind , ctx
46
+ )
38
47
generate_span_name .assert_called_once ()
39
48
tracer .start_span .assert_called_once_with (
40
- context = ctx , name = generate_span_name .return_value
49
+ context = ctx , name = generate_span_name .return_value , kind = span_kind
41
50
)
42
51
enrich_span .assert_called_once ()
43
52
@@ -54,9 +63,12 @@ def test_get_span_suppressed(
54
63
channel = mock .MagicMock ()
55
64
properties = mock .MagicMock ()
56
65
task_name = "test.test"
66
+ span_kind = mock .MagicMock (spec = SpanKind )
57
67
get_value .return_value = True
58
68
ctx = mock .MagicMock ()
59
- span = utils ._get_span (tracer , channel , properties , task_name , ctx )
69
+ span = utils ._get_span (
70
+ tracer , channel , properties , task_name , span_kind , ctx
71
+ )
60
72
self .assertEqual (span , None )
61
73
generate_span_name .assert_not_called ()
62
74
enrich_span .assert_not_called ()
@@ -158,3 +170,117 @@ def test_enrich_span_unique_connection() -> None:
158
170
),
159
171
],
160
172
)
173
+
174
+ @mock .patch ("opentelemetry.instrumentation.pika.utils._get_span" )
175
+ @mock .patch ("opentelemetry.propagate.extract" )
176
+ @mock .patch ("opentelemetry.trace.use_span" )
177
+ def test_decorate_callback (
178
+ self ,
179
+ use_span : mock .MagicMock ,
180
+ extract : mock .MagicMock ,
181
+ get_span : mock .MagicMock ,
182
+ ) -> None :
183
+ callback = mock .MagicMock ()
184
+ mock_task_name = "mock_task_name"
185
+ tracer = mock .MagicMock ()
186
+ channel = mock .MagicMock (spec = Channel )
187
+ method = mock .MagicMock (spec = Basic .Deliver )
188
+ properties = mock .MagicMock ()
189
+ mock_body = b"mock_body"
190
+ decorated_callback = utils ._decorate_callback (
191
+ callback , tracer , mock_task_name
192
+ )
193
+ retval = decorated_callback (channel , method , properties , mock_body )
194
+ extract .assert_called_once_with (
195
+ properties .headers , getter = utils ._pika_getter
196
+ )
197
+ get_span .assert_called_once_with (
198
+ tracer ,
199
+ channel ,
200
+ properties ,
201
+ span_kind = SpanKind .CONSUMER ,
202
+ task_name = mock_task_name ,
203
+ ctx = extract .return_value ,
204
+ operation = MessagingOperationValues .RECEIVE ,
205
+ )
206
+ use_span .assert_called_once_with (
207
+ get_span .return_value , end_on_exit = True
208
+ )
209
+ callback .assert_called_once_with (
210
+ channel , method , properties , mock_body
211
+ )
212
+ self .assertEqual (retval , callback .return_value )
213
+
214
+ @mock .patch ("opentelemetry.instrumentation.pika.utils._get_span" )
215
+ @mock .patch ("opentelemetry.propagate.inject" )
216
+ @mock .patch ("opentelemetry.context.get_current" )
217
+ @mock .patch ("opentelemetry.trace.use_span" )
218
+ def test_decorate_basic_publish (
219
+ self ,
220
+ use_span : mock .MagicMock ,
221
+ get_current : mock .MagicMock ,
222
+ inject : mock .MagicMock ,
223
+ get_span : mock .MagicMock ,
224
+ ) -> None :
225
+ callback = mock .MagicMock ()
226
+ tracer = mock .MagicMock ()
227
+ channel = mock .MagicMock (spec = Channel )
228
+ method = mock .MagicMock (spec = Basic .Deliver )
229
+ properties = mock .MagicMock ()
230
+ mock_body = b"mock_body"
231
+ decorated_basic_publish = utils ._decorate_basic_publish (
232
+ callback , channel , tracer
233
+ )
234
+ retval = decorated_basic_publish (
235
+ channel , method , mock_body , properties
236
+ )
237
+ get_current .assert_called_once ()
238
+ get_span .assert_called_once_with (
239
+ tracer ,
240
+ channel ,
241
+ properties ,
242
+ span_kind = SpanKind .PRODUCER ,
243
+ task_name = "(temporary)" ,
244
+ ctx = get_current .return_value ,
245
+ operation = None ,
246
+ )
247
+ use_span .assert_called_once_with (
248
+ get_span .return_value , end_on_exit = True
249
+ )
250
+ get_span .return_value .is_recording .assert_called_once ()
251
+ inject .assert_called_once_with (properties .headers )
252
+ callback .assert_called_once_with (
253
+ channel , method , mock_body , properties , False
254
+ )
255
+ self .assertEqual (retval , callback .return_value )
256
+
257
+ @mock .patch ("opentelemetry.instrumentation.pika.utils._get_span" )
258
+ @mock .patch ("opentelemetry.propagate.inject" )
259
+ @mock .patch ("opentelemetry.context.get_current" )
260
+ @mock .patch ("opentelemetry.trace.use_span" )
261
+ @mock .patch ("pika.spec.BasicProperties.__new__" )
262
+ def test_decorate_basic_publish_no_properties (
263
+ self ,
264
+ basic_properties : mock .MagicMock ,
265
+ use_span : mock .MagicMock ,
266
+ get_current : mock .MagicMock ,
267
+ inject : mock .MagicMock ,
268
+ get_span : mock .MagicMock ,
269
+ ) -> None :
270
+ callback = mock .MagicMock ()
271
+ tracer = mock .MagicMock ()
272
+ channel = mock .MagicMock (spec = Channel )
273
+ method = mock .MagicMock (spec = Basic .Deliver )
274
+ mock_body = b"mock_body"
275
+ decorated_basic_publish = utils ._decorate_basic_publish (
276
+ callback , channel , tracer
277
+ )
278
+ retval = decorated_basic_publish (channel , method , body = mock_body )
279
+ basic_properties .assert_called_once_with (BasicProperties , headers = {})
280
+ get_current .assert_called_once ()
281
+ use_span .assert_called_once_with (
282
+ get_span .return_value , end_on_exit = True
283
+ )
284
+ get_span .return_value .is_recording .assert_called_once ()
285
+ inject .assert_called_once_with (basic_properties .return_value .headers )
286
+ self .assertEqual (retval , callback .return_value )
0 commit comments