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

1import asyncio (empty)

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

3 

4import pytest (empty)

5from redis.asyncio import Redis (empty)

6 

7from faststream import AckPolicy (empty)

8from faststream.redis import ( (empty)

9 ListSub, 

10 PubSub, 

11 RedisMessage, 

12 RedisStreamMessage, 

13 StreamSub, 

14) 

15from faststream.redis.exceptions import StreamGroupNotFoundError (empty)

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

17from tests.tools import spy_decorator (empty)

18 

19from .basic import RedisTestcaseConfig (empty)

20 

21 

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

23@pytest.mark.redis() (empty)

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

25class TestConsume(RedisTestcaseConfig, BrokerRealConsumeTestcase): (empty)

26 async def test_consume_native( (empty)

27 self, 

28 mock: MagicMock, 

29 queue: str, 

30 ) -> None: 

31 event = asyncio.Event() (empty)

32 

33 consume_broker = self.get_broker() (empty)

34 

35 @consume_broker.subscriber(queue) (empty)

36 async def handler(msg) -> None: (empty)

37 mock(msg) (empty)

38 event.set() (empty)

39 

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

41 await br.start() (empty)

42 

43 result = await br._connection.publish(queue, "hello") (empty)

44 await asyncio.wait( (empty)

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

46 timeout=3, 

47 ) 

48 assert result == 1, result (empty)

49 

50 mock.assert_called_once_with(b"hello") (empty)

51 

52 async def test_pattern_with_path( (empty)

53 self, 

54 mock: MagicMock, 

55 ) -> None: 

56 event = asyncio.Event() (empty)

57 

58 consume_broker = self.get_broker() (empty)

59 

60 @consume_broker.subscriber("test.{name}") (empty)

61 async def handler(msg) -> None: (empty)

62 mock(msg) (empty)

63 event.set() (empty)

64 

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

66 await br.start() (empty)

67 

68 await asyncio.wait( (empty)

69 ( 

70 asyncio.create_task(br.publish("hello", "test.name")), 

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

72 ), 

73 timeout=3, 

74 ) 

75 

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

77 

78 async def test_pattern_without_path( (empty)

79 self, 

80 mock: MagicMock, 

81 ) -> None: 

82 event = asyncio.Event() (empty)

83 

84 consume_broker = self.get_broker() (empty)

85 

86 @consume_broker.subscriber(PubSub("test.*", pattern=True)) (empty)

87 async def handler(msg) -> None: (empty)

88 mock(msg) (empty)

89 event.set() (empty)

90 

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

92 await br.start() (empty)

93 

94 await asyncio.wait( (empty)

95 ( 

96 asyncio.create_task(br.publish("hello", "test.name")), 

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

98 ), 

99 timeout=3, 

100 ) 

101 

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

103 

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

105 async def test_concurrent_consume_channel( (empty)

106 self, 

107 queue: str, 

108 mock: MagicMock, 

109 ) -> None: 

110 event = asyncio.Event() (empty)

111 event2 = asyncio.Event() (empty)

112 

113 consume_broker = self.get_broker() (empty)

114 

115 @consume_broker.subscriber(channel=PubSub(queue), max_workers=2) (empty)

116 async def handler(msg): (empty)

117 mock() (empty)

118 if event.is_set(): (empty)

119 event2.set() (empty)

120 else: 

121 event.set() (empty)

122 await asyncio.sleep(0.1) (empty)

123 

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

125 await br.start() (empty)

126 

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

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

129 

130 await asyncio.wait( (empty)

131 ( 

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

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

134 ), 

135 timeout=3, 

136 ) 

137 

138 assert event.is_set() (empty)

139 assert event2.is_set() (empty)

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

141 

142 

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

144@pytest.mark.redis() (empty)

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

146class TestConsumeList(RedisTestcaseConfig): (empty)

147 async def test_consume_list( (empty)

148 self, 

149 queue: str, 

150 mock: MagicMock, 

151 ) -> None: 

152 event = asyncio.Event() (empty)

153 

154 consume_broker = self.get_broker() (empty)

155 

156 @consume_broker.subscriber(list=queue) (empty)

157 async def handler(msg) -> None: (empty)

158 mock(msg) (empty)

159 event.set() (empty)

160 

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

162 await br.start() (empty)

163 

