Coverage for tests / brokers / nats / test_consume.py: 99%

1import asyncio (empty)

2from unittest.mock import MagicMock, patch (empty)

3 

4import pytest (empty)

5from nats.aio.msg import Msg (empty)

6 

7from faststream import AckPolicy (empty)

8from faststream.exceptions import AckMessage (empty)

9from faststream.nats import ConsumerConfig, JStream, PubAck, PullSub (empty)

10from faststream.nats.annotations import NatsMessage (empty)

11from faststream.nats.message import NatsMessage as StreamMessage (empty)

12from tests.brokers.base.consume import BrokerRealConsumeTestcase (empty)

13from tests.tools import spy_decorator (empty)

14 

15from .basic import NatsTestcaseConfig (empty)

16 

17 

18@pytest.mark.connected() (empty)

19@pytest.mark.nats() (empty)

20class TestConsume(NatsTestcaseConfig, BrokerRealConsumeTestcase): (empty)

21 async def test_concurrent_subscriber( (empty)

22 self, 

23 queue: str, 

24 mock: MagicMock, 

25 ) -> None: 

26 event = asyncio.Event() (empty)

27 event2 = asyncio.Event() (empty)

28 

29 broker = self.get_broker() (empty)

30 

31 args, kwargs = self.get_subscriber_params(queue, max_workers=2) (empty)

32 

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

34 async def handler(msg): (empty)

35 mock() (empty)

36 

37 if event.is_set(): (empty)

38 event2.set() (empty)

39 else: 

40 event.set() (empty)

41 

42 await asyncio.sleep(1.0) (empty)

43 

44 async with self.patch_broker(broker) as br: (empty)

45 await br.start() (empty)

46 

47 for i in range(5): (empty)

48 await br.publish(i, queue) (empty)

49 

50 await asyncio.wait( (empty)

51 ( 

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

53 asyncio.create_task(event2.wait()), 

54 ), 

55 timeout=3, 

56 ) 

57 

58 assert event.is_set() (empty)

59 assert event2.is_set() (empty)

60 assert mock.call_count == 2, mock.call_count (empty)

61 

62 async def test_consume_js( (empty)

63 self, 

64 queue: str, 

65 stream: JStream, 

66 ) -> None: 

67 event = asyncio.Event() (empty)

68 

69 consume_broker = self.get_broker() (empty)

70 

71 args, kwargs = self.get_subscriber_params(queue, stream=stream) (empty)

72 

73 @consume_broker.subscriber(*args, **kwargs) (empty)

74 def subscriber(m) -> None: (empty)

75 event.set() (empty)

76 

77 async with self.patch_broker(consume_broker) as br: (empty)

78 await br.start() (empty)

79 

80 result = await br.publish("hello", queue, stream=stream.name) (empty)

81 

82 await asyncio.wait( (empty)

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

84 timeout=3, 

85 ) 

86 

87 assert isinstance(result, PubAck), result (empty)

88 assert event.is_set() (empty)

89 

90 async def test_consume_with_filter( (empty)

91 self, 

92 queue: str, 

93 mock: MagicMock, 

94 ) -> None: 

95 event = asyncio.Event() (empty)

96 

97 consume_broker = self.get_broker() (empty)

98 

99 @consume_broker.subscriber( (empty)

100 config=ConsumerConfig(filter_subjects=[f"{queue}.a"]), 

101 stream=JStream(queue, subjects=[f"{queue}.*"]), 

102 ) 

103 def subscriber(m) -> None: (empty)

104 mock(m) (empty)

105 event.set() (empty)

106 

107 async with self.patch_broker(consume_broker) as br: (empty)

108 await br.start() (empty)

109 await asyncio.wait( (empty)

110 ( 

111 asyncio.create_task(br.publish(2, f"{queue}.a")), 

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

113 ), 

114 timeout=3, 

115 ) 

116 

117 assert event.is_set() (empty)

118 mock.assert_called_once_with(2) (empty)

119 

120 async def test_consume_pull( (empty)

121 self, 

122 queue: str, 

123 stream: JStream, 

124 mock, 

125 ) -> None: 

