|
14 | 14 | import asyncio
|
15 | 15 | from typing import Type
|
16 | 16 | from unittest import TestCase, mock, skipIf
|
| 17 | +from unittest.mock import MagicMock |
17 | 18 |
|
18 | 19 | from aio_pika import Exchange, RobustExchange
|
19 | 20 |
|
@@ -92,6 +93,36 @@ def test_publish(self):
|
92 | 93 | def test_robust_publish(self):
|
93 | 94 | self._test_publish(RobustExchange)
|
94 | 95 |
|
| 96 | + def _test_publish_works_with_not_recording_span(self, exchange_type): |
| 97 | + exchange = exchange_type(CONNECTION_7, CHANNEL_7, EXCHANGE_NAME) |
| 98 | + with mock.patch.object( |
| 99 | + PublishDecorator, "_get_publish_span" |
| 100 | + ) as mock_get_publish_span: |
| 101 | + mocked_not_recording_span = MagicMock() |
| 102 | + mocked_not_recording_span.is_recording.return_value = False |
| 103 | + mock_get_publish_span.return_value = mocked_not_recording_span |
| 104 | + with mock.patch.object( |
| 105 | + Exchange, "publish", return_value=asyncio.sleep(0) |
| 106 | + ) as mock_publish: |
| 107 | + with mock.patch( |
| 108 | + "opentelemetry.instrumentation.aio_pika.publish_decorator.propagate.inject" |
| 109 | + ) as mock_inject: |
| 110 | + decorated_publish = PublishDecorator( |
| 111 | + self.tracer, exchange |
| 112 | + ).decorate(mock_publish) |
| 113 | + self.loop.run_until_complete( |
| 114 | + decorated_publish(MESSAGE, ROUTING_KEY) |
| 115 | + ) |
| 116 | + mock_publish.assert_called_once() |
| 117 | + mock_get_publish_span.assert_called_once() |
| 118 | + mock_inject.assert_called_once() |
| 119 | + |
| 120 | + def test_publish_works_with_not_recording_span(self): |
| 121 | + self._test_publish_works_with_not_recording_span(Exchange) |
| 122 | + |
| 123 | + def test_publish_works_with_not_recording_span_robust(self): |
| 124 | + self._test_publish_works_with_not_recording_span(RobustExchange) |
| 125 | + |
95 | 126 |
|
96 | 127 | @skipIf(AIOPIKA_VERSION_INFO <= (8, 0), "Only for aio_pika 8")
|
97 | 128 | class TestInstrumentedExchangeAioRmq8(TestCase):
|
@@ -144,3 +175,33 @@ def test_publish(self):
|
144 | 175 |
|
145 | 176 | def test_robust_publish(self):
|
146 | 177 | self._test_publish(RobustExchange)
|
| 178 | + |
| 179 | + def _test_publish_works_with_not_recording_span(self, exchange_type): |
| 180 | + exchange = exchange_type(CONNECTION_7, CHANNEL_7, EXCHANGE_NAME) |
| 181 | + with mock.patch.object( |
| 182 | + PublishDecorator, "_get_publish_span" |
| 183 | + ) as mock_get_publish_span: |
| 184 | + mocked_not_recording_span = MagicMock() |
| 185 | + mocked_not_recording_span.is_recording.return_value = False |
| 186 | + mock_get_publish_span.return_value = mocked_not_recording_span |
| 187 | + with mock.patch.object( |
| 188 | + Exchange, "publish", return_value=asyncio.sleep(0) |
| 189 | + ) as mock_publish: |
| 190 | + with mock.patch( |
| 191 | + "opentelemetry.instrumentation.aio_pika.publish_decorator.propagate.inject" |
| 192 | + ) as mock_inject: |
| 193 | + decorated_publish = PublishDecorator( |
| 194 | + self.tracer, exchange |
| 195 | + ).decorate(mock_publish) |
| 196 | + self.loop.run_until_complete( |
| 197 | + decorated_publish(MESSAGE, ROUTING_KEY) |
| 198 | + ) |
| 199 | + mock_publish.assert_called_once() |
| 200 | + mock_get_publish_span.assert_called_once() |
| 201 | + mock_inject.assert_called_once() |
| 202 | + |
| 203 | + def test_publish_works_with_not_recording_span(self): |
| 204 | + self._test_publish_works_with_not_recording_span(Exchange) |
| 205 | + |
| 206 | + def test_publish_works_with_not_recording_span_robust(self): |
| 207 | + self._test_publish_works_with_not_recording_span(RobustExchange) |
0 commit comments