164 await asyncio.wait( (empty)

165 ( 

166 asyncio.create_task(br.publish("hello", list=queue)), 

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

168 ), 

169 timeout=3, 

170 ) 

171 

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

173 

174 async def test_consume_list_native( (empty)

175 self, 

176 queue: str, 

177 mock: MagicMock, 

178 ) -> None: 

179 event = asyncio.Event() (empty)

180 

181 consume_broker = self.get_broker() (empty)

182 

183 @consume_broker.subscriber(list=queue) (empty)

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

185 mock(msg) (empty)

186 event.set() (empty)

187 

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

189 await br.start() (empty)

190 

191 await asyncio.wait( (empty)

192 ( 

193 asyncio.create_task(br._connection.rpush(queue, "hello")), 

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

195 ), 

196 timeout=3, 

197 ) 

198 

199 mock.assert_called_once_with(b"hello") (empty)

200 

201 @pytest.mark.slow() (empty)

202 async def test_consume_list_batch_with_one( (empty)

203 self, 

204 queue: str, 

205 mock: MagicMock, 

206 ) -> None: 

207 event = asyncio.Event() (empty)

208 

209 consume_broker = self.get_broker() (empty)

210 

211 @consume_broker.subscriber( (empty)

212 list=ListSub(queue, batch=True, polling_interval=0.01), 

213 ) 

214 async def handler(msg) -> None: (empty)

215 mock(msg) (empty)

216 event.set() (empty)

217 

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

219 await br.start() (empty)

220 await asyncio.wait( (empty)

221 ( 

222 asyncio.create_task(br.publish("hi", list=queue)), 

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

224 ), 

225 timeout=3, 

226 ) 

227 

228 assert event.is_set() (empty)

229 mock.assert_called_once_with(["hi"]) (empty)

230 

231 @pytest.mark.slow() (empty)

232 async def test_consume_list_batch_headers( (empty)

233 self, 

234 queue: str, 

235 mock: MagicMock, 

236 ) -> None: 

237 event = asyncio.Event() (empty)

238 

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

240 

241 @consume_broker.subscriber( (empty)

242 list=ListSub(queue, batch=True, polling_interval=0.01), 

243 ) 

244 def subscriber(m, msg: RedisMessage) -> None: (empty)

245 check = all( (empty)

246 ( 

247 msg.headers, 

248 msg.headers["correlation_id"] 

249 == msg.batch_headers[0]["correlation_id"], 

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

251 ), 

252 ) 

253 mock(check) (empty)

254 event.set() (empty)

255 

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

257 await br.start() (empty)

258 await asyncio.wait( (empty)

259 ( 

260 asyncio.create_task( 

261 br.publish("", list=queue, headers={"custom": "1"}), 

262 ), 

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

264 ), 

265 timeout=3, 

266 ) 

267 

268 assert event.is_set() (empty)

269 mock.assert_called_once_with(True) (empty)

270 

271 @pytest.mark.slow() (empty)

272 async def test_consume_list_batch( (empty)

273 self, 

274 queue: str, 

275 ) -> None: 

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

277 

278 msgs_queue = asyncio.Queue(maxsize=1) (empty)

279 

280 @consume_broker.subscriber( (empty)

281 list=ListSub(queue, batch=True, polling_interval=0.01), 

282 ) 

283 async def handler(msg) -> None: (empty)

284 await msgs_queue.put(msg) (empty)

285 

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

287 await br.start() (empty)

288 

289 await br.publish_batch(1, "hi", list=queue) (empty)

290 

291 result, _ = await asyncio.wait( (empty)

292 (asyncio.create_task(msgs_queue.get()),), 

293 timeout=3, 

294 ) 

295 

296 assert [{1, "hi"}] == [set(r.result()) for r in result] (empty)

297 

298 @pytest.mark.slow() (empty)

299 async def test_consume_list_batch_complex( (empty)

300 self, 

301 queue: str, 

302 ) -> None: 

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

304 

305 from pydantic import BaseModel (empty)

306 

307 class Data(BaseModel): (empty)

308 m: str (empty)

309 

310 def __hash__(self): (empty)

311 return hash(self.m) (empty)

312 

313 msgs_queue = asyncio.Queue(maxsize=1) (empty)

314 

315 @consume_broker.subscriber( (empty)

316 list=ListSub(queue, batch=True, polling_interval=0.01), 

317 ) 