126 event = asyncio.Event() (empty)

127 

128 consume_broker = self.get_broker() (empty)

129 

130 @consume_broker.subscriber( (empty)

131 queue, 

132 stream=stream, 

133 pull_sub=PullSub(1), 

134 ) 

135 def subscriber(m) -> None: (empty)

136 mock(m) (empty)

137 event.set() (empty)

138 

139 async with self.patch_broker(consume_broker) as br: (empty)

140 await br.start() (empty)

141 

142 await asyncio.wait( (empty)

143 ( 

144 asyncio.create_task(br.publish("hello", queue)), 

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

146 ), 

147 timeout=3, 

148 ) 

149 

150 assert event.is_set() (empty)

151 mock.assert_called_once_with("hello") (empty)

152 

153 async def test_consume_batch( (empty)

154 self, 

155 queue: str, 

156 stream: JStream, 

157 mock, 

158 ) -> None: 

159 event = asyncio.Event() (empty)

160 

161 consume_broker = self.get_broker() (empty)

162 

163 @consume_broker.subscriber( (empty)

164 queue, 

165 stream=stream, 

166 pull_sub=PullSub(1, batch=True), 

167 ) 

168 def subscriber(m) -> None: (empty)

169 mock(m) (empty)

170 event.set() (empty)

171 

172 async with self.patch_broker(consume_broker) as br: (empty)

173 await br.start() (empty)

174 

175 await asyncio.wait( (empty)

176 ( 

177 asyncio.create_task(br.publish(b"hello", queue)), 

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

179 ), 

180 timeout=3, 

181 ) 

182 

183 assert event.is_set() (empty)

184 mock.assert_called_once_with([b"hello"]) (empty)

185 

186 async def test_core_consume_no_ack( (empty)

187 self, 

188 queue: str, 

189 mock: MagicMock, 

190 ) -> None: 

191 event = asyncio.Event() (empty)

192 

193 consume_broker = self.get_broker(apply_types=True) (empty)

194 

195 args, kwargs = self.get_subscriber_params( (empty)

196 queue, 

197 ack_policy=AckPolicy.MANUAL, 

198 ) 

199 

200 @consume_broker.subscriber(*args, **kwargs) (empty)

201 async def handler(msg: NatsMessage) -> None: (empty)

202 mock(msg.raw_message._ackd) (empty)

203 event.set() (empty)

204 

205 async with self.patch_broker(consume_broker) as br: (empty)

206 await br.start() (empty)

207 

208 # Check, that Core Subscriber doesn't call Acknowledgement automatically 

209 with patch.object( (empty)

210 StreamMessage, 

211 "ack", 

212 spy_decorator(StreamMessage.ack), 

213 ) as m: 

214 await asyncio.wait( (empty)

215 ( 

216 asyncio.create_task(br.publish("hello", queue)), 

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

218 ), 

219 timeout=3, 

220 ) 

221 assert not m.mock.called (empty)

222 

223 assert event.is_set() (empty)

224 mock.assert_called_once_with(True) # True was set by parser (empty)

225 

226 async def test_consume_ack( (empty)

227 self, 

228 queue: str, 

229 stream: JStream, 

230 ) -> None: 

231 event = asyncio.Event() (empty)

232 

233 consume_broker = self.get_broker(apply_types=True) (empty)

234 

235 @consume_broker.subscriber(queue, stream=stream) (empty)

236 async def handler(msg: NatsMessage) -> None: (empty)

237 event.set() (empty)

238 

239 async with self.patch_broker(consume_broker) as br: (empty)

240 await br.start() (empty)

241 

242 with patch.object(Msg, "ack", spy_decorator(Msg.ack)) as m: (empty)

243 await asyncio.wait( (empty)

244 ( 

245 asyncio.create_task(br.publish("hello", queue)), 

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

247 ), 

248 timeout=3, 

249 ) 

250 m.mock.assert_called_once() (empty)

251 

252 assert event.is_set() (empty)

253 

254 async def test_consume_ack_manual( (empty)

255 self, 

256 queue: str, 

257 stream: JStream, 

258 ) -> None: 

