Coverage for tests / opentelemetry / basic.py: 99%

1import asyncio (empty)

2from typing import Any, cast (empty)

3from unittest.mock import MagicMock (empty)

4 

5import pytest (empty)

6from dirty_equals import IsFloat, IsUUID (empty)

7from opentelemetry import baggage, context (empty)

8from opentelemetry.baggage.propagation import W3CBaggagePropagator (empty)

9from opentelemetry.sdk.metrics import MeterProvider (empty)

10from opentelemetry.sdk.metrics._internal.point import Metric (empty)

11from opentelemetry.sdk.metrics.export import InMemoryMetricReader (empty)

12from opentelemetry.sdk.resources import Resource (empty)

13from opentelemetry.sdk.trace import Span, TracerProvider (empty)

14from opentelemetry.sdk.trace.export import SimpleSpanProcessor (empty)

15from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter (empty)

16from opentelemetry.semconv.trace import SpanAttributes as SpanAttr (empty)

17from opentelemetry.trace import SpanKind, get_current_span (empty)

18 

19from faststream._internal.broker import BrokerUsecase (empty)

20from faststream.opentelemetry import Baggage, CurrentBaggage, CurrentSpan (empty)

21from faststream.opentelemetry.consts import ( (empty)

22 ERROR_TYPE, 

23 MESSAGING_DESTINATION_PUBLISH_NAME, 

24) 

25from faststream.opentelemetry.middleware import ( (empty)

26 MessageAction as Action, 

27 TelemetryMiddleware, 

28) 

29from tests.brokers.base.basic import BaseTestcaseConfig (empty)

30 

31 

32@pytest.mark.asyncio() (empty)

33class LocalTelemetryTestcase(BaseTestcaseConfig): (empty)

34 messaging_system: str (empty)

35 include_messages_counters: bool (empty)

36 resource: Resource = Resource.create(attributes={"service.name": "faststream.test"}) (empty)

37 telemetry_middleware_class: TelemetryMiddleware (empty)

38 

39 def get_broker( (empty)

40 self, 

41 apply_types: bool = False, 

42 **kwargs: Any, 

43 ) -> BrokerUsecase[Any, Any]: 

44 raise NotImplementedError 

45 

46 def patch_broker( (empty)

47 self, 

48 broker: BrokerUsecase[Any, Any], 

49 **kwargs: Any, 

50 ) -> BrokerUsecase[Any, Any]: 

51 return broker 

52 

53 def destination_name(self, queue: str) -> str: (empty)

54 return queue (empty)

55 

56 @staticmethod (empty)

57 def get_spans(exporter: InMemorySpanExporter) -> list[Span]: (empty)

58 spans = cast("tuple[Span, ...]", exporter.get_finished_spans()) (empty)

59 return sorted(spans, key=lambda s: s.start_time or 0) (empty)

60 

61 @staticmethod (empty)

62 def get_metrics( (empty)

63 reader: InMemoryMetricReader, 

64 ) -> list[Metric]: 

65 """Get sorted metrics. 

66 

67 Return order: 

68 - messaging.process.duration 

69 - messaging.process.messages 

70 - messaging.publish.duration 

71 - messaging.publish.messages 

72 """ 

73 metrics = reader.get_metrics_data() (empty)

74 metrics = metrics.resource_metrics[0].scope_metrics[0].metrics (empty)

75 metrics = sorted(metrics, key=lambda m: m.name) (empty)

76 return cast("list[Metric]", metrics) (empty)

77 

78 @pytest.fixture() (empty)

79 def tracer_provider(self) -> TracerProvider: (empty)

80 return TracerProvider(resource=self.resource) (empty)

81 

82 @pytest.fixture() (empty)

83 def trace_exporter(self, tracer_provider: TracerProvider) -> InMemorySpanExporter: (empty)

84 exporter = InMemorySpanExporter() (empty)

85 tracer_provider.add_span_processor(SimpleSpanProcessor(exporter)) (empty)

86 return exporter (empty)

87 

88 @pytest.fixture() (empty)

89 def metric_reader(self) -> InMemoryMetricReader: (empty)