318 async def handler(msg: list[Data]) -> None: (empty)

319 await msgs_queue.put(msg) (empty)

320 

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

322 await br.start() (empty)

323 

324 await br.publish_batch(Data(m="hi"), Data(m="again"), list=queue) (empty)

325 

326 result, _ = await asyncio.wait( (empty)

327 (asyncio.create_task(msgs_queue.get()),), 

328 timeout=3, 

329 ) 

330 

331 assert [{Data(m="hi"), Data(m="again")}] == [set(r.result()) for r in result] (empty)

332 

333 @pytest.mark.slow() (empty)

334 async def test_consume_list_batch_native( (empty)

335 self, 

336 queue: str, 

337 ) -> None: 

338 consume_broker = self.get_broker() (empty)

339 

340 msgs_queue = asyncio.Queue(maxsize=1) (empty)

341 

342 @consume_broker.subscriber( (empty)

343 list=ListSub(queue, batch=True, polling_interval=0.01), 

344 ) 

345 async def handler(msg) -> None: (empty)

346 await msgs_queue.put(msg) (empty)

347 

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

349 await br.start() (empty)

350 

351 await br._connection.rpush(queue, 1, "hi") (empty)

352 

353 result, _ = await asyncio.wait( (empty)

354 (asyncio.create_task(msgs_queue.get()),), 

355 timeout=3, 

356 ) 

357 

358 assert [{1, "hi"}] == [set(r.result()) for r in result] (empty)

359 

360 async def test_get_one( (empty)

361 self, 

362 queue: str, 

363 ) -> None: 

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

365 subscriber = broker.subscriber(list=queue) (empty)

366 

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

368 await br.start() (empty)

369 

370 message = None (empty)

371 

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

373 nonlocal message 

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

375 

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

377 await br.publish("test_message", list=queue) (empty)

378 

379 await asyncio.wait( (empty)

380 ( 

381 asyncio.create_task(consume()), 

382 asyncio.create_task(publish()), 

383 ), 

384 timeout=10, 

385 ) 

386 

387 assert message is not None (empty)

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

389 

390 async def test_get_one_timeout( (empty)

391 self, 

392 queue: str, 

393 mock: MagicMock, 

394 ) -> None: 

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

396 subscriber = broker.subscriber(list=queue) (empty)

397 

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

399 await br.start() (empty)

400 

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

402 mock.assert_called_once_with(None) (empty)

403 

404 async def test_concurrent_consume_list( (empty)

405 self, 

406 queue: str, 

407 mock: MagicMock, 

408 ) -> None: 

409 event = asyncio.Event() (empty)

410 event2 = asyncio.Event() (empty)

411 

412 consume_broker = self.get_broker() (empty)

413 

414 @consume_broker.subscriber(list=ListSub(queue), max_workers=2) (empty)

415 async def handler(msg): (empty)

416 mock() (empty)

417 if event.is_set(): (empty)

418 event2.set() (empty)

419 else: 

420 event.set() (empty)

421 await asyncio.sleep(0.1) (empty)

422 

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

424 await br.start() (empty)

425 

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

427 await br.publish(i, list=queue) (empty)

428 

429 await asyncio.wait( (empty)

430 ( 

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

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

433 ), 

434 timeout=3, 

435 ) 

436 

437 assert event.is_set() (empty)

438 assert event2.is_set() (empty)

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

440 

441 async def test_iterator( (empty)

442 self, 

443 queue: str, 

444 ) -> None: 

445 expected_messages = ("test_message_1", "test_message_2") (empty)

446 

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

448 subscriber = broker.subscriber(list=queue) (empty)

449 

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

451 await br.start() (empty)

452 

453 async def publish_test_message(): (empty)

454 for msg in expected_messages: (empty)

455 await br.publish(msg, list=queue) (empty)

456 

457 _ = await asyncio.create_task(publish_test_message()) (empty)

458 

459 index_message = 0 (empty)

460 async for msg in subscriber: 460 ↛ exitline 460 didn't jump to the function exit(empty)

461 result_message = await msg.decode() (empty)

462 

463 assert result_message == expected_messages[index_message] (empty)

464 

465 index_message += 1 (empty)

466 if index_message >= len(expected_messages): (empty)

467 break (empty)

468 

469 

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

