def test_instrument(
self,
instrument_consumers: mock.MagicMock,
instrument_channel_functions: mock.MagicMock,
):
PikaInstrumentor.instrument_channel(channel=self.channel)
assert hasattr(
self.channel, "_is_instrumented_by_opentelemetry"
), "channel is not marked as instrumented!"
instrument_consumers.assert_called_once()
instrument_channel_functions.assert_called_once()
def test_instrument_basic_publish(
self, decorate_basic_publish: mock.MagicMock
) -> None:
tracer = mock.MagicMock(spec=Tracer)
original_function = self.channel.basic_publish
PikaInstrumentor._instrument_basic_publish(self.channel, tracer)
decorate_basic_publish.assert_called_once_with(
original_function, self.channel, tracer
)
self.assertEqual(
self.channel.basic_publish, decorate_basic_publish.return_value
)
def test_instrument_consumers(
self, decorate_callback: mock.MagicMock
) -> None:
tracer = mock.MagicMock(spec=Tracer)
expected_decoration_calls = [
mock.call(value, tracer, key)
for key, value in self.channel._impl._consumers.items()
]
PikaInstrumentor._instrument_consumers(
self.channel._impl._consumers, tracer
)
decorate_callback.assert_has_calls(
calls=expected_decoration_calls, any_order=True
)
assert all(
hasattr(callback, "_original_callback")
for callback in self.channel._impl._consumers.values()
)
def test_instrument_consumers(
self, decorate_callback: mock.MagicMock
) -> None:
tracer = mock.MagicMock(spec=Tracer)
callback_attr = PikaInstrumentor.CONSUMER_CALLBACK_ATTR
expected_decoration_calls = [
mock.call(
getattr(value, callback_attr), tracer, key, dummy_callback
)
for key, value in self.channel._consumer_infos.items()
]
PikaInstrumentor._instrument_blocking_channel_consumers(
self.channel, tracer
)
decorate_callback.assert_has_calls(
calls=expected_decoration_calls, any_order=True
)
assert all(
hasattr(callback, "_original_callback")
for callback in self.channel._consumer_infos.values()
)
def test_instrument_api(self) -> None:
instrumentation = PikaInstrumentor()
instrumentation.instrument()
self.assertTrue(
isinstance(BlockingConnection.channel, BoundFunctionWrapper)
)
assert hasattr(
instrumentation, "__opentelemetry_tracer_provider"
), "Tracer not stored for the object!"
instrumentation.uninstrument(channel=self.channel)
self.assertFalse(
isinstance(BlockingConnection.channel, BoundFunctionWrapper)
)
def test_uninstrument_channel_functions(self) -> None:
original_function = self.channel.basic_publish
self.channel.basic_publish = mock.MagicMock()
self.channel.basic_publish._original_function = original_function
PikaInstrumentor._uninstrument_channel_functions(self.channel)
self.assertEqual(self.channel.basic_publish, original_function)