90 return InMemoryMetricReader() (empty)

91 

92 @pytest.fixture() (empty)

93 def meter_provider(self, metric_reader: InMemoryMetricReader) -> MeterProvider: (empty)

94 return MeterProvider(metric_readers=(metric_reader,), resource=self.resource) (empty)

95 

96 def assert_span( (empty)

97 self, 

98 span: Span, 

99 action: str, 

100 queue: str, 

101 msg: str, 

102 parent_span_id: str | None = None, 

103 ) -> None: 

104 attrs = span.attributes or {} (empty)

105 assert attrs[SpanAttr.MESSAGING_SYSTEM] == self.messaging_system, attrs[ (empty)

106 SpanAttr.MESSAGING_SYSTEM 

107 ] 

108 assert attrs[SpanAttr.MESSAGING_MESSAGE_CONVERSATION_ID] == IsUUID, attrs[ (empty)

109 SpanAttr.MESSAGING_MESSAGE_CONVERSATION_ID 

110 ] 

111 assert span.name == f"{self.destination_name(queue)} {action}", span.name (empty)

112 assert span.kind in {SpanKind.CONSUMER, SpanKind.PRODUCER}, span.kind (empty)

113 

114 if span.kind == SpanKind.PRODUCER and action in {Action.CREATE, Action.PUBLISH}: (empty)

115 assert attrs[SpanAttr.MESSAGING_DESTINATION_NAME] == queue, attrs[ (empty)

116 SpanAttr.MESSAGING_DESTINATION_NAME 

117 ] 

118 

119 if span.kind == SpanKind.CONSUMER and action in {Action.CREATE, Action.PROCESS}: (empty)

120 assert attrs[MESSAGING_DESTINATION_PUBLISH_NAME] == queue, attrs[ (empty)

121 MESSAGING_DESTINATION_PUBLISH_NAME 

122 ] 

123 assert attrs[SpanAttr.MESSAGING_MESSAGE_ID] == IsUUID, attrs[ (empty)

124 SpanAttr.MESSAGING_MESSAGE_ID 

125 ] 

126 

127 if action == Action.PROCESS: (empty)

128 assert attrs[SpanAttr.MESSAGING_MESSAGE_PAYLOAD_SIZE_BYTES] == len(msg), ( (empty)

129 attrs[SpanAttr.MESSAGING_MESSAGE_PAYLOAD_SIZE_BYTES] 

130 ) 

131 assert attrs[SpanAttr.MESSAGING_OPERATION] == action, attrs[ (empty)

132 SpanAttr.MESSAGING_OPERATION 

133 ] 

134 

135 if action == Action.PUBLISH: (empty)

136 assert attrs[SpanAttr.MESSAGING_OPERATION] == action, attrs[ (empty)

137 SpanAttr.MESSAGING_OPERATION 

138 ] 

139 

140 if parent_span_id: (empty)

141 assert span.parent (empty)

142 assert span.parent.span_id == parent_span_id, span.parent.span_id (empty)

143 

144 def assert_metrics( (empty)

145 self, 

146 metrics: list[Metric], 

147 count: int = 1, 

148 error_type: str | None = None, 

149 ) -> None: 

150 if self.include_messages_counters: (empty)

151 assert len(metrics) == 4 (empty)

152 proc_dur, proc_msg, pub_dur, pub_msg = metrics (empty)

153 

154 assert proc_msg.data.data_points[0].value == count (empty)

155 assert pub_msg.data.data_points[0].value == count (empty)

156 

157 else: 

158 assert len(metrics) == 2 (empty)

159 proc_dur, pub_dur = metrics (empty)

160 

161 if error_type: (empty)

162 assert proc_dur.data.data_points[0].attributes[ERROR_TYPE] == error_type (empty)

163 

164 assert proc_dur.data.data_points[0].count == 1 (empty)

165 assert proc_dur.data.data_points[0].sum == IsFloat (empty)

166 

167 assert pub_dur.data.data_points[0].count == 1 (empty)

168 assert pub_dur.data.data_points[0].sum == IsFloat (empty)

169 