259 event = asyncio.Event() (empty)

260 

261 consume_broker = self.get_broker(apply_types=True) (empty)

262 

263 @consume_broker.subscriber(queue, stream=stream) (empty)

264 async def handler(msg: NatsMessage) -> None: (empty)

265 await msg.ack() (empty)

266 event.set() (empty)

267 

268 async with self.patch_broker(consume_broker) as br: (empty)

269 await br.start() (empty)

270 

271 with patch.object(Msg, "ack", spy_decorator(Msg.ack)) as m: (empty)

272 await asyncio.wait( (empty)

273 ( 

274 asyncio.create_task(br.publish("hello", queue)), 

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

276 ), 

277 timeout=3, 

278 ) 

279 m.mock.assert_called_once() (empty)

280 

281 assert event.is_set() (empty)

282 

283 async def test_consume_ack_sync_manual( (empty)

284 self, 

285 queue: str, 

286 event: asyncio.Event, 

287 stream: JStream, 

288 ): 

289 consume_broker = self.get_broker(apply_types=True) (empty)

290 

291 @consume_broker.subscriber(queue, stream=stream) (empty)

292 async def handler(msg: NatsMessage): (empty)

293 await msg.ack_sync() (empty)

294 event.set() (empty)

295 

296 async with self.patch_broker(consume_broker) as br: (empty)

297 await br.start() (empty)

298 

299 with patch.object(Msg, "ack_sync", spy_decorator(Msg.ack_sync)) as m: (empty)

300 await asyncio.wait( (empty)

301 ( 

302 asyncio.create_task(br.publish("hello", queue)), 

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

304 ), 

305 timeout=3, 

306 ) 

307 m.mock.assert_called_once() (empty)

308 

309 assert event.is_set() (empty)

310 

311 async def test_consume_ack_raise( (empty)

312 self, 

313 queue: str, 

314 stream: JStream, 

315 ) -> None: 

316 event = asyncio.Event() (empty)

317 

318 consume_broker = self.get_broker(apply_types=True) (empty)

319 

320 @consume_broker.subscriber(queue, stream=stream) (empty)

321 async def handler(msg: NatsMessage): (empty)

322 event.set() (empty)

323 raise AckMessage (empty)

324 

325 async with self.patch_broker(consume_broker) as br: (empty)

326 await br.start() (empty)

327 

328 with patch.object(Msg, "ack", spy_decorator(Msg.ack)) as m: (empty)

329 await asyncio.wait( (empty)

330 ( 

331 asyncio.create_task(br.publish("hello", queue)), 

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

333 ), 

334 timeout=3, 

335 ) 

336 m.mock.assert_called_once() (empty)

337 

338 assert event.is_set() (empty)

339 

340 async def test_nack( (empty)

341 self, 

342 queue: str, 

343 stream: JStream, 

344 ) -> None: 

345 event = asyncio.Event() (empty)

346 

347 consume_broker = self.get_broker(apply_types=True) (empty)

348 

349 @consume_broker.subscriber(queue, stream=stream) (empty)

350 async def handler(msg: NatsMessage) -> None: (empty)

351 await msg.nack() (empty)

352 event.set() (empty)

353 

354 async with self.patch_broker(consume_broker) as br: (empty)

355 await br.start() (empty)

356 

357 with patch.object(Msg, "nak", spy_decorator(Msg.nak)) as m: (empty)

358 await asyncio.wait( (empty)

359 ( 

360 asyncio.create_task(br.publish("hello", queue)), 

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

362 ), 

363 timeout=3, 

364 ) 

365 m.mock.assert_called_once() (empty)

366 

367 assert event.is_set() (empty)

368 

369 async def test_consume_no_ack( (empty)

370 self, 

371 queue: str, 

372 stream: str, 

373 ) -> None: 

374 event = asyncio.Event() (empty)

375 

376 consume_broker = self.get_broker(apply_types=True) (empty)

377 

378 @consume_broker.subscriber( (empty)

379 queue, 

380 stream=stream, 

381 ack_policy=AckPolicy.MANUAL, 

382 ) 

383 async def handler(msg: NatsMessage) -> None: (empty)

