def test_decorate_callback(
self,
use_span: mock.MagicMock,
extract: mock.MagicMock,
get_span: mock.MagicMock,
) -> None:
callback = mock.MagicMock()
mock_task_name = "mock_task_name"
tracer = mock.MagicMock()
channel = mock.MagicMock(spec=Channel)
method = mock.MagicMock(spec=Basic.Deliver)
properties = mock.MagicMock()
mock_body = b"mock_body"
decorated_callback = utils._decorate_callback(callback, tracer,
mock_task_name)
retval = decorated_callback(channel, method, properties, mock_body)
extract.assert_called_once_with(properties.headers,
getter=utils._pika_getter)
get_span.assert_called_once_with(
tracer,
channel,
properties,
span_kind=SpanKind.CONSUMER,
task_name=mock_task_name,
ctx=extract.return_value,
operation=MessagingOperationValues.RECEIVE,
)
use_span.assert_called_once_with(get_span.return_value,
end_on_exit=True)
callback.assert_called_once_with(channel, method, properties,
mock_body)
self.assertEqual(retval, callback.return_value)
def _instrument_consumers(consumers_dict: Dict[str, Callable[..., Any]],
tracer: Tracer) -> Any:
for key, callback in consumers_dict.items():
decorated_callback = utils._decorate_callback(
callback, tracer, key)
setattr(decorated_callback, "_original_callback", callback)
consumers_dict[key] = decorated_callback
def _instrument_blocking_channel_consumers(
channel: BlockingChannel,
tracer: Tracer,
consume_hook: utils.HookT = utils.dummy_callback,
) -> Any:
for consumer_tag, consumer_info in channel._consumer_infos.items():
callback_attr = PikaInstrumentor.CONSUMER_CALLBACK_ATTR
consumer_callback = getattr(consumer_info, callback_attr, None)
if consumer_callback is None:
continue
decorated_callback = utils._decorate_callback(
consumer_callback,
tracer,
consumer_tag,
consume_hook,
)
setattr(
decorated_callback,
"_original_callback",
consumer_callback,
)
setattr(consumer_info, callback_attr, decorated_callback)