170 async def test_subscriber_create_publish_process_span( (empty)

171 self, 

172 queue: str, 

173 mock: MagicMock, 

174 tracer_provider: TracerProvider, 

175 trace_exporter: InMemorySpanExporter, 

176 event: asyncio.Event, 

177 ) -> None: 

178 mid = self.telemetry_middleware_class(tracer_provider=tracer_provider) (empty)

179 broker = self.get_broker(middlewares=(mid,)) (empty)

180 

181 args, kwargs = self.get_subscriber_params(queue) (empty)

182 

183 @broker.subscriber(*args, **kwargs) (empty)

184 async def handler(m) -> None: (empty)

185 mock(m) (empty)

186 event.set() (empty)

187 

188 broker = self.patch_broker(broker) (empty)

189 msg = "start" (empty)

190 

191 async with broker: (empty)

192 await broker.start() (empty)

193 tasks = ( (empty)

194 asyncio.create_task(broker.publish(msg, queue)), 

195 asyncio.create_task(event.wait()), 

196 ) 

197 await asyncio.wait(tasks, timeout=self.timeout) (empty)

198 

199 create, publish, process = self.get_spans(trace_exporter) (empty)

200 parent_span_id = create.context.span_id (empty)

201 

202 self.assert_span(create, Action.CREATE, queue, msg) (empty)

203 self.assert_span(publish, Action.PUBLISH, queue, msg, parent_span_id) (empty)

204 self.assert_span(process, Action.PROCESS, queue, msg, parent_span_id) (empty)

205 

206 mock.assert_called_once_with(msg) (empty)

207 

208 async def test_chain_subscriber_publisher( (empty)

209 self, 

210 queue: str, 

211 mock: MagicMock, 

212 tracer_provider: TracerProvider, 

213 trace_exporter: InMemorySpanExporter, 

214 event: asyncio.Event, 

215 ) -> None: 

216 mid = self.telemetry_middleware_class(tracer_provider=tracer_provider) (empty)

217 broker = self.get_broker(middlewares=(mid,)) (empty)

218 

219 first_queue = queue (empty)

220 second_queue = queue + "2" (empty)

221 

222 args, kwargs = self.get_subscriber_params(first_queue) (empty)

223 

224 @broker.subscriber(*args, **kwargs) (empty)

225 @broker.publisher(second_queue) (empty)

226 async def handler1(m): (empty)

227 return m (empty)

228 

229 args2, kwargs2 = self.get_subscriber_params(second_queue) (empty)

230 

231 @broker.subscriber(*args2, **kwargs2) (empty)

232 async def handler2(m) -> None: (empty)

233 mock(m) (empty)

234 event.set() (empty)

235 

236 broker = self.patch_broker(broker) (empty)

237 msg = "start" (empty)

238 

239 async with broker: (empty)

240 await broker.start() (empty)

241 tasks = ( (empty)

242 asyncio.create_task(broker.publish(msg, queue)), 

243 asyncio.create_task(event.wait()), 

244 ) 

245 await asyncio.wait(tasks, timeout=self.timeout) (empty)

246 

247 spans = self.get_spans(trace_exporter) (empty)

248 create, pub1, proc1, pub2, proc2 = spans (empty)

249 parent_span_id = create.context.span_id (empty)

250 

251 self.assert_span(create, Action.CREATE, first_queue, msg) (empty)

252 self.assert_span(pub1, Action.PUBLISH, first_queue, msg, parent_span_id) (empty)

253 self.assert_span(proc1, Action.PROCESS, first_queue, msg, parent_span_id) (empty)

254 self.assert_span(pub2, Action.PUBLISH, second_queue, msg, proc1.context.span_id) (empty)

255 self.assert_span(proc2, Action.PROCESS, second_queue, msg, parent_span_id) (empty)

256 

257 assert ( 

258 create.start_time 

259 < pub1.start_time 

260 < proc1.start_time 

261 < pub2.start_time 

262 < proc2.start_time 

263 ) 

264 

265 mock.assert_called_once_with(msg) (empty)

266 

267 @pytest.mark.flaky(reruns=3, reruns_delay=1) (empty)