384 event.set() (empty)

385 

386 async with self.patch_broker(consume_broker) as br: (empty)

387 await br.start() (empty)

388 

389 with patch.object(Msg, "ack", spy_decorator(Msg.ack)) as m: (empty)

390 await asyncio.wait( (empty)

391 ( 

392 asyncio.create_task(br.publish("hello", queue)), 

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

394 ), 

395 timeout=3, 

396 ) 

397 m.mock.assert_not_called() (empty)

398 

399 assert event.is_set() (empty)

400 

401 async def test_consume_batch_headers( (empty)

402 self, 

403 queue: str, 

404 stream: JStream, 

405 mock, 

406 ) -> None: 

407 event = asyncio.Event() (empty)

408 

409 consume_broker = self.get_broker(apply_types=True) (empty)

410 

411 @consume_broker.subscriber( (empty)

412 queue, 

413 stream=stream, 

414 pull_sub=PullSub(1, batch=True), 

415 ) 

416 def subscriber(m, msg: NatsMessage) -> None: (empty)

417 check = all( (empty)

418 ( 

419 msg.headers, 

420 [msg.headers] == msg.batch_headers, 

421 msg.headers.get("custom") == "1", 

422 ), 

423 ) 

424 mock(check) (empty)

425 event.set() (empty)

426 

427 async with self.patch_broker(consume_broker) as br: (empty)

428 await br.start() (empty)

429 await asyncio.wait( (empty)

430 ( 

431 asyncio.create_task(br.publish("", queue, headers={"custom": "1"})), 

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

433 ), 

434 timeout=3, 

435 ) 

436 

437 assert event.is_set() (empty)

438 mock.assert_called_once_with(True) (empty)

439 

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

441 async def test_consume_kv( (empty)

442 self, 

443 queue: str, 

444 mock, 

445 ) -> None: 

446 event = asyncio.Event() (empty)

447 

448 consume_broker = self.get_broker(apply_types=True) (empty)

449 

450 @consume_broker.subscriber(queue, kv_watch=queue + "1") (empty)

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

452 mock(m) (empty)

453 event.set() (empty)

454 

455 async with self.patch_broker(consume_broker) as br: (empty)

456 await br.start() (empty)

457 bucket = await br.key_value(queue + "1") (empty)

458 

459 await asyncio.wait( (empty)

460 ( 

461 asyncio.create_task( 

462 bucket.put( 

463 queue, 

464 b"world", 

465 ), 

466 ), 

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

468 ), 

469 timeout=3, 

470 ) 

471 

472 assert event.is_set() (empty)

473 mock.assert_called_with(b"world") (empty)

474 

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

476 async def test_consume_os( (empty)

477 self, 

478 queue: str, 

479 mock, 

480 ) -> None: 

481 event = asyncio.Event() (empty)

482 

483 consume_broker = self.get_broker(apply_types=True) (empty)

484 

485 @consume_broker.subscriber(queue, obj_watch=True) (empty)

486 async def handler(filename: str) -> None: (empty)

487 event.set() (empty)

488 mock(filename) (empty)

489 

490 async with self.patch_broker(consume_broker) as br: (empty)

491 await br.start() (empty)

492 bucket = await br.object_storage(queue) (empty)

493 

494 await asyncio.wait( (empty)

495 ( 

496 asyncio.create_task( 

497 bucket.put( 

498 "hello", 

499 b"world", 

500 ), 

501 ), 

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

503 ), 

504 timeout=3, 

505 ) 

506 

507 assert event.is_set() (empty)

508 mock.assert_called_once_with("hello") (empty)

509 

510 async def test_get_one_js( (empty)

511 self, 

512 queue: str, 

513 stream: JStream, 

514 ) -> None: 

515 broker = self.get_broker(apply_types=True) (empty)

516 subscriber = broker.subscriber(queue, stream=stream) (empty)

517 

518 async with self.patch_broker(broker) as br: (empty)

519 await br.start() (empty)

520 

521 message = None (empty)

522 

523 async def consume() -> None: (empty)

524 nonlocal message 

525 message = await subscriber.get_one(timeout=5) (empty)