471@pytest.mark.redis() (empty)

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

473class TestConsumeStream(RedisTestcaseConfig): (empty)

474 @pytest.mark.slow() (empty)

475 async def test_consume_stream( (empty)

476 self, 

477 mock: MagicMock, 

478 queue: str, 

479 ) -> None: 

480 event = asyncio.Event() (empty)

481 

482 consume_broker = self.get_broker() (empty)

483 

484 @consume_broker.subscriber(stream=StreamSub(queue, polling_interval=10)) (empty)

485 async def handler(msg) -> None: (empty)

486 mock(msg) (empty)

487 event.set() (empty)

488 

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

490 await br.start() (empty)

491 

492 await asyncio.wait( (empty)

493 ( 

494 asyncio.create_task(br.publish("hello", stream=queue)), 

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

496 ), 

497 timeout=3, 

498 ) 

499 

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

501 

502 @pytest.mark.slow() (empty)

503 async def test_consume_stream_with_big_interval( (empty)

504 self, 

505 event: asyncio.Event, 

506 mock: MagicMock, 

507 queue: str, 

508 ) -> None: 

509 consume_broker = self.get_broker() (empty)

510 

511 @consume_broker.subscriber(stream=StreamSub(queue, polling_interval=100000)) (empty)

512 async def handler(msg): (empty)

513 mock(msg) (empty)

514 event.set() (empty)

515 

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

517 await br.start() (empty)

518 await asyncio.wait( (empty)

519 ( 

520 asyncio.create_task(br.publish("hello", stream=queue)), 

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

522 ), 

523 timeout=3, 

524 ) 

525 

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

527 

528 @pytest.mark.slow() (empty)

529 async def test_consume_stream_native( (empty)

530 self, 

531 mock: MagicMock, 

532 queue: str, 

533 ) -> None: 

534 event = asyncio.Event() (empty)

535 

536 consume_broker = self.get_broker() (empty)

537 

538 @consume_broker.subscriber(stream=StreamSub(queue, polling_interval=10)) (empty)

539 async def handler(msg) -> None: (empty)

540 mock(msg) (empty)

541 event.set() (empty)

542 

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

544 await br.start() (empty)

545 

546 await asyncio.wait( (empty)

547 ( 

548 asyncio.create_task( 

549 br._connection.xadd(queue, {"message": "hello"}), 

550 ), 

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

552 ), 

553 timeout=3, 

554 ) 

555 

556 mock.assert_called_once_with({"message": "hello"}) (empty)

557 

558 @pytest.mark.slow() (empty)

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

560 async def test_consume_stream_batch( (empty)

561 self, 

562 mock: MagicMock, 

563 queue: str, 

564 ) -> None: 

565 event = asyncio.Event() (empty)

566 

567 consume_broker = self.get_broker() (empty)

568 

569 @consume_broker.subscriber( (empty)

570 stream=StreamSub(queue, polling_interval=10, batch=True), 

571 ) 

572 async def handler(msg) -> None: (empty)

573 mock(msg) (empty)

574 event.set() (empty)

575 

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

577 await br.start() (empty)

578 

579 await asyncio.wait( (empty)

580 ( 

581 asyncio.create_task(br.publish("hello", stream=queue)), 

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

583 ), 

584 timeout=3, 

585 ) 

586 

587 mock.assert_called_once_with(["hello"]) (empty)

588 

589 @pytest.mark.slow() (empty)

590 async def test_consume_stream_batch_headers( (empty)

591 self, 

592 queue: str, 

593 mock: MagicMock, 

594 ) -> None: 

595 event = asyncio.Event() (empty)

596 

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

598 

599 @consume_broker.subscriber( (empty)

600 stream=StreamSub(queue, polling_interval=10, batch=True), 

601 ) 

602 def subscriber(m, msg: RedisMessage) -> None: (empty)

603 check = all( (empty)

604 ( 

605 msg.headers, 

606 msg.headers["correlation_id"] 

607 == msg.batch_headers[0]["correlation_id"], 

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

609 ), 

610 ) 

611 mock(check) (empty)

612 event.set() (empty)

613 

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

615 await br.start() (empty)

616 await asyncio.wait( (empty)

617 ( 

618 asyncio.create_task( 

619 br.publish("", stream=queue, headers={"custom": "1"}), 

620 ), 

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

622 ), 

623 timeout=3, 

624 ) 