268 async def test_no_trace_context_create_process_span( (empty)

269 self, 

270 queue: str, 

271 tracer_provider: TracerProvider, 

272 trace_exporter: InMemorySpanExporter, 

273 event: asyncio.Event, 

274 ) -> None: 

275 mid = self.telemetry_middleware_class(tracer_provider=tracer_provider) (empty)

276 broker = self.get_broker(middlewares=(mid,)) (empty)

277 broker_without_middlewares = self.get_broker() (empty)

278 

279 args, kwargs = self.get_subscriber_params(queue) (empty)

280 

281 @broker.subscriber(*args, **kwargs) (empty)

282 async def handler(m) -> None: (empty)

283 event.set() (empty)

284 

285 broker = self.patch_broker(broker) (empty)

286 msg = "start" (empty)

287 

288 async with broker, broker_without_middlewares: (empty)

289 await broker.start() (empty)

290 await broker_without_middlewares.start() (empty)

291 

292 broker.config.broker_config.middlewares = () (empty)

293 tasks = ( (empty)

294 asyncio.create_task(broker_without_middlewares.publish(msg, queue)), 

295 asyncio.create_task(event.wait()), 

296 ) 

297 await asyncio.wait(tasks, timeout=self.timeout) (empty)

298 

299 create, process = self.get_spans(trace_exporter) (empty)

300 parent_span_id = create.context.span_id (empty)

301 

302 self.assert_span(create, Action.CREATE, queue, msg) (empty)

303 self.assert_span(process, Action.PROCESS, queue, msg, parent_span_id) (empty)

304 

305 async def test_metrics( (empty)

306 self, 

307 queue: str, 

308 mock: MagicMock, 

309 meter_provider: MeterProvider, 

310 metric_reader: InMemoryMetricReader, 

311 event: asyncio.Event, 

312 ) -> None: 

313 mid = self.telemetry_middleware_class(meter_provider=meter_provider) (empty)

314 broker = self.get_broker(middlewares=(mid,)) (empty)

315 

316 args, kwargs = self.get_subscriber_params(queue) (empty)

317 

318 @broker.subscriber(*args, **kwargs) (empty)

319 async def handler(m) -> None: (empty)

320 mock(m) (empty)

321 event.set() (empty)

322 

323 broker = self.patch_broker(broker) (empty)

324 msg = "start" (empty)

325 

326 async with broker: (empty)

327 await broker.start() (empty)

328 tasks = ( (empty)

329 asyncio.create_task(broker.publish(msg, queue)), 

330 asyncio.create_task(event.wait()), 

331 ) 

332 await asyncio.wait(tasks, timeout=self.timeout) (empty)

333 

334 metrics = self.get_metrics(metric_reader) (empty)

335 

336 self.assert_metrics(metrics) (empty)

337 mock.assert_called_once_with(msg) (empty)

338 

339 async def test_error_metrics( (empty)

340 self, 

341 queue: str, 

342 mock: MagicMock, 

343 meter_provider: MeterProvider, 

344 metric_reader: InMemoryMetricReader, 

345 event: asyncio.Event, 

346 ) -> None: 

347 mid = self.telemetry_middleware_class(meter_provider=meter_provider) (empty)

348 broker = self.get_broker(middlewares=(mid,)) (empty)

349 expected_value_type = "ValueError" (empty)

350 

351 args, kwargs = self.get_subscriber_params(queue) (empty)

352 

353 @broker.subscriber(*args, **kwargs) (empty)

354 async def handler(m) -> None: (empty)

355 try: (empty)

356 raise ValueError (empty)

357 finally: 

358 mock(m) (empty)

359 event.set() (empty)

360 

361 broker = self.patch_broker(broker) (empty)

362 msg = "start" (empty)

363 

364 async with broker: (empty)

365 await broker.start() (empty)

366 tasks = ( (empty)

367 asyncio.create_task(broker.publish(msg, queue)), 

368 asyncio.create_task(event.wait()), 

369 ) 

370 await asyncio.wait(tasks, timeout=self.timeout) (empty)

371 

372 metrics = self.get_metrics(metric_reader) (empty)