526 

527 async def publish() -> None: (empty)

528 await br.publish("test_message", queue, stream=stream.name) (empty)

529 

530 await asyncio.wait( (empty)

531 ( 

532 asyncio.create_task(consume()), 

533 asyncio.create_task(publish()), 

534 ), 

535 timeout=10, 

536 ) 

537 

538 assert message is not None (empty)

539 assert await message.decode() == "test_message" (empty)

540 

541 async def test_get_one_timeout_js( (empty)

542 self, 

543 queue: str, 

544 stream: JStream, 

545 mock, 

546 ) -> None: 

547 broker = self.get_broker(apply_types=True) (empty)

548 subscriber = broker.subscriber(queue, stream=stream) (empty)

549 

550 async with self.patch_broker(broker) as br: (empty)

551 await br.start() (empty)

552 

553 mock(await subscriber.get_one(timeout=1e-24)) (empty)

554 mock.assert_called_once_with(None) (empty)

555 

556 async def test_get_one_pull( (empty)

557 self, 

558 queue: str, 

559 stream: JStream, 

560 ) -> None: 

561 broker = self.get_broker(apply_types=True) (empty)

562 subscriber = broker.subscriber( (empty)

563 queue, 

564 stream=stream, 

565 pull_sub=PullSub(1), 

566 ) 

567 

568 async with self.patch_broker(broker) as br: (empty)

569 await br.start() (empty)

570 

571 message = None (empty)

572 

573 async def consume() -> None: (empty)

574 nonlocal message 

575 message = await subscriber.get_one(timeout=5) (empty)

576 

577 async def publish() -> None: (empty)

578 await br.publish("test_message", queue) (empty)

579 

580 await asyncio.wait( (empty)

581 ( 

582 asyncio.create_task(consume()), 

583 asyncio.create_task(publish()), 

584 ), 

585 timeout=10, 

586 ) 

587 

588 assert message is not None (empty)

589 assert await message.decode() == "test_message" (empty)

590 

591 async def test_get_one_pull_timeout( (empty)

592 self, 

593 queue: str, 

594 stream: JStream, 

595 mock: MagicMock, 

596 ) -> None: 

597 broker = self.get_broker(apply_types=True) (empty)

598 subscriber = broker.subscriber( (empty)

599 queue, 

600 stream=stream, 

601 pull_sub=PullSub(1), 

602 ) 

603 

604 async with self.patch_broker(broker) as br: (empty)

605 await br.start() (empty)

606 

607 mock(await subscriber.get_one(timeout=1e-24)) (empty)

608 mock.assert_called_once_with(None) (empty)

609 

610 async def test_get_one_batch( (empty)

611 self, 

612 queue: str, 

613 stream: JStream, 

614 ) -> None: 

615 broker = self.get_broker(apply_types=True) (empty)

616 subscriber = broker.subscriber( (empty)

617 queue, 

618 stream=stream, 

619 pull_sub=PullSub(1, batch=True), 

620 ) 

621 

622 async with self.patch_broker(broker) as br: (empty)

623 await br.start() (empty)

624 

625 message = None (empty)

626 

627 async def consume() -> None: (empty)

628 nonlocal message 

629 message = await subscriber.get_one(timeout=5) (empty)

630 

631 async def publish() -> None: (empty)

632 await br.publish("test_message", queue) (empty)

633 

634 await asyncio.wait( (empty)

635 ( 

636 asyncio.create_task(consume()), 

637 asyncio.create_task(publish()), 

638 ), 

639 timeout=10, 

640 ) 

641 

642 assert message is not None (empty)

643 assert await message.decode() == ["test_message"] (empty)

644 

645 async def test_get_one_batch_timeout( (empty)

646 self, 

647 queue: str, 

648 stream: JStream, 

649 mock: MagicMock, 

650 ) -> None: 

651 broker = self.get_broker(apply_types=True) (empty)

652 subscriber = broker.subscriber( (empty)

653 queue, 

654 stream=stream, 

655 pull_sub=PullSub(1, batch=True), 

656 ) 

657 

658 async with self.patch_broker(broker) as br: (empty)