625 

626 assert event.is_set() (empty)

627 mock.assert_called_once_with(True) (empty)

628 

629 @pytest.mark.slow() (empty)

630 async def test_consume_stream_batch_complex( (empty)

631 self, 

632 queue: str, 

633 ) -> None: 

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

635 

636 from pydantic import BaseModel (empty)

637 

638 class Data(BaseModel): (empty)

639 m: str (empty)

640 

641 msgs_queue = asyncio.Queue(maxsize=1) (empty)

642 

643 @consume_broker.subscriber( (empty)

644 stream=StreamSub(queue, polling_interval=10, batch=True), 

645 ) 

646 async def handler(msg: list[Data]) -> None: (empty)

647 await msgs_queue.put(msg) (empty)

648 

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

650 await br.start() (empty)

651 

652 await br.publish(Data(m="hi"), stream=queue) (empty)

653 

654 result, _ = await asyncio.wait( (empty)

655 (asyncio.create_task(msgs_queue.get()),), 

656 timeout=3, 

657 ) 

658 

659 assert next(iter(result)).result() == [Data(m="hi")] (empty)

660 

661 @pytest.mark.slow() (empty)

662 async def test_consume_stream_batch_native( (empty)

663 self, 

664 mock: MagicMock, 

665 queue: str, 

666 ) -> None: 

667 event = asyncio.Event() (empty)

668 

669 consume_broker = self.get_broker() (empty)

670 

671 @consume_broker.subscriber( (empty)

672 stream=StreamSub(queue, polling_interval=10, batch=True), 

673 ) 

674 async def handler(msg) -> None: (empty)

675 mock(msg) (empty)

676 event.set() (empty)

677 

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

679 await br.start() (empty)

680 

681 await asyncio.wait( (empty)

682 ( 

683 asyncio.create_task( 

684 br._connection.xadd(queue, {"message": "hello"}), 

685 ), 

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

687 ), 

688 timeout=3, 

689 ) 

690 

691 mock.assert_called_once_with([{"message": "hello"}]) (empty)

692 

693 async def test_consume_group( (empty)

694 self, 

695 queue: str, 

696 ) -> None: 

697 consume_broker = self.get_broker() (empty)

698 

699 @consume_broker.subscriber( (empty)

700 stream=StreamSub(queue, group="group", consumer=queue), (empty)

701 ) 

702 async def handler(msg: RedisMessage) -> None: ... (empty)

703 

704 assert next(iter(consume_broker.subscribers)).last_id == ">" (empty)

705 

706 async def test_consume_group_with_last_id( (empty)

707 self, 

708 queue: str, 

709 ) -> None: 

710 consume_broker = self.get_broker() (empty)

711 

712 @consume_broker.subscriber( (empty)

713 stream=StreamSub(queue, group="group", consumer=queue, last_id="0"), (empty)

714 ) 

715 async def handler(msg: RedisMessage) -> None: ... (empty)

716 

717 assert next(iter(consume_broker.subscribers)).last_id == "0" (empty)

718 

719 async def test_consume_nack( (empty)

720 self, 

721 queue: str, 

722 ) -> None: 

723 event = asyncio.Event() (empty)

724 

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

726 

727 @consume_broker.subscriber( (empty)

728 stream=StreamSub(queue, group="group", consumer=queue), 

729 ) 

730 async def handler(msg: RedisMessage) -> None: (empty)

731 event.set() (empty)

732 await msg.nack() (empty)

733 

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

735 await br.start() (empty)

736 

737 with patch.object(Redis, "xack", spy_decorator(Redis.xack)) as m: (empty)

738 await asyncio.wait( (empty)

739 ( 

740 asyncio.create_task(br.publish("hello", stream=queue)), 

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

742 ), 

743 timeout=3, 

744 ) 

745 

746 assert not m.mock.called (empty)

747 

748 assert event.is_set() (empty)

749 

750 async def test_consume_ack( (empty)

751 self, 

752 queue: str, 

753 ) -> None: 

754 event = asyncio.Event() (empty)

755 

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

757 

758 @consume_broker.subscriber( (empty)

759 stream=StreamSub(queue, group="group", consumer=queue), 

760 ) 

761 async def handler(msg: RedisMessage) -> None: (empty)

762 event.set() (empty)