373 

374 self.assert_metrics(metrics, error_type=expected_value_type) (empty)

375 mock.assert_called_once_with(msg) (empty)

376 

377 async def test_span_in_context( (empty)

378 self, 

379 queue: str, 

380 mock: MagicMock, 

381 tracer_provider: TracerProvider, 

382 trace_exporter: InMemorySpanExporter, 

383 event: asyncio.Event, 

384 ) -> None: 

385 mid = self.telemetry_middleware_class(tracer_provider=tracer_provider) (empty)

386 broker = self.get_broker(middlewares=(mid,), apply_types=True) (empty)

387 

388 args, kwargs = self.get_subscriber_params(queue) (empty)

389 

390 @broker.subscriber(*args, **kwargs) (empty)

391 async def handler(m, span: CurrentSpan) -> None: (empty)

392 assert span is get_current_span() (empty)

393 mock(m) (empty)

394 event.set() (empty)

395 

396 broker = self.patch_broker(broker) (empty)

397 msg = "start" (empty)

398 

399 async with broker: (empty)

400 await broker.start() (empty)

401 tasks = ( (empty)

402 asyncio.create_task(broker.publish(msg, queue)), 

403 asyncio.create_task(event.wait()), 

404 ) 

405 await asyncio.wait(tasks, timeout=self.timeout) (empty)

406 

407 mock.assert_called_once_with(msg) (empty)

408 

409 async def test_get_baggage( (empty)

410 self, 

411 queue: str, 

412 mock: MagicMock, 

413 event: asyncio.Event, 

414 ) -> None: 

415 mid = self.telemetry_middleware_class() (empty)

416 broker = self.get_broker(middlewares=(mid,), apply_types=True) (empty)

417 expected_baggage = {"foo": "bar"} (empty)

418 

419 args, kwargs = self.get_subscriber_params(queue) (empty)

420 

421 @broker.subscriber(*args, **kwargs) (empty)

422 async def handler1(m, baggage: CurrentBaggage) -> None: (empty)

423 assert baggage.get("foo") == "bar" (empty)

424 assert baggage.get_all() == expected_baggage (empty)

425 assert baggage.get_all_batch() == [] (empty)

426 assert baggage.__repr__() == expected_baggage.__repr__() (empty)

427 mock(m) (empty)

428 event.set() (empty)

429 

430 broker = self.patch_broker(broker) (empty)

431 msg = "start" (empty)

432 

433 async with broker: (empty)

434 await broker.start() (empty)

435 tasks = ( (empty)

436 asyncio.create_task( 

437 broker.publish( 

438 msg, 

439 queue, 

440 headers=Baggage({"foo": "bar"}).to_headers(), 

441 ), 

442 ), 

443 asyncio.create_task(event.wait()), 

444 ) 

445 await asyncio.wait(tasks, timeout=self.timeout) (empty)

446 

447 mock.assert_called_once_with(msg) (empty)

448 

449 async def test_clear_baggage( (empty)

450 self, 

451 queue: str, 

452 mock: MagicMock, 

453 event: asyncio.Event, 

454 ) -> None: 

455 mid = self.telemetry_middleware_class() (empty)

456 broker = self.get_broker(middlewares=(mid,), apply_types=True) (empty)

457 

458 first_queue = queue + "1" (empty)

459 second_queue = queue + "2" (empty)

460 

461 args, kwargs = self.get_subscriber_params(first_queue) (empty)

462 

463 @broker.subscriber(*args, **kwargs) (empty)

464 @broker.publisher(second_queue) (empty)

465 async def handler1(m, baggage: CurrentBaggage): (empty)

466 baggage.clear() (empty)

467 assert baggage.get_all() == {} (empty)

468 return m (empty)

469 

470 args2, kwargs2 = self.get_subscriber_params(second_queue) (empty)

471 

472 @broker.subscriber(*args2, **kwargs2) (empty)

473 async def handler2(m, baggage: CurrentBaggage) -> None: (empty)

474 assert baggage.get_all() == {} (empty)

475 mock(m) (empty)

476 event.set() (empty)