659 await br.start() (empty)

660 

661 mock(await subscriber.get_one(timeout=1e-24)) (empty)

662 mock.assert_called_once_with(None) (empty)

663 

664 async def test_get_one_with_filter( (empty)

665 self, 

666 queue: str, 

667 stream: JStream, 

668 ) -> None: 

669 broker = self.get_broker(apply_types=True) (empty)

670 subscriber = broker.subscriber( (empty)

671 config=ConsumerConfig(filter_subjects=[f"{queue}.a"]), 

672 stream=JStream(queue, subjects=[f"{queue}.*"]), 

673 ) 

674 

675 async with self.patch_broker(broker) as br: (empty)

676 await br.start() (empty)

677 

678 message = None (empty)

679 

680 async def consume() -> None: (empty)

681 nonlocal message 

682 message = await subscriber.get_one(timeout=5) (empty)

683 

684 async def publish() -> None: (empty)

685 await br.publish("test_message", f"{queue}.a") (empty)

686 

687 await asyncio.wait( (empty)

688 ( 

689 asyncio.create_task(publish()), 

690 asyncio.create_task(consume()), 

691 ), 

692 timeout=10, 

693 ) 

694 

695 assert message is not None (empty)

696 assert await message.decode() == "test_message" (empty)

697 

698 async def test_get_one_kv( (empty)

699 self, 

700 queue: str, 

701 stream: JStream, 

702 ) -> None: 

703 broker = self.get_broker(apply_types=True) (empty)

704 subscriber = broker.subscriber(queue, kv_watch=queue + "1") (empty)

705 

706 async with self.patch_broker(broker) as br: (empty)

707 await br.start() (empty)

708 bucket = await br.key_value(queue + "1") (empty)

709 

710 message = None (empty)

711 

712 async def consume() -> None: (empty)

713 nonlocal message 

714 message = await subscriber.get_one(timeout=5) (empty)

715 

716 async def publish() -> None: (empty)

717 await bucket.put(queue, b"test_message") (empty)

718 

719 await asyncio.wait( (empty)

720 ( 

721 asyncio.create_task(consume()), 

722 asyncio.create_task(publish()), 

723 ), 

724 timeout=10, 

725 ) 

726 

727 assert message is not None (empty)

728 assert await message.decode() == b"test_message" (empty)

729 

730 async def test_get_one_kv_timeout( (empty)

731 self, 

732 queue: str, 

733 stream: JStream, 

734 mock: MagicMock, 

735 ) -> None: 

736 broker = self.get_broker(apply_types=True) (empty)

737 subscriber = broker.subscriber(queue, kv_watch=queue + "1") (empty)

738 

739 async with self.patch_broker(broker) as br: (empty)

740 await br.start() (empty)

741 

742 mock(await subscriber.get_one(timeout=1e-24)) (empty)

743 mock.assert_called_once_with(None) (empty)

744 

745 async def test_get_one_os( (empty)

746 self, 

747 queue: str, 

748 stream: JStream, 

749 ) -> None: 

750 broker = self.get_broker(apply_types=True) (empty)

751 subscriber = broker.subscriber(queue, obj_watch=True) (empty)

752 

753 async with self.patch_broker(broker) as br: (empty)

754 await br.start() (empty)

755 bucket = await br.object_storage(queue) (empty)

756 

757 new_object_id = None (empty)

758 

759 async def consume() -> None: (empty)

760 nonlocal new_object_id 

761 new_object_event = await subscriber.get_one(timeout=5) (empty)

762 new_object_id = await new_object_event.decode() (empty)

763 

764 async def publish() -> None: (empty)

765 await bucket.put(queue, b"test_message") (empty)

766 

767 await asyncio.wait( (empty)

768 ( 

769 asyncio.create_task(consume()), 

770 asyncio.create_task(publish()), 

771 ), 

772 timeout=10, 

773 ) 

774 

775 new_object = await bucket.get(new_object_id) (empty)

776 assert new_object.data == b"test_message" (empty)

777 

778 async def test_get_one_os_timeout( (empty)

779 self, 

780 queue: str, 

781 stream: JStream,