477 

478 broker = self.patch_broker(broker) (empty)

479 msg = "start" (empty)

480 

481 async with broker: (empty)

482 await broker.start() (empty)

483 tasks = ( (empty)

484 asyncio.create_task( 

485 broker.publish( 

486 msg, 

487 first_queue, 

488 headers=Baggage({"foo": "bar"}).to_headers(), 

489 ), 

490 ), 

491 asyncio.create_task(event.wait()), 

492 ) 

493 await asyncio.wait(tasks, timeout=self.timeout) (empty)

494 

495 mock.assert_called_once_with(msg) (empty)

496 

497 async def test_modify_baggage( (empty)

498 self, 

499 queue: str, 

500 mock: MagicMock, 

501 event: asyncio.Event, 

502 ) -> None: 

503 mid = self.telemetry_middleware_class() (empty)

504 broker = self.get_broker(middlewares=(mid,), apply_types=True) (empty)

505 expected_baggage = {"baz": "bar", "bar": "baz"} (empty)

506 

507 first_queue = queue + "1" (empty)

508 second_queue = queue + "2" (empty)

509 

510 args, kwargs = self.get_subscriber_params(first_queue) (empty)

511 

512 @broker.subscriber(*args, **kwargs) (empty)

513 @broker.publisher(second_queue) (empty)

514 async def handler1(m, baggage: CurrentBaggage): (empty)

515 baggage.set("bar", "baz") (empty)

516 baggage.set("baz", "bar") (empty)

517 baggage.remove("foo") (empty)

518 return m (empty)

519 

520 args2, kwargs2 = self.get_subscriber_params(second_queue) (empty)

521 

522 @broker.subscriber(*args2, **kwargs2) (empty)

523 async def handler2(m, baggage: CurrentBaggage) -> None: (empty)

524 assert baggage.get_all() == expected_baggage (empty)

525 mock(m) (empty)

526 event.set() (empty)

527 

528 broker = self.patch_broker(broker) (empty)

529 msg = "start" (empty)

530 

531 async with broker: (empty)

532 await broker.start() (empty)

533 tasks = ( (empty)

534 asyncio.create_task( 

535 broker.publish( 

536 msg, 

537 first_queue, 

538 headers=Baggage({"foo": "bar"}).to_headers(), 

539 ), 

540 ), 

541 asyncio.create_task(event.wait()), 

542 ) 

543 await asyncio.wait(tasks, timeout=self.timeout) (empty)

544 

545 mock.assert_called_once_with(msg) (empty)

546 

547 async def test_get_baggage_from_headers( (empty)

548 self, 

549 queue: str, 

550 event: asyncio.Event, 

551 ) -> None: 

552 mid = self.telemetry_middleware_class() (empty)

553 broker = self.get_broker(middlewares=(mid,), apply_types=True) (empty)

554 

555 args, kwargs = self.get_subscriber_params(queue) (empty)

556 

557 expected_baggage = {"foo": "bar", "bar": "baz"} (empty)

558 

559 ctx = context.Context() (empty)

560 for key, value in expected_baggage.items(): (empty)

561 ctx = baggage.set_baggage(key, value, context=ctx) (empty)

562 

563 propagator = W3CBaggagePropagator() (empty)

564 headers = {} (empty)

565 propagator.inject(headers, context=ctx) (empty)

566 

567 @broker.subscriber(*args, **kwargs) (empty)

568 async def handler() -> None: (empty)

569 baggage_instance = Baggage.from_headers(headers) (empty)

570 extracted_baggage = baggage_instance.get_all() (empty)

571 assert extracted_baggage == expected_baggage (empty)

572 event.set() (empty)

573 

574 broker = self.patch_broker(broker) (empty)

575 msg = "start" (empty)

576 

577 async with broker: (empty)

578 await broker.start() (empty)

579 tasks = ( (empty)

580 asyncio.create_task(broker.publish(msg, queue, headers=headers)), 

581 asyncio.create_task(event.wait()), 

582 ) 

583 await asyncio.wait(tasks, timeout=self.timeout) (empty)

584 

585 assert event.is_set() (empty)